copilot-swe-agent[bot]
c7d3ec57c3
Commitd2217531(google-labs-jules SSRF fix) catastrophically deleted 11,500+ lines across 100+ files while fixing an unrelated IMAP issue. Restored from d2217531^ (pre-bad-commit state): Deleted files (fully restored): - app/api/{automation,classification_rules,comments,sharing}.py - app/middleware/upload_rate_limit.py - app/tasks/{automation_tasks,classify_document}.py - app/utils/{automation_hooks,classification_rules}.py - docs/AppleAppStoreCompliance.md - frontend/input.css, package.json, package-lock.json, tailwind.config.js - frontend/static/js/{annotations,claim,comments,sharing}.js - frontend/templates/{admin_connections,file_annotations,file_summary}.html - tests/{test_api_files_comprehensive,test_auth_extended,test_sharing, test_comments,test_connections,test_imap_profiles,test_api_sessions, test_automation,test_classification_rules,test_api_advanced_filters, test_api_classification_rules,test_upload_rate_limit,test_api_dropbox, test_classify_document,test_comments_ui,test_upload_to_icloud, test_api_onedrive_comprehensive,test_frontend_build,test_sentry, test_diagnostic,test_database,test_views_dropbox,test_local_auth}.py Truncated files (content restored): - app/{auth,config,main,models,celery_worker,database}.py - app/api/{__init__,api_tokens,diagnostic,dropbox,files,google_drive, integrations,local_auth,mobile,onedrive,pipelines,qr_auth, settings,url_upload}.py - app/middleware/upload_rate_limit.py - app/tasks/upload_to_nextcloud.py - app/utils/{allowed_types,settings_service,settings_sync,user_scope,webhook}.py - app/views/{base,dropbox,files,google_drive,onedrive,settings}.py - docs/{API,AuthenticationSetup,ConfigurationGuide,DatabaseConfiguration, DeploymentGuide,DropboxSetup,GoogleDriveSetup,KubernetesDeployment, MobileApp,OneDriveSetup,ProductionReadiness,SentrySetup, SocialLoginSetup,UserGuide}.md - frontend/static/{js/upload.js,styles.css} - frontend/templates/{api_tokens,base,devices,dropbox,dropbox_callback, file_view,files,google_drive,onedrive,onedrive_callback, signup}.html - frontend/translations/en.json - migrations/env.py - tests/{conftest,test_api_integrations,test_api_mobile,test_api_settings, test_api_tokens,test_audit_logs,test_duplicates,test_imap_tasks, test_setup_wizard,test_views_files_comprehensive}.py Security fixes kept from post-d2217531 commits: - app/utils/network.py: DNS SSRF fail-secure fix (06b0fced) - app/utils/file_operations.py: path traversal fix (1018ea17) - tests/test_imap_tasks.py: re-applied 4 is_private_ip mock patches Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com> Agent-Logs-Url: https://github.com/christianlouis/DocuElevate/sessions/51133dd8-9bec-41ab-aa10-3de753634187
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)
|