copilot-swe-agent[bot]
df051e8b81
Implements the classify pipeline step with: - Classification rules engine (app/utils/classification_rules.py) with pre-built categories (invoice, contract, receipt, letter, report, bank_statement, tax_document, insurance, payslip) and support for filename patterns, content keywords, and metadata matching rules - Celery task (app/tasks/classify_document.py) that runs as a pipeline step - CRUD API (app/api/classification_rules.py) for managing custom rules - ClassificationRuleModel in app/models.py with migration 027 - Updated pipeline step config_schema and stage mapping - Comprehensive tests for engine, API, and task Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
233 lines
9.3 KiB
Python
233 lines
9.3 KiB
Python
"""Tests for the classify_document Celery task.
|
|
|
|
Covers the ``classify_document_task`` in ``app/tasks/classify_document.py``.
|
|
"""
|
|
|
|
import json
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from app.models import ClassificationRuleModel, FileRecord
|
|
from app.tasks.classify_document import _load_custom_rules
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_file_record(db_session, **overrides):
|
|
"""Insert a minimal FileRecord and return it."""
|
|
defaults = {
|
|
"owner_id": "test-user",
|
|
"filehash": "abc123",
|
|
"original_filename": "Invoice_2024.pdf",
|
|
"local_filename": "/tmp/test.pdf",
|
|
"file_size": 1024,
|
|
"mime_type": "application/pdf",
|
|
"ocr_text": "Invoice number: 12345. Amount due: $500.",
|
|
"ai_metadata": None,
|
|
}
|
|
defaults.update(overrides)
|
|
fr = FileRecord(**defaults)
|
|
db_session.add(fr)
|
|
db_session.commit()
|
|
db_session.refresh(fr)
|
|
return fr
|
|
|
|
|
|
def _make_rule(db_session, **overrides):
|
|
"""Insert a ClassificationRuleModel and return it."""
|
|
defaults = {
|
|
"owner_id": None,
|
|
"name": "test_rule",
|
|
"category": "test_category",
|
|
"rule_type": "filename_pattern",
|
|
"pattern": r"(?i)test",
|
|
"priority": 0,
|
|
"case_sensitive": False,
|
|
"enabled": True,
|
|
}
|
|
defaults.update(overrides)
|
|
rule = ClassificationRuleModel(**defaults)
|
|
db_session.add(rule)
|
|
db_session.commit()
|
|
db_session.refresh(rule)
|
|
return rule
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _load_custom_rules
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestLoadCustomRules:
|
|
"""Test the custom rule loading helper."""
|
|
|
|
@patch("app.tasks.classify_document.SessionLocal")
|
|
def test_loads_enabled_rules(self, mock_session_local):
|
|
"""Should load enabled rules from the database."""
|
|
mock_rule = MagicMock()
|
|
mock_rule.name = "rule1"
|
|
mock_rule.category = "invoice"
|
|
mock_rule.rule_type = "filename_pattern"
|
|
mock_rule.pattern = r"(?i)invoice"
|
|
mock_rule.priority = 10
|
|
mock_rule.case_sensitive = False
|
|
|
|
mock_db = MagicMock()
|
|
mock_query = MagicMock()
|
|
mock_db.query.return_value = mock_query
|
|
mock_query.filter.return_value = mock_query
|
|
mock_query.order_by.return_value = mock_query
|
|
mock_query.all.return_value = [mock_rule]
|
|
mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db)
|
|
mock_session_local.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
rules = _load_custom_rules(owner_id="test-user")
|
|
assert len(rules) == 1
|
|
assert rules[0].name == "rule1"
|
|
assert rules[0].category == "invoice"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# classify_document_task (integration-style with mocked DB and Celery)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestClassifyDocumentTask:
|
|
"""Test the Celery classify_document_task."""
|
|
|
|
@patch("app.tasks.classify_document.log_task_progress")
|
|
@patch("app.tasks.classify_document._load_custom_rules", return_value=[])
|
|
@patch("app.tasks.classify_document.SessionLocal")
|
|
def test_classify_invoice_file(self, mock_session_local, mock_load_rules, mock_log):
|
|
"""Should classify a file with invoice filename and text as 'invoice'."""
|
|
mock_file = MagicMock(spec=FileRecord)
|
|
mock_file.id = 1
|
|
mock_file.original_filename = "Invoice_2024.pdf"
|
|
mock_file.ocr_text = "Invoice number: 12345. Amount due: $500."
|
|
mock_file.ai_metadata = None
|
|
mock_file.owner_id = "test-user"
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query.return_value.filter.return_value.first.return_value = mock_file
|
|
mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db)
|
|
mock_session_local.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
from app.tasks.classify_document import classify_document_task
|
|
|
|
# Call the underlying function directly via .run(), bypassing Celery
|
|
result = classify_document_task.run(1, owner_id="test-user")
|
|
|
|
assert result["status"] == "success"
|
|
assert result["category"] == "invoice"
|
|
assert result["confidence"] > 0
|
|
|
|
# Verify ai_metadata was updated
|
|
assert mock_file.ai_metadata is not None
|
|
metadata = json.loads(mock_file.ai_metadata)
|
|
assert "classification" in metadata
|
|
assert metadata["classification"]["category"] == "invoice"
|
|
|
|
@patch("app.tasks.classify_document.log_task_progress")
|
|
@patch("app.tasks.classify_document.SessionLocal")
|
|
def test_classify_file_not_found(self, mock_session_local, mock_log):
|
|
"""Should return error when file record is not found."""
|
|
mock_db = MagicMock()
|
|
mock_db.query.return_value.filter.return_value.first.return_value = None
|
|
mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db)
|
|
mock_session_local.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
from app.tasks.classify_document import classify_document_task
|
|
|
|
result = classify_document_task.run(99999)
|
|
assert result["status"] == "error"
|
|
|
|
@patch("app.tasks.classify_document.log_task_progress")
|
|
@patch("app.tasks.classify_document._load_custom_rules", return_value=[])
|
|
@patch("app.tasks.classify_document.SessionLocal")
|
|
def test_classify_preserves_existing_metadata(self, mock_session_local, mock_load_rules, mock_log):
|
|
"""Should preserve existing ai_metadata fields and add classification."""
|
|
existing_meta = json.dumps({"document_type": "Invoice", "tags": ["finance"]})
|
|
|
|
mock_file = MagicMock(spec=FileRecord)
|
|
mock_file.id = 2
|
|
mock_file.original_filename = "doc.pdf"
|
|
mock_file.ocr_text = ""
|
|
mock_file.ai_metadata = existing_meta
|
|
mock_file.owner_id = "test-user"
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query.return_value.filter.return_value.first.return_value = mock_file
|
|
mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db)
|
|
mock_session_local.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
from app.tasks.classify_document import classify_document_task
|
|
|
|
classify_document_task.run(2)
|
|
|
|
# Check that existing fields are preserved
|
|
metadata = json.loads(mock_file.ai_metadata)
|
|
assert metadata["tags"] == ["finance"]
|
|
assert metadata["document_type"] == "Invoice"
|
|
assert "classification" in metadata
|
|
|
|
@patch("app.tasks.classify_document.log_task_progress")
|
|
@patch("app.tasks.classify_document._load_custom_rules", return_value=[])
|
|
@patch("app.tasks.classify_document.SessionLocal")
|
|
def test_classify_sets_document_type_when_missing(self, mock_session_local, mock_load_rules, mock_log):
|
|
"""Should set document_type from classification when not already present."""
|
|
mock_file = MagicMock(spec=FileRecord)
|
|
mock_file.id = 3
|
|
mock_file.original_filename = "Invoice_2024.pdf"
|
|
mock_file.ocr_text = "Invoice number: 12345"
|
|
mock_file.ai_metadata = json.dumps({"tags": ["test"]})
|
|
mock_file.owner_id = "test-user"
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query.return_value.filter.return_value.first.return_value = mock_file
|
|
mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db)
|
|
mock_session_local.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
from app.tasks.classify_document import classify_document_task
|
|
|
|
classify_document_task.run(3)
|
|
|
|
metadata = json.loads(mock_file.ai_metadata)
|
|
assert metadata["document_type"] == "Invoice"
|
|
|
|
@patch("app.tasks.classify_document.log_task_progress")
|
|
@patch("app.tasks.classify_document._load_custom_rules", return_value=[])
|
|
@patch("app.tasks.classify_document.SessionLocal")
|
|
def test_classify_unknown_document(self, mock_session_local, mock_load_rules, mock_log):
|
|
"""Should classify as 'unknown' when no rules match."""
|
|
mock_file = MagicMock(spec=FileRecord)
|
|
mock_file.id = 4
|
|
mock_file.original_filename = "random_file.pdf"
|
|
mock_file.ocr_text = "Lorem ipsum dolor sit amet."
|
|
mock_file.ai_metadata = None
|
|
mock_file.owner_id = "test-user"
|
|
|
|
mock_db = MagicMock()
|
|
mock_db.query.return_value.filter.return_value.first.return_value = mock_file
|
|
mock_session_local.return_value.__enter__ = MagicMock(return_value=mock_db)
|
|
mock_session_local.return_value.__exit__ = MagicMock(return_value=False)
|
|
|
|
from app.tasks.classify_document import classify_document_task
|
|
|
|
result = classify_document_task.run(4)
|
|
|
|
assert result["category"] == "unknown"
|
|
assert result["confidence"] == 0
|
|
|
|
def test_classify_document_task_is_celery_task(self):
|
|
"""Task should be registered as a Celery task."""
|
|
from app.tasks.classify_document import classify_document_task
|
|
|
|
assert hasattr(classify_document_task, "apply_async")
|
|
assert hasattr(classify_document_task, "delay")
|
|
assert callable(classify_document_task)
|