From 013b8dbc3253555fc1bfeec065fcb1a7c1ae94b1 Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:58:34 -0700 Subject: [PATCH] Codex please write backend tests thx --- src/documents/tests/test_api_workflows.py | 139 ++++++++ src/documents/tests/test_tasks.py | 108 ++++++ src/documents/tests/test_workflows.py | 388 ++++++++++++++++++++++ 3 files changed, 635 insertions(+) diff --git a/src/documents/tests/test_api_workflows.py b/src/documents/tests/test_api_workflows.py index 58c9379d7..cb2fc6f66 100644 --- a/src/documents/tests/test_api_workflows.py +++ b/src/documents/tests/test_api_workflows.py @@ -466,6 +466,145 @@ class TestApiWorkflows(DirectoriesMixin, APITestCase): self.assertEqual(response.status_code, status.HTTP_201_CREATED) + def _post_ai_suggestions_workflow(self, *, trigger_types, action: dict): + def trigger(trigger_type): + # consumption triggers require a filter of their own + if trigger_type == WorkflowTrigger.WorkflowTriggerType.CONSUMPTION: + return {"type": trigger_type, "filter_filename": "*.pdf"} + return {"type": trigger_type} + + return self.client.post( + self.ENDPOINT, + json.dumps( + { + "name": "Apply AI suggestions", + "order": 1, + "triggers": [trigger(t) for t in trigger_types], + "actions": [ + { + "type": WorkflowAction.WorkflowActionType.APPLY_AI_SUGGESTIONS, + **action, + }, + ], + }, + ), + content_type="application/json", + ) + + def test_api_create_apply_ai_suggestions_action(self) -> None: + """ + GIVEN: + - API request to create a workflow with an apply AI suggestions + action and a valid set of fields + WHEN: + - API is called + THEN: + - The workflow is created with the chosen options + """ + response = self._post_ai_suggestions_workflow( + trigger_types=[WorkflowTrigger.WorkflowTriggerType.DOCUMENT_ADDED], + action={ + "ai_suggestion_fields": ["title", "tags", "correspondent"], + "ai_create_missing": True, + "ai_overwrite_existing": True, + }, + ) + + self.assertEqual(response.status_code, status.HTTP_201_CREATED) + action = Workflow.objects.get(name="Apply AI suggestions").actions.first() + self.assertEqual( + action.ai_suggestion_fields, + ["title", "tags", "correspondent"], + ) + self.assertTrue(action.ai_create_missing) + self.assertTrue(action.ai_overwrite_existing) + + def test_api_create_apply_ai_suggestions_action_requires_fields(self) -> None: + """ + GIVEN: + - API request to create an apply AI suggestions action with no + fields selected, which could never do anything + WHEN: + - API is called + THEN: + - Correct HTTP 400 response + - No objects are created + """ + existing_count = Workflow.objects.count() + + response = self._post_ai_suggestions_workflow( + trigger_types=[WorkflowTrigger.WorkflowTriggerType.DOCUMENT_ADDED], + action={"ai_suggestion_fields": []}, + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertEqual(Workflow.objects.count(), existing_count) + + def test_api_create_apply_ai_suggestions_action_rejects_unknown_field( + self, + ) -> None: + """ + GIVEN: + - API request to create an apply AI suggestions action naming a + field that does not exist + WHEN: + - API is called + THEN: + - Correct HTTP 400 response + """ + response = self._post_ai_suggestions_workflow( + trigger_types=[WorkflowTrigger.WorkflowTriggerType.DOCUMENT_ADDED], + action={"ai_suggestion_fields": ["title", "not_a_field"]}, + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + + def test_api_create_apply_ai_suggestions_action_rejects_consumption_only( + self, + ) -> None: + """ + GIVEN: + - API request to create an apply AI suggestions action whose only + trigger is consumption started, so there is no document content + to make suggestions from yet + WHEN: + - API is called + THEN: + - Correct HTTP 400 response + - No objects are created + """ + existing_count = Workflow.objects.count() + + response = self._post_ai_suggestions_workflow( + trigger_types=[WorkflowTrigger.WorkflowTriggerType.CONSUMPTION], + action={"ai_suggestion_fields": ["title"]}, + ) + + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertEqual(Workflow.objects.count(), existing_count) + + def test_api_create_apply_ai_suggestions_action_allows_extra_consumption_trigger( + self, + ) -> None: + """ + GIVEN: + - API request to create an apply AI suggestions action with a + consumption trigger alongside a usable one + WHEN: + - API is called + THEN: + - The workflow is created, the action applies to the other trigger + """ + response = self._post_ai_suggestions_workflow( + trigger_types=[ + WorkflowTrigger.WorkflowTriggerType.CONSUMPTION, + WorkflowTrigger.WorkflowTriggerType.DOCUMENT_ADDED, + ], + action={"ai_suggestion_fields": ["title"]}, + ) + + self.assertEqual(response.status_code, status.HTTP_201_CREATED) + def test_api_create_workflow_trigger_action_empty_fields(self) -> None: """ GIVEN: diff --git a/src/documents/tests/test_tasks.py b/src/documents/tests/test_tasks.py index 2954d6442..107c1dc6b 100644 --- a/src/documents/tests/test_tasks.py +++ b/src/documents/tests/test_tasks.py @@ -14,6 +14,7 @@ from documents.models import Correspondent from documents.models import Document from documents.models import DocumentType from documents.models import Tag +from documents.models import WorkflowAction from documents.sanity_checker import SanityCheckFailedException from documents.sanity_checker import SanityCheckMessages from documents.tests.test_classifier import dummy_preprocess @@ -447,3 +448,110 @@ class TestAIIndex(DirectoriesMixin, TestCase): rebuild=False, document_ids=doc_ids, ) + + +class TestApplyAISuggestionsTask(DirectoriesMixin, TestCase): + def setUp(self) -> None: + super().setUp() + self.doc = Document.objects.create( + title="doc", + content="content", + checksum="apply-ai-suggestions", + ) + self.action = WorkflowAction.objects.create( + type=WorkflowAction.WorkflowActionType.APPLY_AI_SUGGESTIONS, + ai_suggestion_fields=[WorkflowAction.AISuggestionField.TITLE], + ) + + def test_reindexes_without_sending_document_updated(self) -> None: + """ + GIVEN: + - An apply AI suggestions action that changes the document + WHEN: + - The task runs + THEN: + - The search index and caches are refreshed directly, deliberately + not via the document_updated signal: that re-runs updated + workflows, which for this action means queueing another LLM + query for a document it just changed, forever + """ + with ( + mock.patch( + "documents.workflows.ai.apply_ai_suggestions_to_document", + return_value=["title"], + ), + mock.patch("documents.tasks.index_document") as index_document, + mock.patch("documents.tasks.clear_document_caches") as clear_caches, + mock.patch("documents.tasks.document_updated") as document_updated, + ): + tasks.apply_ai_suggestions(self.action.pk, self.doc.pk) + + index_document.delay.assert_called_once_with(self.doc.pk) + clear_caches.assert_called_once_with(self.doc.pk) + document_updated.send.assert_not_called() + + def test_no_changes_skips_reindex(self) -> None: + """ + GIVEN: + - An apply AI suggestions action that changes nothing + WHEN: + - The task runs + THEN: + - No reindexing work is queued + """ + with ( + mock.patch( + "documents.workflows.ai.apply_ai_suggestions_to_document", + return_value=[], + ), + mock.patch("documents.tasks.index_document") as index_document, + ): + tasks.apply_ai_suggestions(self.action.pk, self.doc.pk) + + index_document.delay.assert_not_called() + + @override_settings(AI_ENABLED=True, LLM_EMBEDDING_BACKEND="huggingface") + def test_updates_llm_index_when_enabled(self) -> None: + """ + GIVEN: + - An apply AI suggestions action that changes the document + - The LLM index is enabled + WHEN: + - The task runs + THEN: + - The document is updated in the LLM index too + """ + with ( + mock.patch( + "documents.workflows.ai.apply_ai_suggestions_to_document", + return_value=["title"], + ), + mock.patch("documents.tasks.index_document"), + mock.patch( + "documents.tasks.update_document_in_llm_index", + ) as update_in_llm_index, + ): + tasks.apply_ai_suggestions(self.action.pk, self.doc.pk) + + update_in_llm_index.apply_async.assert_called_once() + + def test_deleted_document_is_a_noop(self) -> None: + """ + GIVEN: + - A document that was deleted between the workflow running and the + queued task starting + WHEN: + - The task runs + THEN: + - It logs and exits rather than raising + """ + with ( + mock.patch( + "documents.workflows.ai.apply_ai_suggestions_to_document", + ) as apply_suggestions, + self.assertLogs("paperless.tasks", level="WARNING") as cm, + ): + tasks.apply_ai_suggestions(self.action.pk, self.doc.pk + 1000) + + apply_suggestions.assert_not_called() + self.assertIn("no longer exists", "".join(cm.output)) diff --git a/src/documents/tests/test_workflows.py b/src/documents/tests/test_workflows.py index 9e91151c0..04afbc6e3 100644 --- a/src/documents/tests/test_workflows.py +++ b/src/documents/tests/test_workflows.py @@ -31,6 +31,7 @@ from documents.file_handling import create_source_path_directory from documents.file_handling import generate_filename from documents.file_handling import generate_unique_filename from documents.signals.handlers import run_workflows +from documents.workflows.ai import apply_ai_suggestions_to_document from documents.workflows.webhooks import send_webhook if TYPE_CHECKING: @@ -5439,3 +5440,390 @@ class TestRemoteOCRWorkflowAction(DirectoriesMixin, SampleDirMixin, APITestCase) ) self.assertIn("only applies to consumption triggers", "".join(cm.output)) + + +SUGGESTIONS = { + "title": "Suggested Title", + "tags": ["Existing Tag", "Suggested Tag"], + "correspondents": ["Existing Correspondent", "Suggested Correspondent"], + "document_types": ["Suggested Document Type"], + "storage_paths": ["Suggested Storage Path"], + "dates": ["2024-03-05"], +} + +ALL_SUGGESTION_FIELDS = [ + WorkflowAction.AISuggestionField.TITLE, + WorkflowAction.AISuggestionField.TAGS, + WorkflowAction.AISuggestionField.CORRESPONDENT, + WorkflowAction.AISuggestionField.DOCUMENT_TYPE, + WorkflowAction.AISuggestionField.STORAGE_PATH, + WorkflowAction.AISuggestionField.CREATED, +] + + +@override_settings(AI_ENABLED=True) +class TestApplyAISuggestionsWorkflowAction( + DirectoriesMixin, + SampleDirMixin, + APITestCase, +): + def setUp(self) -> None: + super().setUp() + self.user = User.objects.create(username="ai-user") + self.doc = Document.objects.create( + title="original.pdf", + content="the document content", + checksum="ai-suggestions-checksum", + mime_type="application/pdf", + created=datetime.date(2020, 1, 1), + owner=self.user, + ) + + def make_action(self, **kwargs) -> WorkflowAction: + return WorkflowAction.objects.create( + type=WorkflowAction.WorkflowActionType.APPLY_AI_SUGGESTIONS, + ai_suggestion_fields=kwargs.pop( + "ai_suggestion_fields", + ALL_SUGGESTION_FIELDS, + ), + **kwargs, + ) + + def make_workflow(self, action: WorkflowAction, trigger_type) -> Workflow: + trigger = WorkflowTrigger.objects.create(type=trigger_type) + w = Workflow.objects.create(name="Apply AI suggestions", order=0) + w.triggers.add(trigger) + w.actions.add(action) + w.save() + return w + + def apply(self, action: WorkflowAction) -> list[str]: + with mock.patch( + "documents.workflows.ai.get_ai_document_classification", + return_value=SUGGESTIONS, + ): + changed = apply_ai_suggestions_to_document(action, self.doc) + self.doc.refresh_from_db() + return changed + + def test_document_added_trigger_queues_task(self) -> None: + """ + GIVEN: + - A document added workflow with an apply AI suggestions action + WHEN: + - A matching document is added + THEN: + - The work is queued rather than run inline, so a slow LLM query + cannot stall the rest of the workflow run + """ + action = self.make_action() + self.make_workflow(action, WorkflowTrigger.WorkflowTriggerType.DOCUMENT_ADDED) + + with mock.patch("documents.tasks.apply_ai_suggestions.delay") as delay: + run_workflows( + WorkflowTrigger.WorkflowTriggerType.DOCUMENT_ADDED, + self.doc, + ) + + delay.assert_called_once_with(action.pk, self.doc.pk) + + def test_consumption_trigger_is_ignored(self) -> None: + """ + GIVEN: + - A workflow with an apply AI suggestions action and a consumption + trigger alongside a valid one + WHEN: + - The consumption trigger fires + THEN: + - The action is skipped, since the document has not been parsed + yet and so has no content to make suggestions from + """ + action = self.make_action() + w = self.make_workflow( + action, + WorkflowTrigger.WorkflowTriggerType.DOCUMENT_ADDED, + ) + w.triggers.add( + WorkflowTrigger.objects.create( + type=WorkflowTrigger.WorkflowTriggerType.CONSUMPTION, + ), + ) + + test_file = shutil.copy( + self.SAMPLE_DIR / "simple.pdf", + self.dirs.scratch_dir / "simple.pdf", + ) + + with ( + mock.patch("documents.tasks.apply_ai_suggestions.delay") as delay, + self.assertLogs("paperless.handlers", level="DEBUG") as cm, + ): + run_workflows( + WorkflowTrigger.WorkflowTriggerType.CONSUMPTION, + ConsumableDocument( + source=DocumentSource.ConsumeFolder, + original_file=test_file, + ), + overrides=DocumentMetadataOverrides(), + ) + + delay.assert_not_called() + self.assertIn("does not apply to consumption triggers", "".join(cm.output)) + + def test_no_selected_fields_does_nothing(self) -> None: + """ + GIVEN: + - An action with no suggestion fields selected + WHEN: + - The action is applied + THEN: + - Nothing is changed and it is logged + """ + action = self.make_action(ai_suggestion_fields=[]) + + with self.assertLogs("paperless.workflows.ai", level="WARNING") as cm: + changed = self.apply(action) + + self.assertEqual(changed, []) + self.assertIn("no AI suggestion fields selected", "".join(cm.output)) + + @override_settings(AI_ENABLED=False) + def test_ai_disabled_does_nothing(self) -> None: + """ + GIVEN: + - An action on an install where AI has since been disabled + WHEN: + - The action is applied + THEN: + - Nothing is changed and it is logged + """ + action = self.make_action() + + with self.assertLogs("paperless.workflows.ai", level="ERROR") as cm: + changed = self.apply(action) + + self.assertEqual(changed, []) + self.assertIn("AI is not enabled", "".join(cm.output)) + + def test_llm_failure_leaves_document_untouched(self) -> None: + """ + GIVEN: + - An LLM backend that errors out + WHEN: + - The action is applied + THEN: + - The failure is logged and the document is left alone, rather + than the error taking down the whole task + """ + action = self.make_action() + + with ( + mock.patch( + "documents.workflows.ai.get_ai_document_classification", + side_effect=ValueError("nope"), + ), + self.assertLogs("paperless.workflows.ai", level="ERROR") as cm, + ): + changed = apply_ai_suggestions_to_document(action, self.doc) + + self.assertEqual(changed, []) + self.doc.refresh_from_db() + self.assertEqual(self.doc.title, "original.pdf") + self.assertIn("Error getting AI suggestions", "".join(cm.output)) + + def test_only_matching_objects_are_applied(self) -> None: + """ + GIVEN: + - An action without create missing, and only some of the suggested + objects existing + WHEN: + - The action is applied + THEN: + - Only the existing objects are assigned, unmatched suggestions are + dropped rather than creating anything + """ + tag = Tag.objects.create(name="Existing Tag", owner=self.user) + correspondent = Correspondent.objects.create( + name="Existing Correspondent", + owner=self.user, + ) + action = self.make_action(ai_overwrite_existing=True) + + changed = self.apply(action) + + self.assertEqual(self.doc.correspondent, correspondent) + self.assertEqual(list(self.doc.tags.all()), [tag]) + # Nothing matched for these and create missing is off + self.assertIsNone(self.doc.document_type) + self.assertIsNone(self.doc.storage_path) + self.assertNotIn("document_type", changed) + self.assertEqual(Tag.objects.count(), 1) + self.assertEqual(Correspondent.objects.count(), 1) + + def test_create_missing_creates_objects_owned_by_document_owner(self) -> None: + """ + GIVEN: + - An action with create missing enabled + WHEN: + - The action is applied and suggestions match nothing + THEN: + - Tags, correspondents and document types are created, owned by the + document owner so they stay private to them + - Storage paths are never created, since a path template cannot be + inferred from a name + """ + action = self.make_action( + ai_create_missing=True, + ai_overwrite_existing=True, + ) + + changed = self.apply(action) + + self.assertEqual( + sorted(t.name for t in self.doc.tags.all()), + ["Existing Tag", "Suggested Tag"], + ) + self.assertEqual(self.doc.correspondent.name, "Existing Correspondent") + self.assertEqual(self.doc.correspondent.owner, self.user) + self.assertEqual(self.doc.document_type.name, "Suggested Document Type") + self.assertEqual(self.doc.document_type.owner, self.user) + + self.assertIsNone(self.doc.storage_path) + self.assertFalse(StoragePath.objects.exists()) + self.assertNotIn("storage_path", changed) + + def test_overwrite_disabled_keeps_existing_values(self) -> None: + """ + GIVEN: + - An action without overwrite existing + - A document that already has a title, created date and + correspondent + WHEN: + - The action is applied + THEN: + - The existing values are kept, only the empty document type is + filled in + """ + existing = Correspondent.objects.create(name="Mine", owner=self.user) + self.doc.correspondent = existing + self.doc.save() + action = self.make_action(ai_create_missing=True) + + changed = self.apply(action) + + self.assertEqual(self.doc.title, "original.pdf") + self.assertEqual(self.doc.created, datetime.date(2020, 1, 1)) + self.assertEqual(self.doc.correspondent, existing) + self.assertEqual(self.doc.document_type.name, "Suggested Document Type") + self.assertNotIn("title", changed) + self.assertNotIn("correspondent", changed) + + def test_overwrite_enabled_replaces_existing_values(self) -> None: + """ + GIVEN: + - An action with overwrite existing + - A document that already has a title and created date + WHEN: + - The action is applied + THEN: + - The suggested values replace them + """ + action = self.make_action( + ai_create_missing=True, + ai_overwrite_existing=True, + ) + + changed = self.apply(action) + + self.assertEqual(self.doc.title, "Suggested Title") + self.assertEqual(self.doc.created, datetime.date(2024, 3, 5)) + self.assertIn("title", changed) + self.assertIn("created", changed) + + def test_tags_are_added_not_replaced(self) -> None: + """ + GIVEN: + - A document that already has a tag unrelated to the suggestions + WHEN: + - The action is applied with overwrite existing enabled + THEN: + - The existing tag is kept, since suggested tags are always + additive regardless of the overwrite setting + """ + kept = Tag.objects.create(name="Do Not Remove", owner=self.user) + self.doc.tags.add(kept) + Tag.objects.create(name="Existing Tag", owner=self.user) + action = self.make_action(ai_overwrite_existing=True) + + self.apply(action) + + self.assertEqual( + sorted(t.name for t in self.doc.tags.all()), + ["Do Not Remove", "Existing Tag"], + ) + + def test_unselected_fields_are_untouched(self) -> None: + """ + GIVEN: + - An action that only selects the title + WHEN: + - The action is applied + THEN: + - Only the title changes, even though the LLM suggested everything + """ + action = self.make_action( + ai_suggestion_fields=[WorkflowAction.AISuggestionField.TITLE], + ai_create_missing=True, + ai_overwrite_existing=True, + ) + + changed = self.apply(action) + + self.assertEqual(changed, ["title"]) + self.assertEqual(self.doc.title, "Suggested Title") + self.assertEqual(self.doc.tags.count(), 0) + self.assertIsNone(self.doc.correspondent) + self.assertEqual(self.doc.created, datetime.date(2020, 1, 1)) + + def test_another_users_private_objects_are_not_matched(self) -> None: + """ + GIVEN: + - A suggested tag name that exists, but is owned by someone else + WHEN: + - The action is applied + THEN: + - It is not assigned, because the document owner cannot see it + """ + other = User.objects.create(username="someone-else") + Tag.objects.create(name="Existing Tag", owner=other) + action = self.make_action( + ai_suggestion_fields=[WorkflowAction.AISuggestionField.TAGS], + ) + + self.apply(action) + + self.assertEqual(self.doc.tags.count(), 0) + + def test_unparsable_dates_are_skipped(self) -> None: + """ + GIVEN: + - Suggested dates that are not all valid + WHEN: + - The action is applied + THEN: + - The first usable date is applied and the rest ignored + """ + action = self.make_action( + ai_suggestion_fields=[WorkflowAction.AISuggestionField.CREATED], + ai_overwrite_existing=True, + ) + + with mock.patch( + "documents.workflows.ai.get_ai_document_classification", + return_value={**SUGGESTIONS, "dates": ["not a date", "2019-07-04"]}, + ): + changed = apply_ai_suggestions_to_document(action, self.doc) + + self.doc.refresh_from_db() + self.assertEqual(changed, ["created"]) + self.assertEqual(self.doc.created, datetime.date(2019, 7, 4))