nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
b3e4855abb6dcd7965a0911e9b622295fd39bed9
ComfyAssets-ComfyUI_PromptM…/tests/test_vocabulary.py
T
Vito cdb771d722 feat(images): filter, inspect and count by each image's effective category (#332, PR 2) (#406) 
* feat(images): store images' own categories in the library database (#332)

One new table, image_categories, and its category index, created empty
by a forward-only migration that reads and writes only its own objects.
Startup checks the structure; restore also checks the values. The store
sets, clears and reads own categories by file key (the deepest output
source holding the file and the path below it), and copies a removed
duplicate's category to a kept copy that has none.

* feat(images): duplicate cleanup keeps images' own categories (#332)

The kept copy's own category wins; a kept copy without one receives the
removed copy's, keyed by the deepest output source holding each file.
The removed copy's row is never deleted, so a give-back or a failed
delete loses nothing. The job log reports how many were kept.

* feat(images): set, remove and list images' own categories (#332)

PUT /prompt_manager/images/category sets the own category of up to 500
listed files, DELETE removes it (listed or not), and
GET /prompt_manager/images/categories lists the values whose file is
listed. Targets are the Images listing identity of an output source,
checked by media_target (shared with Collections). The routes never
touch prompts, recipes or links; POST /prompt_manager/bulk/category is
unchanged.

