Merge branch 'work/332-vocab' into feat/332-image-categories-pr2

tests/test_api_vocabulary.py

@@ -62,12 +62,16 @@ class TestVocabularyAPI(ImageAPITestCase):
             },
         )
         self.assertEqual(
-            categories["categories"], [{"value": "Animal", "prompt_count": 1}]
+            categories["categories"],
+            [{"value": "Animal", "prompt_count": 1, "image_count": 0}],
         )
         _, rest = await self.get(
             "categories", limit=1, cursor=categories["next_cursor"]
         )
-        self.assertEqual(rest["categories"], [{"value": "Pet", "prompt_count": 1}])
+        self.assertEqual(
+            rest["categories"],
+            [{"value": "Pet", "prompt_count": 1, "image_count": 0}],
+        )
         self.assertFalse(rest["has_more"])
 
     async def test_vocabulary_cursor_is_signed_epoch_scoped(self):
@@ -488,6 +492,7 @@ class TestVocabularyAPI(ImageAPITestCase):
             "other_prompts": 0,
             "negative_variants": 1,
             "archived_negative_variants": 0,
+            "images": 0,
         }
         self.assertEqual(
             preview,
@@ -528,6 +533,73 @@ class TestVocabularyAPI(ImageAPITestCase):
         self.assertEqual(response.status, 409)
         self.assertEqual(replay["code"], "vocabulary-conflict")
 
+    async def test_category_rows_report_images_and_image_only_values(self):
+        self._save_prompt("Owner", category="Prompt value")
+        _, options_before = await self.get_options()
+        self.api.db.set_image_categories(
+            [("root", "a.png"), ("root", "b.png")], "Prompt value"
+        )
+        self.api.db.set_image_categories([("root", "c.png")], "Image value")
+        response, first = await self.get("categories", limit=1)
+        self.assertEqual(response.status, 200)
+        self.assertEqual(
+            first["categories"],
+            [{"value": "Prompt value", "prompt_count": 1, "image_count": 2}],
+        )
+        self.assertEqual(first["total"], 2)
+        _, rest = await self.get("categories", limit=1, cursor=first["next_cursor"])
+        self.assertEqual(
+            rest["categories"],
+            [{"value": "Image value", "prompt_count": 0, "image_count": 1}],
+        )
+        self.assertFalse(rest["has_more"])
+        # The Prompt library's category list does not gain image values.
+        _, options_after = await self.get_options()
+        self.assertEqual(options_after, options_before)
+        self.assertEqual(options_after["categories"], ["Prompt value"])
+
+    async def get_options(self):
+        response = await self.client.get("/prompt_manager/library/options")
+        return response, await response.json()
+
+    async def test_category_routes_rename_image_only_value(self):
+        epoch = self.api.db.get_execution_epoch()
+        self._save_prompt("Bystander", category="Kept")
+        self.api.db.set_image_categories(
+            [("root", "a.png"), ("second", "b.png")], "Image value"
+        )
+        prompts_before = self.category_rows()
+        body = {"expected_epoch": epoch, "source": "Image value", "target": "Moved"}
+        response, preview = await self.send("POST", "categories/preview", body)
+        self.assertEqual(response.status, 200, preview)
+        impact = {
+            "positive_prompts": 0,
+            "other_prompts": 0,
+            "negative_variants": 0,
+            "archived_negative_variants": 0,
+            "images": 2,
+        }
+        self.assertEqual(preview["impact"], impact)
+        stale = preview["impact_token"]
+        self.api.db.set_image_categories([("root", "late.png")], "Image value")
+        response, refused = await self.send(
+            "POST", "categories/apply", {**body, "impact_token": stale}
+        )
+        self.assertEqual(response.status, 409)
+        self.assertEqual(refused["code"], "vocabulary-conflict")
+        self.assertEqual(self.api.db.image_category("root", "a.png"), "Image value")
+        _, preview = await self.send("POST", "categories/preview", body)
+        response, applied = await self.send(
+            "POST",
+            "categories/apply",
+            {**body, "impact_token": preview["impact_token"]},
+        )
+        self.assertEqual(response.status, 200, applied)
+        self.assertEqual(applied["impact"], {**impact, "images": 3})
+        for key in (("root", "a.png"), ("second", "b.png"), ("root", "late.png")):
+            self.assertEqual(self.api.db.image_category(*key), "Moved")
+        self.assertEqual(self.category_rows(), prompts_before)
+
     async def test_preview_validation(self):
         epoch = self.api.db.get_execution_epoch()
         self._save_prompt("Owner", category="Scenery")

tests/test_vocabulary.py

@@ -13,6 +13,11 @@ from database.library_commands import LibraryCommandError
 from database.operations import PromptDatabase
 from utils.hashing import generate_prompt_hash
 
+try:
+    from test_image_categories import category_checksum
+except ImportError:
+    from tests.test_image_categories import category_checksum
+
 
 class TestVocabularyReads(unittest.TestCase):
     def setUp(self):
@@ -168,7 +173,8 @@ class TestVocabularyReads(unittest.TestCase):
         self.assertEqual(page["total"], 4)
         self.assertFalse(page["has_more"])
         for row in page["categories"]:
-            self.assertEqual(set(row), {"value", "prompt_count"})
+            self.assertEqual(set(row), {"value", "prompt_count", "image_count"})
+            self.assertEqual(row["image_count"], 0)
             self.assertEqual(
                 row["prompt_count"],
                 self.db.list_library(category=row["value"])["total"],
@@ -238,6 +244,111 @@ class TestVocabularyReads(unittest.TestCase):
             with self.assertRaises(ValueError):
                 self.read("list_vocabulary_tags", **options)
 
+    def own(self, value, *paths, root="root"):
+        """Give each relative path under ``root`` the own category ``value``."""
+        self.db.set_image_categories([(root, path) for path in paths], value)
+
+    def all_categories(self, limit=200):
+        rows, pages, after = [], [], None
+        while True:
+            page = self.read("list_vocabulary_categories", limit=limit, after=after)
+            pages.append(page)
+            rows.extend(page["categories"])
+            if not page["has_more"]:
+                return rows, pages
+            after = page["next_key"]
+
+    def test_category_list_includes_image_only_values_last(self):
+        self.save("Oldest", category="Zeta")
+        self.save("Newer", category="Alpha")
+        self.own("Zeta", "a.png", "b.png")
+        self.own("Only images B", "c.png")
+        self.own("Only images A", "d.png", "e.png", "gone.png")
+        rows, _pages = self.all_categories()
+        self.assertEqual(
+            rows,
+            [
+                {"value": "Zeta", "prompt_count": 1, "image_count": 2},
+                {"value": "Alpha", "prompt_count": 1, "image_count": 0},
+                {"value": "Only images A", "prompt_count": 0, "image_count": 3},
+                {"value": "Only images B", "prompt_count": 0, "image_count": 1},
+            ],
+        )
+        page = self.read("list_vocabulary_categories")
+        self.assertEqual(page["total"], 4)
+
+    def test_category_cursor_crosses_into_image_only_values(self):
+        for index in range(3):
+            self.save(f"Owner {index}", category=f"P{index}")
+        self.own("P1", "shared.png")
+        self.own("I-b", "x.png")
+        self.own("I-a", "y.png")
+        self.own("I-c", "z.png", root="other")
+        for limit in (1, 2, 3, 4):
+            with self.subTest(limit=limit):
+                rows, pages = self.all_categories(limit=limit)
+                self.assertEqual(
+                    [row["value"] for row in rows],
+                    ["P0", "P1", "P2", "I-a", "I-b", "I-c"],
+                )
+                self.assertEqual({page["total"] for page in pages}, {6})
+                self.assertEqual(
+                    [row["image_count"] for row in rows], [0, 1, 0, 1, 1, 1]
+                )
+        first = self.read("list_vocabulary_categories", limit=3)
+        self.assertEqual(first["next_key"], [3, "P2"])
+        second = self.read("list_vocabulary_categories", limit=1, after=[3, "P2"])
+        self.assertEqual(second["next_key"], [2**63 - 1, "I-a"])
+
+    def test_image_rows_on_a_prompt_value_count_once(self):
+        self.save("Owner", category="Shared")
+        self.db.save_prompt_pair("Owner", "Noise", category="Shared")
+        self.own("Shared", *[f"{n}.png" for n in range(5)])
+        self.own("Shared", "elsewhere.png", root="second")
+        rows, _pages = self.all_categories()
+        self.assertEqual(
+            rows, [{"value": "Shared", "prompt_count": 1, "image_count": 6}]
+        )
+
+    def test_library_options_ignore_image_categories(self):
+        self.save("Owner", category="Prompt value")
+        before = self.db.get_library_options()
+        self.own("Image value", "a.png")
+        self.own("Prompt value", "b.png")
+        self.assertEqual(self.db.get_library_options(), before)
+        self.assertEqual(before["categories"], ["Prompt value"])
+
+    def test_category_reads_use_the_image_category_index(self):
+        from database import vocabulary
+
+        self.own("Value", "a.png")
+        statements = [
+            vocabulary._CATEGORY_ROWS + "SELECT COUNT(*) FROM counted",
+            vocabulary._CATEGORY_ROWS
+            + "SELECT value,first_pid,prompt_count,image_count FROM counted "
+            "WHERE first_pid>? OR (first_pid=? AND value>?) "
+            "ORDER BY first_pid,value LIMIT ?",
+        ]
+        plans = []
+        with self.db.model.get_connection() as conn:
+            for sql in statements:
+                params = [1] * sql.count("?")
+                plans.extend(
+                    row[-1] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)
+                )
+            for sql in (
+                vocabulary._IMAGE_CATEGORY_ROWIDS,
+                *vocabulary._CATEGORY_UPDATES,
+            ):
+                params = [1] * sql.count("?")
+                plans.extend(
+                    row[-1] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)
+                )
+        image_steps = [step for step in plans if "image_categories" in step]
+        self.assertTrue(image_steps)
+        for step in image_steps:
+            self.assertIn("idx_image_categories_category", step, step)
+
     def test_reads_share_one_snapshot_and_report_epoch(self):
         self.save("Cat", category="Animal", tags=["cat"])
         epoch = self.db.get_execution_epoch()
