Upload updates

backend/apps/library/tests/test_e621.py

@@ -0,0 +1,99 @@
+"""The server-side e621 client: batching, retries and IQDB stream handling."""
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest import mock
+
+from django.test import SimpleTestCase
+
+from apps.library import e621
+
+
+class FakeResponse:
+    def __init__(self, status_code, payload=None, headers=None):
+        self.status_code = status_code
+        self._payload = payload
+        self.headers = headers or {}
+
+    def json(self):
+        if self._payload is None:
+            raise ValueError("not json")
+        return self._payload
+
+
+class E621ClientTests(SimpleTestCase):
+    def test_iqdb_search_reopens_the_file_on_retry(self):
+        bodies = []
+
+        def fake_request(method, url, **kwargs):
+            handle = kwargs["files"]["search[file]"][1]
+            bodies.append(handle.read())
+            if len(bodies) == 1:
+                return FakeResponse(
+                    429, {"success": False, "message": "Throttled"}
+                )
+            return FakeResponse(200, [{"post_id": 1, "score": 90.0}])
+
+        with TemporaryDirectory() as tmp:
+            path = Path(tmp) / "x.png"
+            path.write_bytes(b"image-bytes")
+            with mock.patch.object(
+                e621.requests, "request", side_effect=fake_request
+            ), mock.patch.object(e621.time, "sleep"), mock.patch.object(
+                e621, "_wait_for_slot"
+            ):
+                result = e621.iqdb_search(None, path)
+
+        # Both attempts must send the full body, not the consumed handle.
+        self.assertEqual(bodies, [b"image-bytes", b"image-bytes"])
+        self.assertEqual(result, [{"post_id": 1, "score": 90.0}])
+
+    def test_check_md5_batch_keys_by_md5(self):
+        payload = {
+            "posts": [
+                {"id": 5, "file": {"md5": "a" * 32}},
+                {"id": 6, "file": {"md5": "b" * 32}},
+            ]
+        }
+        with mock.patch.object(e621, "get", return_value=payload) as getter:
+            found = e621.check_md5_batch(None, ["A" * 32, "b" * 32])
+        self.assertEqual(set(found), {"a" * 32, "b" * 32})
+        params = getter.call_args.kwargs["params"]
+        self.assertTrue(params["tags"].startswith("md5:"))
+        self.assertEqual(params["limit"], 2)
+
+    def test_auth_errors_are_not_retried(self):
+        calls = []
+
+        def fake_request(*args, **kwargs):
+            calls.append(1)
+            return FakeResponse(403, {"error": "nope"})
+
+        with mock.patch.object(
+            e621.requests, "request", side_effect=fake_request
+        ), mock.patch.object(e621, "_wait_for_slot"):
+            with self.assertRaises(e621.E621AuthError):
+                e621._request(None, "GET", "/posts.json", require_auth=False)
+        self.assertEqual(len(calls), 1)
+
+    def test_load_shedding_html_raises_rate_limited_after_retries(self):
+        with mock.patch.object(
+            e621.requests,
+            "request",
+            return_value=FakeResponse(200, None),
+        ), mock.patch.object(e621.time, "sleep"), mock.patch.object(
+            e621, "_wait_for_slot"
+        ):
+            with self.assertRaises(e621.E621RateLimited):
+                e621._request(
+                    None, "GET", "/posts.json", require_auth=False, attempts=2
+                )
+
+    def test_fetch_posts_by_ids_queries_with_id_tag(self):
+        with mock.patch.object(
+            e621, "get", return_value={"posts": [{"id": 9}]}
+        ) as getter:
+            posts = e621.fetch_posts_by_ids(None, [9])
+        self.assertEqual(posts, [{"id": 9}])
+        params = getter.call_args.kwargs["params"]
+        self.assertEqual(params["tags"], "id:9")

backend/apps/library/tests/test_uploads.py

@@ -6,17 +6,21 @@ import io
 import json
 import shutil
 import tempfile
+from datetime import timedelta
 from pathlib import Path
+from unittest import mock
 
 from django.contrib.auth import get_user_model
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.test import Client, TestCase, override_settings
+from django.utils import timezone
 
 from PIL import Image
 
 from rest_framework.authtoken.models import Token
 
-from apps.library.models import MediaItem, TempUpload
+from apps.library import upload_pipeline
+from apps.library.models import MediaItem, TempUpload, UploadRun
 
 User = get_user_model()
 
@@ -95,6 +99,7 @@ class TempUploadListTests(TestCase):
         self.assertEqual(Client().get("/api/uploads/").status_code, 401)
 
 
+@override_settings(UPLOAD_PIPELINE_AUTOSTART=False)
 class StagedUploadWorkflowTests(TestCase):
     @classmethod
     def setUpClass(cls):
@@ -400,3 +405,293 @@ class IqdbRecordingTests(TestCase):
             self.client, f"/api/uploads/{temp.id}/iqdb/", {"results": "nope"}
         )
         self.assertEqual(response.status_code, 400)
+
+
+@override_settings(UPLOAD_PIPELINE_AUTOSTART=False)
+class UploadPipelineTests(TestCase):
+    """The server-side MD5 -> visual -> IQDB queue and its board API."""
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls._tmp = tempfile.mkdtemp(prefix="j621-pipeline-")
+        cls._watched = Path(cls._tmp) / "library"
+        cls._watched.mkdir(parents=True, exist_ok=True)
+        cls._settings = override_settings(
+            MEDIA_ROOT=cls._tmp, WATCHED_FOLDER=str(cls._watched)
+        )
+        cls._settings.enable()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls._settings.disable()
+        shutil.rmtree(cls._tmp, ignore_errors=True)
+        super().tearDownClass()
+
+    def setUp(self):
+        self.uploader = User.objects.create_user(
+            username="pipe-uploader", password="pipe-pass-123456"
+        )
+        self.uploader.role = "uploader"
+        self.uploader.save(update_fields=["role"])
+        self.other = User.objects.create_user(
+            username="pipe-other", password="pipe-pass-123456"
+        )
+        self.other.role = "uploader"
+        self.other.save(update_fields=["role"])
+
+    def api_client(self, user):
+        client = Client()
+        client.defaults["HTTP_AUTHORIZATION"] = (
+            f"Token {Token.objects.create(user=user).key}"
+        )
+        return client
+
+    def stage(self, user=None, label="file", payload=None, filename=None):
+        user = user or self.uploader
+        payload = payload or TINY_PNG
+        name = filename or f"{label}.png"
+        return TempUpload.objects.create(
+            user=user,
+            file=SimpleUploadedFile(name, payload, content_type="image/png"),
+            original_filename=name,
+            md5=hashlib.md5(payload + label.encode()).hexdigest(),
+            size=len(payload),
+        )
+
+    def test_md5_match_auto_imports_the_file(self):
+        temp = self.stage(label="match")
+        post = {
+            "id": 123456,
+            "rating": "s",
+            "file": {
+                "md5": temp.md5,
+                "url": "https://static1.e621.net/data/m.png",
+            },
+            "tags": {"general": ["canine"]},
+        }
+        with mock.patch.object(
+            upload_pipeline.e621,
+            "check_md5_batch",
+            return_value={temp.md5: post},
+        ):
+            upload_pipeline.run_pipeline(self.uploader.id)
+
+        temp.refresh_from_db()
+        self.assertEqual(temp.status, TempUpload.STATUS_COMPLETED)
+        self.assertEqual(temp.resolution, TempUpload.RESOLUTION_AUTO_MD5)
+        self.assertEqual(temp.e621_post_id, 123456)
+        self.assertIsNotNone(temp.library_item_id)
+        self.assertIsNotNone(temp.e621_checked_at)
+        self.assertIsNone(temp.claimed_at)
+        run = UploadRun.objects.get(user=self.uploader)
+        self.assertEqual(run.status, UploadRun.STATUS_IDLE)
+        self.assertEqual(run.matched, 1)
+        self.assertEqual(run.processed, 1)
+
+    def test_unmatched_file_runs_every_phase(self):
+        temp = self.stage(label="nomatch")
+        raw_iqdb = [
+            {
+                "post_id": 777,
+                "score": 91.0,
+                "post": {
+                    "id": 777,
+                    "rating": "q",
+                    "md5": "a" * 32,
+                    "score": 5,
+                    "fav_count": 2,
+                    "image_width": 800,
+                    "image_height": 600,
+                },
+            }
+        ]
+        modern = [
+            {
+                "id": 777,
+                "rating": "q",
+                "fav_count": 4,
+                "score": {"total": 9},
+                "preview": {"url": "https://static1.e621.net/data/preview/x.jpg"},
+                "file": {"md5": "a" * 32, "width": 801, "height": 601},
+                "tags": {"general": ["canine", "solo"]},
+            }
+        ]
+        with mock.patch.object(
+            upload_pipeline.e621, "check_md5_batch", return_value={}
+        ), mock.patch.object(
+            upload_pipeline.e621, "iqdb_search", return_value=raw_iqdb
+        ), mock.patch.object(
+            upload_pipeline.e621, "fetch_posts_by_ids", return_value=modern
+        ):
+            upload_pipeline.run_pipeline(self.uploader.id)
+
+        temp.refresh_from_db()
+        self.assertIsNotNone(temp.e621_checked_at)
+        self.assertIsNotNone(temp.visual_checked_at)
+        self.assertEqual(temp.status, TempUpload.STATUS_VISUAL_MATCH)
+        self.assertEqual(len(temp.iqdb_data), 1)
+        candidate = temp.iqdb_data[0]
+        self.assertEqual(candidate["post_id"], 777)
+        self.assertEqual(candidate["score_total"], 9)
+        self.assertEqual(candidate["fav_count"], 4)
+        self.assertEqual(candidate["width"], 801)
+        self.assertEqual(
+            candidate["preview_url"], "https://static1.e621.net/data/preview/x.jpg"
+        )
+        self.assertEqual(candidate["tags_preview"], ["canine", "solo"])
+        run = UploadRun.objects.get(user=self.uploader)
+        self.assertEqual(run.status, UploadRun.STATUS_IDLE)
+        self.assertEqual(run.processed, 1)
+
+    def test_visual_match_flags_similar_library_items(self):
+        from apps.library.uploads import complete_temp_upload
+
+        seed = self.stage(label="seed")
+        complete_temp_upload(seed)
+        self.assertTrue(MediaItem.objects.exists())
+
+        temp = self.stage(label="similar")
+        with mock.patch.object(
+            upload_pipeline.e621, "check_md5_batch", return_value={}
+        ), mock.patch.object(
+            upload_pipeline.e621, "iqdb_search", return_value=[]
+        ):
+            upload_pipeline.run_pipeline(self.uploader.id)
+
+        temp.refresh_from_db()
+        self.assertGreaterEqual(len(temp.visual_matches or []), 1)
+        self.assertEqual(temp.status, TempUpload.STATUS_VISUAL_MATCH)
+
+    def test_iqdb_rate_limit_pauses_the_run(self):
+        temp = self.stage(label="throttled")
+        with mock.patch.object(
+            upload_pipeline.e621, "check_md5_batch", return_value={}
+        ), mock.patch.object(
+            upload_pipeline.e621,
+            "iqdb_search",
+            side_effect=upload_pipeline.e621.E621RateLimited("Throttled"),
+        ):
+            upload_pipeline.run_pipeline(self.uploader.id)
+
+        run = UploadRun.objects.get(user=self.uploader)
+        self.assertEqual(run.status, UploadRun.STATUS_PAUSED)
+        self.assertIn("e621", run.error)
+        temp.refresh_from_db()
+        self.assertEqual(temp.attempts, 1)
+        self.assertNotEqual(temp.pipeline_error, "")
+        self.assertIsNone(temp.iqdb_data)
+        self.assertIsNone(temp.claimed_at)
+
+        client = self.api_client(self.uploader)
+        response = jpost(client, f"/api/uploads/{temp.id}/retry/", {})
+        self.assertEqual(response.status_code, 200)
+        temp.refresh_from_db()
+        self.assertEqual(temp.attempts, 0)
+        self.assertEqual(temp.pipeline_error, "")
+
+    def test_failed_rows_stop_after_the_attempt_cap(self):
+        temp = self.stage(label="broken")
+        for _ in range(upload_pipeline.MAX_ATTEMPTS):
+            with mock.patch.object(
+                upload_pipeline.e621, "check_md5_batch", return_value={}
+            ), mock.patch.object(
+                upload_pipeline.e621,
+                "iqdb_search",
+                side_effect=upload_pipeline.e621.E621Error("boom"),
+            ):
+                upload_pipeline.run_pipeline(self.uploader.id)
+        temp.refresh_from_db()
+        self.assertEqual(temp.attempts, upload_pipeline.MAX_ATTEMPTS)
+        self.assertEqual(upload_pipeline.count_outstanding(self.uploader), 0)
+        run = UploadRun.objects.get(user=self.uploader)
+        self.assertGreaterEqual(run.failed, 1)
+
+    def test_status_process_and_compact_board_payload(self):
+        temp = self.stage(label="board")
+        client = self.api_client(self.uploader)
+
+        status = client.get("/api/uploads/status/").json()
+        self.assertEqual(status["status"], UploadRun.STATUS_IDLE)
+        self.assertTrue(status["active"])
+        self.assertEqual(status["outstanding"], 1)
+        self.assertEqual(status["waiting"]["md5"], 1)
+
+        self.assertEqual(client.post("/api/uploads/process/").status_code, 200)
+
+        rows = client.get("/api/uploads/").json()
+        self.assertEqual(len(rows), 1)
+        row = rows[0]
+        for key in (
+            "md5_checked",
+            "visual_checked",
+            "iqdb_checked",
+            "processing",
+            "similar_count",
+            "pipeline_error",
+        ):
+            self.assertIn(key, row)
+        self.assertNotIn("e621_data", row)
+        self.assertNotIn("iqdb_data", row)
+
+        detail = client.get(f"/api/uploads/{temp.id}/").json()
+        self.assertIn("e621_data", detail)
+        self.assertIn("iqdb_data", detail)
+
+    def test_discard_bulk_removes_only_own_rows(self):
+        client = self.api_client(self.uploader)
+        first = self.stage(label="discard-a")
+        second = self.stage(label="discard-b")
+        theirs = self.stage(self.other, label="discard-theirs")
+        paths = [Path(first.file.path), Path(second.file.path)]
+
+        response = jpost(
+            client,
+            "/api/uploads/discard-bulk/",
+            {"temp_ids": [str(first.id), str(second.id), str(theirs.id)]},
+        )
+        self.assertEqual(response.status_code, 200)
+        body = response.json()
+        self.assertEqual(len(body["discarded"]), 2)
+        self.assertEqual(len(body["errors"]), 1)
+        self.assertEqual(body["errors"][0]["error"], "not found")
+        for path in paths:
+            self.assertFalse(path.exists())
+        self.assertFalse(
+            TempUpload.objects.filter(pk__in=[first.id, second.id]).exists()
+        )
+        self.assertTrue(TempUpload.objects.filter(pk=theirs.id).exists())
+
+    def test_retry_with_phase_rechecks_iqdb(self):
+        temp = self.stage(label="recheck")
+        TempUpload.objects.filter(pk=temp.pk).update(
+            iqdb_data=[{"post_id": 1}],
+            status=TempUpload.STATUS_VISUAL_MATCH,
+        )
+        client = self.api_client(self.uploader)
+        response = jpost(
+            client, f"/api/uploads/{temp.id}/retry/", {"phase": "iqdb"}
+        )
+        self.assertEqual(response.status_code, 200)
+        temp.refresh_from_db()
+        self.assertIsNone(temp.iqdb_data)
+        self.assertEqual(temp.status, TempUpload.STATUS_VISUAL_MATCH)
+
+    def test_stale_claims_are_released(self):
+        temp = self.stage(label="stale")
+        TempUpload.objects.filter(pk=temp.pk).update(
+            claimed_at=timezone.now() - upload_pipeline.STALE_CLAIM_AFTER
+            - timedelta(minutes=1)
+        )
+        UploadRun.objects.create(user=self.uploader, status=UploadRun.STATUS_RUNNING)
+        UploadRun.objects.filter(user=self.uploader).update(
+            updated_at=timezone.now() - upload_pipeline.STALE_RUN_AFTER
+            - timedelta(minutes=1)
+        )
+        released, paused = upload_pipeline.reap_stale_claims()
+        self.assertEqual(released, 1)
+        self.assertEqual(paused, 1)
+        temp.refresh_from_db()
+        self.assertIsNone(temp.claimed_at)
+        run = UploadRun.objects.get(user=self.uploader)
+        self.assertEqual(run.status, UploadRun.STATUS_PAUSED)