Messing around

src/documents/tests/test_tasks.py

@@ -1,21 +1,32 @@
 import os
+import shutil
+import uuid
 from datetime import timedelta
+from pathlib import Path
 from unittest import mock
 
 from django.conf import settings
 from django.test import TestCase
+from django.test import override_settings
 from django.utils import timezone
 
 from documents import tasks
+from documents.data_models import ConsumableDocument
+from documents.data_models import DocumentSource
 from documents.models import Correspondent
 from documents.models import Document
 from documents.models import DocumentType
+from documents.models import PaperlessTask
 from documents.models import Tag
 from documents.sanity_checker import SanityCheckFailedException
 from documents.sanity_checker import SanityCheckMessages
+from documents.signals.handlers import before_task_publish_handler
+from documents.signals.handlers import task_failure_handler
 from documents.tests.test_classifier import dummy_preprocess
 from documents.tests.utils import DirectoriesMixin
+from documents.tests.utils import DummyProgressManager
 from documents.tests.utils import FileSystemAssertsMixin
+from documents.tests.utils import SampleDirMixin
 
 
 class TestIndexReindex(DirectoriesMixin, TestCase):
@@ -184,3 +195,66 @@ class TestEmptyTrashTask(DirectoriesMixin, FileSystemAssertsMixin, TestCase):
 
         tasks.empty_trash()
         self.assertEqual(Document.global_objects.count(), 0)
+
+
+class TestRetryConsumeTask(
+    DirectoriesMixin,
+    SampleDirMixin,
+    FileSystemAssertsMixin,
+    TestCase,
+):
+    @override_settings(CONSUMPTION_FAILED_DIR=Path(__file__).parent / "samples")
+    @mock.patch("documents.consumer.run_subprocess")
+    def test_retry_consume(self, m):
+        test_file = self.SAMPLE_DIR / "corrupted.pdf"
+        temp_copy = self.dirs.scratch_dir / test_file.name
+        shutil.copy(test_file, temp_copy)
+
+        headers = {
+            "id": str(uuid.uuid4()),
+            "task": "documents.tasks.consume_file",
+        }
+        body = (
+            # args
+            (
+                ConsumableDocument(
+                    source=DocumentSource.ConsumeFolder,
+                    original_file=str(temp_copy),
+                ),
+                None,
+            ),
+            # kwargs
+            {},
+            # celery stuff
+            {"callbacks": None, "errbacks": None, "chain": None, "chord": None},
+        )
+        before_task_publish_handler(headers=headers, body=body)
+
+        with mock.patch("documents.tasks.ProgressManager", DummyProgressManager):
+            with self.assertRaises(Exception):
+                tasks.consume_file(
+                    ConsumableDocument(
+                        source=DocumentSource.ConsumeFolder,
+                        original_file=temp_copy,
+                    ),
+                )
+
+        task_failure_handler(
+            task_id=headers["id"],
+            exception="Example failure",
+        )
+
+        task = PaperlessTask.objects.first()
+        # Ensure the file is moved to the failed dir
+        self.assertIsFile(settings.CONSUMPTION_FAILED_DIR / task.task_file_name)
+
+        tasks.retry_failed_file(task_id=task.task_id)
+
+        m.assert_called_once()
+
+        args, _ = m.call_args
+
+        command = args[0]
+
+        self.assertEqual(command[0], "qpdf")
+        self.assertEqual(command[1], "--replace-input")