* feat(vocabulary): count and rename image categories (#332)

Tags & categories lists every category value with image_count, the
stored image_categories rows holding it. Values only images hold follow
every prompt-backed value (first_pid at SQLite's INTEGER maximum), so
keyset cursors cross the boundary. /library/options keeps its own query
and never lists image-only values.

Manage category counts image rows (impact.images), treats an image-held
target as existing, renames image rows in the same transaction, allows
renaming an image-only value, and hashes the sorted image rowids into
the impact token so a change between preview and apply is refused with
vocabulary-conflict. Every image-category statement reads
idx_image_categories_category.

* feat(images): the category filter matches each image's effective category (#332)

An image's own category wins; without one it inherits its link's: the
recipe's category when set, else the positive prompt's (MC4, one rule
with the inspector chip). Ambiguous and unlinked images inherit nothing,
so an unlinked image with its own category is now listed. The link read
no longer filters by category in SQL; the page filter, select all N
matching and saved views compare the effective value. A failed
own-category read answers the existing listing-failed 500.

* fix(images): key symlinked output folders and keep key work off the event loop (#332)

Review fixes: a source matches by its configured spelling too, so a
detected output folder that is a symlink keys its files (routes and
duplicate cleanup). File keys are computed in the executor. A path no
source holds is a plain 403, a lone surrogate a 400. The install repairs
a table that lacks its category index instead of refusing to open.

* feat(inspector): show an image's own category in the inspector and History (#332)

The inspector's library facts gain category_scope "image": an image's
own category overrides the inherited one for any link status (linked,
ambiguous, none). Images reads it for output sources only; History reads
it once per inspection inside its snapshot, keyed by the same file key
from the symlink-resolved output folder. A failed read hides the chip
rather than show a possibly wrong value. The chip's hidden prefix names
the origin ("Image category: "); the visible origin text comes in PR 3.

* perf(images): keep the category filter's link read narrow; record D141, D142, D145 (#332)

Under a category the link read keeps only links whose raw recipe or
positive category is the value (Python still decides the inherited
one), the links of images whose own category is the value are read by
path, and plain reads skip the inherited-category work. The own-category
map is keyed by normalized path. TESTING.md records D141, D142 and D145
and amends D13, D36 and D39; the filter budget seeds 2,000 own
categories.

* feat(images): set or remove a category on all N matching images (#332)

The category routes also accept the Images listing filter of "Select
all N matching" (#392) instead of explicit items, resolved now and
capped at 500 images like an explicit selection.

* fix(images): Windows path keys for own categories; plan check for the category read (#332)

Review fixes: on Windows the own-category map keys each file with
normalize_image_path (abspath drops trailing dots and spaces), and the
plan test also covers the category's narrowing clause.

* test(vendor): count the #332 host case in the execution matrix

image-attribution.spec.mjs gains the image category route case, so the
execution config lists 92 tests (3 attribution cases per project).

* test(legacy): record the image category routes in the legacy contract snapshot (#332)

#332 adds PUT/DELETE /prompt_manager/images/category and
GET /prompt_manager/images/categories; the legacy pages' routes are
unchanged.
2026-10-28 17:19:11 +00:00

1239 lines
50 KiB
Python
Raw Blame History

"""Bounded vocabulary reads count exactly what Library Browse shows."""
import hashlib
import os
import re
import sqlite3
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
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):
self.temp = tempfile.TemporaryDirectory()
self.db = PromptDatabase(str(Path(self.temp.name) / "vocabulary.db"))
def tearDown(self):
try:
# Windows refuses to remove a database that still has open handles.
self.db.close_all()
finally:
self.temp.cleanup()
def save(self, text, **metadata):
return self.db.save_prompt(
text, prompt_hash=generate_prompt_hash(text), **metadata
)
def read(self, name, **options):
method = getattr(self.db, name, None)
self.assertTrue(callable(method), f"PromptDatabase.{name} is missing")
return method(**options)
def add_unused(self, *names):
with self.db.model.get_connection() as conn:
for name in names:
conn.execute("INSERT INTO tags(name) VALUES(?)", (name,))
def all_tags(self):
rows, after = [], None
while True:
page = self.read("list_vocabulary_tags", limit=200, after=after)
rows.extend(page["tags"])
if not page["has_more"]:
return rows
after = page["next_key"]
def test_tag_rows_are_registry_rows_in_creation_order_with_unused_at_zero(self):
self.save("First", tags=["c"])
self.save("Second", tags=["a"])
self.save("Third", tags=["b", "a"])
self.add_unused("z")
page = self.read("list_vocabulary_tags")
self.assertEqual([row["name"] for row in page["tags"]], ["c", "a", "b", "z"])
self.assertEqual(page["total"], 4)
self.assertFalse(page["has_more"])
self.assertIsNone(page["next_key"])
by_name = {row["name"]: row for row in page["tags"]}
self.assertEqual(by_name["z"]["prompt_count"], 0)
self.assertEqual(by_name["a"]["prompt_count"], 2)
self.assertEqual(
set(by_name["c"]), {"id", "name", "prompt_count", "image_count"}
)
self.assertEqual({row["image_count"] for row in page["tags"]}, {0})
def image(self, name, tags, digest="a", mtime_ns=1):
"""Tag the bytes ``digest`` of the file ``name``; return its locator."""
locator = os.path.normcase(os.path.join(os.path.realpath(self.temp.name), name))
self.read(
"replace_image_tags",
locator=locator,
sha256=hashlib.sha256(digest.encode()).hexdigest(),
size=1,
mtime_ns=mtime_ns,
expected_epoch=self.db.get_execution_epoch(),
expected_tag_revision=0,
tags=tags,
)
return locator
def test_tag_rows_report_image_count_by_locator(self):
self.save("Prompt only", tags=["prompt-only"])
self.image("one.png", ["twice", "shared"], digest="old")
self.image(
"one.png", ["twice"], digest="new", mtime_ns=2
) # two bytes, one file
self.image("two.png", ["shared"], digest="new") # equal digest, other locator
self.image("three.png", ["prior-only"], digest="first")
self.image("three.png", [], digest="second", mtime_ns=2)
rows = {row["name"]: row for row in self.all_tags()}
self.assertEqual(rows["twice"]["image_count"], 1)
row = rows["shared"]
self.assertEqual(row["image_count"], 2)
self.assertEqual(rows["prior-only"]["image_count"], 1)
self.assertEqual(rows["prompt-only"]["image_count"], 0)
self.assertEqual(
{name: r["prompt_count"] for name, r in rows.items()},
{"prompt-only": 1, "twice": 0, "shared": 0, "prior-only": 0},
)
self.assertEqual(
self.db.list_library(tags=["shared"])["total"], row["prompt_count"]
)
def test_prompt_count_equals_library_browse_total_for_every_tag(self):
self.save("Parent only", tags=["parent-only"])
self.db.save_prompt_pair("Pair cat", "Pair noise", tags=["pair-only"])
archived = self.db.save_prompt_pair(
"Archived cat", "Archived noise", tags=["archived-only"]
)
self.db.archive_prompt_pair(archived)
self.save("Legacy blur", tags=["negative-only"])
self.db.save_prompt_pair("Negative owner", "Legacy blur")
self.save("Dual dog", tags=["dual"])
self.db.save_prompt_pair("Dual dog", "Dual noise", tags=["dual"])
retained = self.save("Retained negative", tags=["retained"])
self.db.save_prompt_pair("Retained owner", "Retained negative")
with self.db.model.get_connection() as conn:
conn.execute(
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
(retained, "Kept"),
)
self.add_unused("unused")
rows = self.all_tags()
self.assertEqual(len(rows), 7)
for row in rows:
self.assertEqual(
row["prompt_count"],
self.db.list_library(tags=[row["name"]])["total"],
row["name"],
)
counts = {row["name"]: row["prompt_count"] for row in rows}
self.assertEqual(
counts,
{
"parent-only": 1,
"pair-only": 1,
"archived-only": 0,
"negative-only": 0,
"dual": 1,
"retained": 1,
"unused": 0,
},
)
def test_category_rows_match_library_options_and_browse_totals(self):
self.save("Zeta first", category="Zeta")
self.db.save_prompt_pair("Alpha owner", "Alpha noise", category="Alpha")
self.save("Padded", category=" Landscape ")
self.save("Plain", category="Landscape")
self.save("Blank", category=" ")
self.save("Null")
self.save("Zeta again", category="Zeta")
archived = self.db.save_prompt_pair(
"Archived owner", "Archived noise", category="Archived-only"
)
self.db.archive_prompt_pair(archived)
self.save("Legacy negative", category="Negative-only")
self.db.save_prompt_pair("Negative owner", "Legacy negative")
page = self.read("list_vocabulary_categories")
values = [row["value"] for row in page["categories"]]
self.assertEqual(values, ["Zeta", "Alpha", " Landscape ", "Landscape"])
self.assertEqual(set(values), set(self.db.get_library_options()["categories"]))
self.assertEqual(page["total"], 4)
self.assertFalse(page["has_more"])
for row in page["categories"]:
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"],
row["value"],
)
self.assertEqual(page["categories"][0]["prompt_count"], 2)
def test_category_pages_continue_by_first_use_then_value(self):
for index in range(5):
self.save(f"Owner {index}", category=f"C{4 - index}")
first = self.read("list_vocabulary_categories", limit=2)
second = self.read(
"list_vocabulary_categories", limit=2, after=first["next_key"]
)
third = self.read(
"list_vocabulary_categories", limit=2, after=second["next_key"]
)
values = [
row["value"]
for page in (first, second, third)
for row in page["categories"]
]
self.assertEqual(values, ["C4", "C3", "C2", "C1", "C0"])
self.assertEqual(
[page["has_more"] for page in (first, second, third)], [True, True, False]
)
for after in ([1], ["1", "C4"], {"id": 1}, [True, "C4"]):
with self.assertRaises(ValueError):
self.read("list_vocabulary_categories", after=after)
def test_category_paging_breaks_first_use_ties_by_value(self):
self.save("Shared owner", category="B")
self.db.save_prompt_pair(
"Shared owner", "Pair noise A", category="A", title="A pair"
)
self.db.save_prompt_pair(
"Shared owner", "Pair noise C", category="C", title="C pair"
)
pages, after = [], None
for _ in range(3):
page = self.read("list_vocabulary_categories", limit=1, after=after)
pages.append(page)
after = page["next_key"]
values = [row["value"] for page in pages for row in page["categories"]]
self.assertEqual(values, ["A", "B", "C"])
self.assertEqual([page["has_more"] for page in pages], [True, True, False])
for page in pages:
row = page["categories"][0]
self.assertEqual(
row["prompt_count"],
self.db.list_library(category=row["value"])["total"],
)
def test_pages_are_bounded_and_continue_without_gaps(self):
self.add_unused(*[f"tag-{index:03d}" for index in range(250)])
pages, after = [], None
for _ in range(3):
page = self.read("list_vocabulary_tags", limit=100, after=after)
pages.append(page)
after = page["next_key"]
self.assertEqual([len(page["tags"]) for page in pages], [100, 100, 50])
self.assertEqual([page["has_more"] for page in pages], [True, True, False])
self.assertEqual({page["total"] for page in pages}, {250})
names = [row["name"] for page in pages for row in page["tags"]]
self.assertEqual(names, [f"tag-{index:03d}" for index in range(250)])
for options in ({"limit": 0}, {"limit": 201}, {"after": "1"}, {"after": 0}):
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()
tags = self.read("list_vocabulary_tags")
categories = self.read("list_vocabulary_categories")
self.assertEqual(tags["library_epoch"], epoch)
self.assertEqual(categories["library_epoch"], epoch)
helper = getattr(self.db, "_vocabulary_tags", None)
self.assertTrue(callable(helper), "PromptDatabase._vocabulary_tags is missing")
with self.db.library_snapshot() as (conn, snapshot_epoch, _ceiling):
inner = helper(conn, snapshot_epoch, limit=10, after=None)
self.assertEqual(inner["library_epoch"], epoch)
self.assertTrue(conn.in_transaction)
_TAG_TABLES = ('"tags"', '"prompt_tags"', '"prompt_pair_tags"', '"sqlite_sequence"')
class TestVocabularyTagCommands(unittest.TestCase):
"""Guarded tag create, rename, merge and delete across both junctions."""
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.db = PromptDatabase(str(Path(self.temp.name) / "commands.db"))
self.epoch = self.db.get_execution_epoch()
def tearDown(self):
try:
self.db.close_all()
finally:
self.temp.cleanup()
def command(self, method_name, /, *args, **options):
method = getattr(self.db, method_name, None)
self.assertTrue(callable(method), f"PromptDatabase.{method_name} is missing")
return method(*args, **options)
def save(self, text, **metadata):
return self.db.save_prompt(
text, prompt_hash=generate_prompt_hash(text), **metadata
)
def tag_id(self, name):
with self.db.model.get_connection() as conn:
row = conn.execute("SELECT id FROM tags WHERE name=?", (name,)).fetchone()
return row[0] if row else None
def links(self, tag_name):
with self.db.model.get_connection() as conn:
components = conn.execute(
"SELECT pt.prompt_id FROM prompt_tags pt JOIN tags t "
"ON t.id=pt.tag_id WHERE t.name=? ORDER BY 1",
(tag_name,),
).fetchall()
pairs = conn.execute(
"SELECT ppt.pair_id FROM prompt_pair_tags ppt JOIN tags t "
"ON t.id=ppt.tag_id WHERE t.name=? ORDER BY 1",
(tag_name,),
).fetchall()
return [row[0] for row in components], [row[0] for row in pairs]
def junction_rows(self):
with self.db.model.get_connection() as conn:
return (
conn.execute("SELECT * FROM tags ORDER BY id").fetchall(),
conn.execute("SELECT * FROM prompt_tags ORDER BY 1,2").fetchall(),
conn.execute("SELECT * FROM prompt_pair_tags ORDER BY 1,2").fetchall(),
)
def other_rows(self):
"""Every row outside the tag registry and junctions, timestamps included."""
conn = self.db.model.get_connection()
return [
line
for line in conn.iterdump()
if line.startswith("INSERT INTO")
and not line.startswith(tuple(f"INSERT INTO {t}" for t in _TAG_TABLES))
]
def changes(self):
return self.db.model.get_connection().total_changes
def owners(self, tag="moody"):
"""Parent, active pair, archived pair and negative-only component."""
parent = self.save("Parent sky", category="Scenery", tags=[tag])
self.db.link_image_to_prompt(parent, str(Path(self.temp.name) / "a.png"))
active = self.db.save_prompt_pair("Active sky", "Active noise", tags=[tag])
archived = self.db.save_prompt_pair(
"Archived sky", "Archived noise", tags=[tag]
)
self.db.archive_prompt_pair(archived)
negative = self.save("Legacy blur", tags=[tag])
self.db.save_prompt_pair("Negative owner", "Legacy blur")
return parent, negative, active, archived
def test_create_inserts_unused_tag_and_reuses_exact_existing(self):
self.save("Owner", tags=["other"])
before = self.junction_rows()
changes = self.changes()
created = self.command(
"create_vocabulary_tag", " sunset ", expected_epoch=self.epoch
)
self.assertEqual(self.changes() - changes, 1)
self.assertEqual(created["library_epoch"], self.epoch)
self.assertTrue(created["created"])
self.assertEqual(
created["tag"],
{
"id": self.tag_id("sunset"),
"name": "sunset",
"prompt_count": 0,
"image_count": 0,
},
)
after = self.junction_rows()
self.assertEqual(after[1:], before[1:])
changes = self.changes()
again = self.command(
"create_vocabulary_tag", "sunset", expected_epoch=self.epoch
)
self.assertEqual(self.changes(), changes)
self.assertFalse(again["created"])
self.assertEqual(again["tag"], created["tag"])
other = self.command(
"create_vocabulary_tag", "Sunset", expected_epoch=self.epoch
)
self.assertTrue(other["created"])
self.assertNotEqual(other["tag"]["id"], created["tag"]["id"])
self.assertIn("sunset", self.db.get_all_tags())
def test_create_validates_names(self):
before = self.junction_rows()
for name in ("", " ", "x" * 51, "bad\x01", "a,b", 7, None, ["a"]):
with self.assertRaises(ValueError, msg=repr(name)):
self.command("create_vocabulary_tag", name, expected_epoch=self.epoch)
with self.assertRaises(ValueError):
self.command("create_vocabulary_tag", "ok", expected_epoch="nope")
self.assertEqual(self.junction_rows(), before)
def test_rename_to_new_name_keeps_id_and_all_links(self):
parent, negative, active, archived = self.owners()
tag_id = self.tag_id("moody")
others = self.other_rows()
result = self.command(
"rename_vocabulary_tag",
tag_id,
expected_epoch=self.epoch,
name="moody",
new_name=" dramatic ",
)
self.assertEqual(
result,
{
"library_epoch": self.epoch,
"merged": False,
"tag": {
"id": tag_id,
"name": "dramatic",
"prompt_count": 2,
"image_count": 0,
},
"components_changed": 2,
"pairs_changed": 2,
"images_changed": 0,
},
)
self.assertIsNone(self.tag_id("moody"))
self.assertEqual(
self.links("dramatic"), ([parent, negative], [active, archived])
)
self.assertEqual(self.db.get_prompt_pair(archived)["tags"], ["dramatic"])
self.assertEqual(self.other_rows(), others)
def test_rename_to_existing_merges_both_junctions_without_duplicates(self):
both = self.save("Both", tags=["old", "new"])
only_old = self.save("Old only", tags=["old"])
pair_both = self.db.save_prompt_pair("Pair both", "Noise", tags=["old", "new"])
pair_old = self.db.save_prompt_pair("Pair old", "Noise", tags=["old"])
self.db.archive_prompt_pair(pair_old)
source, target = self.tag_id("old"), self.tag_id("new")
others = self.other_rows()
result = self.command(
"rename_vocabulary_tag",
source,
expected_epoch=self.epoch,
name="old",
new_name="new",
)
self.assertTrue(result["merged"])
self.assertEqual(result["tag"]["id"], target)
self.assertEqual(result["tag"]["name"], "new")
self.assertEqual(result["tag"]["prompt_count"], 3)
self.assertEqual(
(result["components_changed"], result["pairs_changed"]), (2, 2)
)
self.assertIsNone(self.tag_id("old"))
self.assertEqual(self.links("new"), ([both, only_old], [pair_both, pair_old]))
self.assertEqual(self.db.get_prompt_by_id(both)["tags"], ["new"])
self.assertEqual(self.db.get_prompt_pair(pair_both)["tags"], ["new"])
self.assertEqual(self.other_rows(), others)
def test_case_only_rename_is_a_rename(self):
prompt = self.save("Kitten", tags=["Cat"])
tag_id = self.tag_id("Cat")
result = self.command(
"rename_vocabulary_tag",
tag_id,
expected_epoch=self.epoch,
name="Cat",
new_name="cat",
)
self.assertFalse(result["merged"])
self.assertEqual(
result["tag"],
{"id": tag_id, "name": "cat", "prompt_count": 1, "image_count": 0},
)
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["cat"])
def test_stale_identity_and_epoch_write_nothing(self):
self.owners()
tag_id = self.tag_id("moody")
before = (self.junction_rows(), self.other_rows())
changes = self.changes()
for tag, name in ((tag_id, "Moody"), (tag_id + 99, "moody")):
for method, extra in (
("rename_vocabulary_tag", {"new_name": "fresh"}),
("delete_vocabulary_tag", {}),
):
with self.assertRaises(LibraryCommandError) as caught:
self.command(
method, tag, expected_epoch=self.epoch, name=name, **extra
)
self.assertEqual(caught.exception.code, "vocabulary-conflict")
self.assertEqual(caught.exception.epoch, self.epoch)
stale = "00000000-0000-4000-8000-000000000000"
with patch.object(type(self.db), "_vocabulary_tag", create=True) as lookup:
for method, args, extra in (
("create_vocabulary_tag", ("fresh",), {}),
(
"rename_vocabulary_tag",
(tag_id,),
{"name": "moody", "new_name": "x"},
),
("delete_vocabulary_tag", (tag_id,), {"name": "moody"}),
):
with self.assertRaises(LibraryCommandError) as caught:
self.command(method, *args, expected_epoch=stale, **extra)
self.assertEqual(caught.exception.code, "storage-changed")
self.assertEqual(caught.exception.epoch, self.epoch)
lookup.assert_not_called()
for new_name in ("moody", " moody "):
with self.assertRaises(ValueError):
self.command(
"rename_vocabulary_tag",
tag_id,
expected_epoch=self.epoch,
name="moody",
new_name=new_name,
)
for bad_id in (0, -1, "1", True, 2**63):
with self.assertRaises(ValueError):
self.command(
"delete_vocabulary_tag",
bad_id,
expected_epoch=self.epoch,
name="moody",
)
self.assertEqual(self.changes(), changes)
self.assertEqual((self.junction_rows(), self.other_rows()), before)
def test_delete_removes_name_and_links_only(self):
self.owners()
self.save("Bystander", tags=["keep"])
tag_id = self.tag_id("moody")
impact = self.command("vocabulary_tag_impact", tag_id)
self.assertEqual(
impact,
{
"library_epoch": self.epoch,
"impact": {
"positive_prompts": 1,
"other_prompts": 1,
"negative_variants": 1,
"archived_negative_variants": 1,
"images": 0,
},
},
)
self.assertIsNone(self.command("vocabulary_tag_impact", tag_id + 99))
others = self.other_rows()
result = self.command(
"delete_vocabulary_tag", tag_id, expected_epoch=self.epoch, name="moody"
)
self.assertEqual(
result,
{
"library_epoch": self.epoch,
"deleted": True,
"components_changed": 2,
"pairs_changed": 2,
"images_changed": 0,
},
)
self.assertIsNone(self.tag_id("moody"))
self.assertEqual(self.links("moody"), ([], []))
self.assertEqual(self.other_rows(), others)
self.assertEqual(len(self.links("keep")[0]), 1)
with self.assertRaises(LibraryCommandError) as caught:
self.command(
"delete_vocabulary_tag", tag_id, expected_epoch=self.epoch, name="moody"
)
self.assertEqual(caught.exception.code, "vocabulary-conflict")
def image(self, name, tags, digest="a", mtime_ns=1):
locator = os.path.normcase(os.path.join(os.path.realpath(self.temp.name), name))
self.command(
"replace_image_tags",
locator=locator,
sha256=hashlib.sha256(digest.encode()).hexdigest(),
size=1,
mtime_ns=mtime_ns,
expected_epoch=self.epoch,
expected_tag_revision=0,
tags=tags,
)
return locator
def test_impact_and_delete_counts_include_images(self):
self.owners()
self.image("a.png", ["moody", "keep"])
self.image("a.png", ["moody"], digest="prior", mtime_ns=2)
self.image("b.png", ["moody"])
tag_id = self.tag_id("moody")
impact = self.command("vocabulary_tag_impact", tag_id)["impact"]
self.assertEqual(impact["images"], 2)
self.assertEqual(
self.command("vocabulary_tag_impact", self.tag_id("keep"))["impact"],
{
"positive_prompts": 0,
"other_prompts": 0,
"negative_variants": 0,
"archived_negative_variants": 0,
"images": 1,
},
)
result = self.command(
"delete_vocabulary_tag", tag_id, expected_epoch=self.epoch, name="moody"
)
self.assertEqual(
(
result["components_changed"],
result["pairs_changed"],
result["images_changed"],
),
(2, 2, 2),
)
with self.db.model.get_connection() as conn:
self.assertEqual(
conn.execute("SELECT COUNT(*) FROM media_image_tags").fetchone()[0], 1
)
self.assertEqual(
conn.execute("SELECT COUNT(*) FROM media_annotation_bytes").fetchone()[
0
],
3,
)
renamed = self.command(
"rename_vocabulary_tag",
self.tag_id("keep"),
expected_epoch=self.epoch,
name="keep",
new_name="kept",
)
self.assertEqual(renamed["images_changed"], 1)
self.assertEqual(renamed["tag"]["image_count"], 1)
_OLD_STAMP = "2000-01-01T00:00:00+00:00"
_DML = re.compile(r"\s*(INSERT|UPDATE|DELETE|REPLACE)\b", re.IGNORECASE)
class TestVocabularyCategoryCommands(unittest.TestCase):
"""Category preview and token-guarded rename/consolidation."""
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.db = PromptDatabase(str(Path(self.temp.name) / "categories.db"))
self.epoch = self.db.get_execution_epoch()
def tearDown(self):
try:
self.db.close_all()
finally:
self.temp.cleanup()
def command(self, method_name, /, **options):
method = getattr(self.db, method_name, None)
self.assertTrue(callable(method), f"PromptDatabase.{method_name} is missing")
return method(**options)
def preview(self, source="Scenery", target="Landscape"):
return self.command(
"preview_vocabulary_category",
expected_epoch=self.epoch,
source=source,
target=target,
)
def apply(self, token, source="Scenery", target="Landscape"):
return self.command(
"apply_vocabulary_category",
expected_epoch=self.epoch,
source=source,
target=target,
impact_token=token,
)
def save(self, text, **metadata):
return self.db.save_prompt(
text, prompt_hash=generate_prompt_hash(text), **metadata
)
def owners(self, source="Scenery"):
"""2 eligible, 1 negative-only component; 2 active, 1 archived pair."""
first = self.save("Parent one", category=source)
second = self.save("Parent two", category=source)
negative = self.save("Legacy blur", category=source)
self.db.save_prompt_pair("Negative owner", "Legacy blur")
active = [
self.db.save_prompt_pair(
f"Active {n}",
f"Noise {n}",
category=source,
title=f"Title {n}",
tags=[f"tag-{n}"],
rating=n,
notes=f"Notes {n}",
)
for n in (1, 2)
]
archived = self.db.save_prompt_pair(
"Archived sky", "Archived noise", category=source, title="Kept"
)
self.db.archive_prompt_pair(archived)
return [first, second], negative, active, archived
def age_rows(self):
with self.db.model.get_connection() as conn:
conn.execute("UPDATE prompts SET updated_at=?", (_OLD_STAMP,))
conn.execute("UPDATE prompt_pairs SET updated_at=?", (_OLD_STAMP,))
def rows(self, table):
with self.db.model.get_connection() as conn:
return {
row["id"]: dict(row)
for row in conn.execute(f"SELECT * FROM {table} ORDER BY id")
}
def snapshot(self):
return self.rows("prompts"), self.rows("prompt_pairs")
def traced(self, call):
statements = []
original = self.db.model.get_connection
def trace_connection():
connection = original()
connection.set_trace_callback(statements.append)
return connection
with patch.object(
self.db.model, "get_connection", side_effect=trace_connection
):
result = call()
conn = original()
conn.set_trace_callback(None)
return result, statements
def test_preview_counts_every_owner_type(self):
self.owners()
self.save("Elsewhere", category="Other")
before = self.snapshot()
result, statements = self.traced(self.preview)
self.assertEqual(
set(result),
{
"library_epoch",
"source",
"target",
"target_exists",
"impact",
"impact_token",
},
)
self.assertEqual(result["library_epoch"], self.epoch)
self.assertEqual((result["source"], result["target"]), ("Scenery", "Landscape"))
self.assertEqual(
result["impact"],
{
"positive_prompts": 2,
"other_prompts": 1,
"negative_variants": 2,
"archived_negative_variants": 1,
"images": 0,
},
)
self.assertFalse(result["target_exists"])
self.assertRegex(result["impact_token"], r"\A[0-9a-f]{64}\Z")
self.assertEqual([q for q in statements if _DML.match(q)], [])
self.assertEqual(self.snapshot(), before)
trimmed = self.preview(target=" Other ")
self.assertEqual(trimmed["target"], "Other")
self.assertTrue(trimmed["target_exists"])
self.assertNotEqual(trimmed["impact_token"], result["impact_token"])
self.assertEqual(self.preview(), result)
def test_apply_renames_exact_value_on_components_and_all_pairs(self):
components, negative, active, archived = self.owners()
variant = self.save("Padded", category=" Scenery ")
bystander = self.save("Bystander", category="Other")
self.age_rows()
prompts_before, pairs_before = self.snapshot()
preview = self.preview()
receipt = self.apply(preview["impact_token"])
self.assertEqual(
receipt,
{
"library_epoch": self.epoch,
"source": "Scenery",
"target": "Landscape",
"consolidated": False,
"impact": preview["impact"],
},
)
prompts_after, pairs_after = self.snapshot()
changed_prompts = {*components, negative}
changed_pairs = {*active, archived}
for table_before, table_after, changed in (
(prompts_before, prompts_after, changed_prompts),
(pairs_before, pairs_after, changed_pairs),
):
self.assertNotIn("Scenery", [r["category"] for r in table_after.values()])
self.assertEqual(set(table_after), set(table_before))
for row_id, row in table_after.items():
old = table_before[row_id]
if row_id in changed:
self.assertEqual(row["category"], "Landscape")
self.assertNotEqual(row["updated_at"], _OLD_STAMP)
expected = {**old, "category": "Landscape"}
expected["updated_at"] = row["updated_at"]
self.assertEqual(row, expected)
else:
self.assertEqual(row, old)
self.assertEqual(prompts_after[variant]["category"], " Scenery ")
self.assertEqual(prompts_after[bystander]["category"], "Other")
for pair_id in active:
pair = self.db.get_prompt_pair(pair_id)
self.assertEqual(pair["category"], "Landscape")
self.assertEqual(len(pair["tags"]), 1)
self.assertEqual(pairs_after[archived]["archived"], 1)
self.assertEqual(pairs_after[archived]["title"], "Kept")
# Consolidating into an existing value reports it.
merge = self.preview(source=" Scenery ", target="Other")
self.assertTrue(merge["target_exists"])
receipt = self.apply(merge["impact_token"], source=" Scenery ", target="Other")
self.assertTrue(receipt["consolidated"])
self.assertEqual(
receipt["impact"],
{
"positive_prompts": 1,
"other_prompts": 0,
"negative_variants": 0,
"archived_negative_variants": 0,
"images": 0,
},
)
self.assertEqual(self.rows("prompts")[variant]["category"], "Other")
def test_token_changes_with_membership_eligibility_archive_and_target_presence(
self,
):
components, negative, active, _archived = self.owners()
def add_owner():
self.save("Late owner", category="Scenery")
def archive_pair():
self.db.archive_prompt_pair(active[0])
def promote_negative():
with self.db.model.get_connection() as conn:
conn.execute(
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
(negative, "Promoted"),
)
def create_target():
self.save("Target owner", category="Landscape")
for mutate in (add_owner, archive_pair, promote_negative, create_target):
with self.subTest(mutate.__name__):
token = self.preview()["impact_token"]
mutate()
before = self.snapshot()
changes = self.db.model.get_connection().total_changes
with self.assertRaises(LibraryCommandError) as caught:
self.apply(token)
self.assertEqual(caught.exception.code, "vocabulary-conflict")
self.assertEqual(caught.exception.epoch, self.epoch)
self.assertEqual(self.db.model.get_connection().total_changes, changes)
self.assertEqual(self.snapshot(), before)
token = self.preview()["impact_token"]
with self.db.model.get_connection() as conn:
conn.execute(
"UPDATE prompts SET notes='edited' WHERE id=?", (components[0],)
)
self.assertTrue(self.apply(token)["consolidated"])
self.assertEqual(self.rows("prompts")[components[0]]["category"], "Landscape")
def test_category_command_guards(self):
self.owners()
before = self.snapshot()
stale = "00000000-0000-4000-8000-000000000000"
token = self.preview()["impact_token"]
with self.assertRaises(LibraryCommandError) as caught:
self.command(
"preview_vocabulary_category",
expected_epoch=stale,
source="Scenery",
target="Landscape",
)
self.assertEqual(caught.exception.code, "storage-changed")
self.assertEqual(caught.exception.epoch, self.epoch)
with self.assertRaises(LibraryCommandError) as caught:
self.command(
"apply_vocabulary_category",
expected_epoch=stale,
source="Scenery",
target="Landscape",
impact_token=token,
)
self.assertEqual(caught.exception.code, "storage-changed")
for source in ("Missing", "scenery", " Scenery"):
with self.assertRaises(LibraryCommandError) as caught:
self.preview(source=source)
self.assertEqual(caught.exception.code, "vocabulary-conflict")
for source, target in (
("Scenery", ""),
("Scenery", " "),
("Scenery", "x" * 101),
("Scenery", "bad\x01"),
("Scenery", "Scenery"),
("Scenery", " Scenery "),
("Scenery", None),
("Scenery", 7),
(None, "Landscape"),
("", "Landscape"),
(" ", "Landscape"),
):
with self.assertRaises(ValueError, msg=(source, target)):
self.preview(source=source, target=target)
for bad in ("", "A" * 64, token[:-1], None, 7):
with self.assertRaises(ValueError):
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"]
with self.db.model.get_connection() as conn:
conn.execute(
"CREATE TRIGGER abort_pairs BEFORE UPDATE ON prompt_pairs "
"BEGIN SELECT RAISE(ABORT, 'boom'); END"
)
before = self.snapshot()
with self.assertRaises(sqlite3.Error):
self.apply(token)
self.assertEqual(self.snapshot(), before)
self.assertIn("Scenery", [r["category"] for r in before[0].values()])
self.assertFalse(self.db.model.get_connection().in_transaction)
class TestSharedTokenHashing(unittest.TestCase):
def test_both_tokens_keep_their_previous_values(self):
from database.library_tokens import library_edit_token
from database.vocabulary import category_impact_token
record = {
"id": 3,
"text": "t",
"title": None,
"category": "c",
"notes": "",
"rating": 0,
"tags": ["b", "a", "a"],
}
self.assertEqual(
library_edit_token("e", record),
"32534d4c92e4e9a7d8e31c30ecdbb07563703b5bd7ed216a8d2066c0dc44332b",
)
self.assertEqual(
category_impact_token(
"e", "S", "T", True, [(1, True), (2, False)], [(5, False)]
),
"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__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink