feat(tracking): enrich live image pairs and cancel scans safely (#193)

tests/test_pair_image_linking.py

@@ -0,0 +1,55 @@
+"""Live metadata annotation uses the same conservative ingestion as scan."""
+
+import logging
+import tempfile
+import unittest
+from pathlib import Path
+
+from database.operations import PromptDatabase
+from utils.hashing import generate_prompt_hash
+from utils.image_monitor import ImageGenerationHandler
+
+
+def metadata(positive="cat", negative="blur"):
+    return {
+        "prompt": {
+            "1": {"class_type": "PromptManager", "inputs": {"text": positive}},
+            "2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
+            "3": {
+                "class_type": "KSampler",
+                "inputs": {"positive": ["1", 0], "negative": ["2", 0]},
+            },
+        }
+    }
+
+
+class TestPairImageLinking(unittest.TestCase):
+    def setUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        self.db = PromptDatabase(str(Path(self.temp.name) / "db.sqlite"))
+        self.path = str(Path(self.temp.name) / "image.png")
+        self.pid = self.db.save_prompt("cat", prompt_hash=generate_prompt_hash("cat"))
+        self.handler = ImageGenerationHandler.__new__(ImageGenerationHandler)
+        self.handler.db_manager = self.db
+        self.handler.logger = logging.getLogger("test.pairlink")
+
+    def tearDown(self):
+        self.db.close_all()
+        self.temp.cleanup()
+
+    def test_exact_metadata_creates_pair_link_without_counting_a_run(self):
+        self.handler.link_image_to_prompt(self.path, {"id": self.pid}, metadata())
+        self.handler.link_image_to_prompt(self.path, {"id": self.pid}, metadata())
+        pairs = self.db.list_prompt_pairs()["pairs"]
+        self.assertEqual(len(pairs), 1)
+        self.assertEqual(pairs[0]["image_count"], 1)
+        self.assertEqual(pairs[0]["run_count"], 0)
+
+    def test_stale_context_or_ambiguous_metadata_does_not_create_pair(self):
+        meta = metadata(positive="dog")
+        self.handler.link_image_to_prompt(self.path, {"id": self.pid}, meta)
+        self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
+        meta = metadata()
+        meta["prompt"]["4"] = meta["prompt"]["3"]
+        self.handler.link_image_to_prompt(self.path + "2", {"id": self.pid}, meta)
+        self.assertEqual(self.db.list_prompt_pairs()["total"], 0)

tests/test_pair_scan.py

@@ -55,3 +55,29 @@ class TestPairScan(AdminAPITestCase):
         path = self.output_dir / "old.png"
         self.assertEqual(self.api._ingest_scanned_file(path, meta), "added")
         self.assertEqual(self.api._ingest_scanned_file(path, meta), "unchanged")
+
+    async def test_cancel_finishes_at_safe_boundary_and_restart_is_idempotent(self):
+        from py.api.scan_job import ScanJob
+
+        job = ScanJob()
+        self.api._scan_job = job
+        response = await self.client.post("/prompt_manager/scan/cancel")
+        self.assertEqual(response.status, 200)
+        self.assertTrue((await response.json())["cancellation_requested"])
+        counts = self.api._ingest_scan_batch_sync(
+            [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
+        )
+        self.assertEqual(counts["processed"], 0)
+        self.assertEqual(self.api.db.list_prompt_pairs()["total"], 0)
+        job.finish()
+        self.api._scan_job = ScanJob()
+        first = self.api._ingest_scan_batch_sync(
+            [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
+        )
+        self.assertEqual(first["pairs_added"], 1)
+        self.assertEqual(
+            self.api._ingest_scan_batch_sync(
+                [(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
+            )["unchanged"],
+            1,
+        )