@@ -741,6 +852,7 @@ class TestVocabularyCategoryCommands(unittest.TestCase):
                 "other_prompts": 1,
                 "negative_variants": 2,
                 "archived_negative_variants": 1,
+                "images": 0,
             },
         )
         self.assertFalse(result["target_exists"])
@@ -810,6 +922,7 @@ class TestVocabularyCategoryCommands(unittest.TestCase):
                 "other_prompts": 0,
                 "negative_variants": 0,
                 "archived_negative_variants": 0,
+                "images": 0,
             },
         )
         self.assertEqual(self.rows("prompts")[variant]["category"], "Other")
@@ -902,6 +1015,171 @@ class TestVocabularyCategoryCommands(unittest.TestCase):
                 self.apply(bad)
         self.assertEqual(self.snapshot(), before)
 
+    def own(self, value, *paths, root="root"):
+        self.db.set_image_categories([(root, path) for path in paths], value)
+
+    def image_rows(self):
+        with self.db.model.get_connection() as conn:
+            return {
+                (row[0], row[1]): (row[2], row[3], row[4])
+                for row in conn.execute(
+                    "SELECT root_key,relative_path,category,updated_at,rowid "
+                    "FROM image_categories ORDER BY rowid"
+                )
+            }
+
+    def checksum(self):
+        with self.db.model.get_connection() as conn:
+            return category_checksum(conn)
+
+    def test_preview_counts_image_rows(self):
+        self.owners()
+        self.own("Scenery", "a.png", "b.png", "missing.png")
+        self.own("Scenery", "c.png", root="second")
+        self.own("Other", "d.png")
+        before = self.image_rows()
+        result = self.preview()
+        self.assertEqual(result["impact"]["images"], 4)
+        self.assertEqual(self.image_rows(), before)
+        self.assertTrue(self.preview(target="Other")["target_exists"])
+
+    def test_apply_renames_image_categories(self):
+        self.owners()
+        bystander = self.save("Bystander", category="Other")
+        self.own("Scenery", "a.png", "b.png")
+        self.own("Scenery", "c.png", root="second")
+        self.own(" Scenery ", "padded.png")
+        self.own("Other", "d.png")
+        self.age_rows()
+        with self.db.model.get_connection() as conn:
+            conn.execute("UPDATE image_categories SET updated_at=?", (_OLD_STAMP,))
+        images_before = self.image_rows()
+        prompts_before, pairs_before = self.snapshot()
+        preview = self.preview()
+        receipt = self.apply(preview["impact_token"])
+        self.assertEqual(receipt["impact"]["images"], 3)
+        images_after = self.image_rows()
+        self.assertEqual(set(images_after), set(images_before))
+        renamed = {("root", "a.png"), ("root", "b.png"), ("second", "c.png")}
+        for key, (value, stamp, rowid) in images_after.items():
+            old_value, old_stamp, old_rowid = images_before[key]
+            self.assertEqual(rowid, old_rowid, key)
+            if key in renamed:
+                self.assertEqual(value, "Landscape")
+                self.assertNotEqual(stamp, _OLD_STAMP)
+            else:
+                self.assertEqual((value, stamp), (old_value, old_stamp))
+        prompts_after, pairs_after = self.snapshot()
+        self.assertEqual(prompts_after[bystander], prompts_before[bystander])
+        for before, after in (
+            (prompts_before, prompts_after),
+            (pairs_before, pairs_after),
+        ):
+            for row_id, row in after.items():
+                if before[row_id]["category"] != "Scenery":
+                    self.assertEqual(row, before[row_id])
+        # Merging into a value images already hold keeps every file key.
+        merge = self.preview(source="Landscape", target="Other")
+        self.assertTrue(merge["target_exists"])
+        self.assertEqual(merge["impact"]["images"], 3)
+        receipt = self.apply(merge["impact_token"], source="Landscape", target="Other")
+        self.assertTrue(receipt["consolidated"])
+        merged = self.image_rows()
+        self.assertEqual(set(merged), set(images_before))
+        self.assertEqual(
+            sorted(key for key, row in merged.items() if row[0] == "Other"),
+            sorted([*renamed, ("root", "d.png")]),
+        )
+
+    def test_image_only_value_can_be_renamed(self):
+        self.save("Bystander", category="Kept")
+        self.db.save_prompt_pair("Pair", "Noise", category="Kept")
+        self.own("Image only", "a.png", "b.png")
+        self.age_rows()
+        before = self.checksum()
+        preview = self.preview(source="Image only", target="Renamed")
+        self.assertEqual(
+            preview["impact"],
+            {
+                "positive_prompts": 0,
+                "other_prompts": 0,
+                "negative_variants": 0,
+                "archived_negative_variants": 0,
+                "images": 2,
+            },
+        )
+        self.assertFalse(preview["target_exists"])
+        receipt = self.apply(
+            preview["impact_token"], source="Image only", target="Renamed"
+        )
+        self.assertFalse(receipt["consolidated"])
+        self.assertEqual({row[0] for row in self.image_rows().values()}, {"Renamed"})
+        self.assertEqual(self.checksum(), before)
+        # An image-only target counts as existing.
+        self.own("Image target", "c.png")
+        self.assertTrue(
+            self.preview(source="Renamed", target="Image target")["target_exists"]
+        )
+        # Nothing holds the value any more: refused.
+        with self.assertRaises(LibraryCommandError) as caught:
+            self.preview(source="Image only", target="Again")
+        self.assertEqual(caught.exception.code, "vocabulary-conflict")
+
+    def test_token_refused_after_an_image_row_gains_or_loses_the_value(self):
+        self.owners()
+        self.own("Scenery", "a.png", "b.png")
+
+        def gain():
+            self.own("Scenery", "late.png")
+
+        def lose():
+            self.db.clear_image_categories([("root", "a.png")])
+
+        def move_away():
+            self.own("Elsewhere", "b.png")
+
+        def target_on_image():
+            self.own("Landscape", "target.png")
+
+        for mutate in (gain, lose, move_away, target_on_image):
+            with self.subTest(mutate.__name__):
+                token = self.preview()["impact_token"]
+                mutate()
+                before = self.snapshot(), self.image_rows()
+                with self.assertRaises(LibraryCommandError) as caught:
+                    self.apply(token)
+                self.assertEqual(caught.exception.code, "vocabulary-conflict")
+                self.assertEqual((self.snapshot(), self.image_rows()), before)
+
+    def test_image_category_failure_rolls_back_prompts(self):
+        self.owners()
+        self.own("Scenery", "a.png")
+        token = self.preview()["impact_token"]
+        with self.db.model.get_connection() as conn:
+            conn.execute(
+                "CREATE TRIGGER abort_images BEFORE UPDATE ON image_categories "
+                "BEGIN SELECT RAISE(ABORT, 'boom'); END"
+            )
+        before = self.snapshot(), self.image_rows()
+        with self.assertRaises(sqlite3.Error):
+            self.apply(token)
+        self.assertEqual((self.snapshot(), self.image_rows()), before)
+
+    def test_category_exists_reads_the_image_index(self):
+        from database import vocabulary
+
+        with self.db.model.get_connection() as conn:
+            plan = [
+                row[-1]
+                for row in conn.execute(
+                    "EXPLAIN QUERY PLAN " + vocabulary._CATEGORY_EXISTS, {"v": "x"}
+                )
+            ]
+        steps = [step for step in plan if "image_categories" in step]
+        self.assertTrue(steps, plan)
+        for step in steps:
+            self.assertIn("idx_image_categories_category", step)
+
     def test_category_apply_failure_rolls_back(self):
         self.owners()
         token = self.preview()["impact_token"]
@@ -942,6 +1220,18 @@ class TestSharedTokenHashing(unittest.TestCase):
             ),
             "cefe7f6a5544b7f6a203f5527810899c6af568aa7abc1b75adffbfb5ae6251ae",
         )
+        with_images = category_impact_token(
+            "e", "S", "T", True, [(1, True)], [], images=[9, 4]
+        )
+        self.assertRegex(with_images, r"\A[0-9a-f]{64}\Z")
+        self.assertEqual(
+            with_images,
+            category_impact_token("e", "S", "T", True, [(1, True)], [], images=[4, 9]),
+        )
+        self.assertNotEqual(
+            with_images,
+            category_impact_token("e", "S", "T", True, [(1, True)], [], images=[4]),
+        )
 
 
 if __name__ == "__main__":