nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge branch 'feat/331-collections-pr1' into feat/331-collections-pr2

Browse Source
This commit is contained in:
Vito Sansevero
2026-10-27 18:12:22 -07:00
parent 46c5e387c8 1db6088b2f
commit 5b2be11ba7
31 changed files with 2951 additions and 409 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/browser/server.py
+2
View File
@@ -262,6 +262,8 @@ def create_app(root: Path) -> web.Application:
with target.db.model.get_connection() as conn:
for table in (
"saved_views",
"image_collection_items",
"image_collections",
"image_favorites",
"media_review_decisions",
"media_review_items",
tests/js/workspace-system-backups.test.js
+2 -2
View File
@@ -63,9 +63,9 @@ test('status line shows only server-reported backups (D103)', async () => {
const PREVIEW = Object.freeze({
success: true, token: 'r'.repeat(43), expires_at: '2026-10-25T12:15:00+00:00',
summary: { prompts: 1234, positives: 900, pairs: 3, archived_pairs: 0, image_links: 56, runs: 1, saved_views: 2 },
summary: { prompts: 1234, positives: 900, pairs: 3, archived_pairs: 0, image_links: 56, runs: 1, saved_views: 2, collections: 1 },
schema: { migrated: true, library_state: true }, integrity: 'ok', documents: { repaired: 0, cleared: 0, images: 0 },
live: { prompts: 1, positives: 1, pairs: 0, archived_pairs: 0, image_links: 0, runs: 0, saved_views: 0 }, size: 4096,
live: { prompts: 1, positives: 1, pairs: 0, archived_pairs: 0, image_links: 0, runs: 0, saved_views: 0, collections: 0 }, size: 4096,
safety_backup: { destination_label: 'prompts.db.backup_<date and time>' },
});
tests/test_api_admin.py
+13 -21
View File
@@ -846,6 +846,14 @@ class TestSettingsReadMonitorThroughModule(AdminAPITestCase):
self.assertEqual(data["settings"]["gallery_root_paths"], ["output"])
def _wired_prompt(text):
"""A sampler-less graph whose one encoder feeds a guider (#252)."""
return {
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": text}},
"2": {"class_type": "BasicGuider", "inputs": {"conditioning": ["1", 0]}},
}
def _raise(*args, **kwargs):
raise RuntimeError("boom")
@@ -2186,11 +2194,7 @@ class TestOutputScanInternals(AdminAPITestCase):
def test_ingest_outcomes(self):
f = Path("gen.png")
meta = {
"prompt": json.dumps(
{"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
)
}
meta = {"prompt": json.dumps(_wired_prompt("a cat"))}
self.assertIsNone(self.api._ingest_scanned_file(f, {}))
self.assertIsNone(self.api._ingest_scanned_file(f, {"prompt": "{}"}))
self.assertEqual(
@@ -2202,31 +2206,19 @@ class TestOutputScanInternals(AdminAPITestCase):
def test_ingest_link_failures_are_tolerated(self):
f = Path("gen.png")
meta = {
"prompt": json.dumps(
{"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a dog"}}}
)
}
meta = {"prompt": json.dumps(_wired_prompt("a dog"))}
self.api.db.link_image_to_prompt = _raise
self.assertEqual(self.api._ingest_scanned_file(f, meta), "added")
self.assertEqual(self.api._ingest_scanned_file(f, meta), "found")
def test_ingest_save_failure_counts_as_found(self):
f = Path("gen.png")
meta = {
"prompt": json.dumps(
{"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a bird"}}}
)
}
meta = {"prompt": json.dumps(_wired_prompt("a bird"))}
self.api.db.save_prompt = lambda *a, **k: None
self.assertEqual(self.api._ingest_scanned_file(f, meta), "found")
def test_batch_counts_and_per_file_errors(self):
meta = {
"prompt": json.dumps(
{"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a fish"}}}
)
}
meta = {"prompt": json.dumps(_wired_prompt("a fish"))}
self.api.db.get_prompt_by_hash = _raise
counts = self.api._ingest_scan_batch_sync(
[(Path("a.png"), meta), (Path("b.png"), {})]
@@ -2915,7 +2907,7 @@ class TestScanReadsVideos(AdminAPITestCase):
shutil.which("ffmpeg") and shutil.which("ffprobe"), "ffmpeg not installed"
)
async def test_scan_adds_prompts_found_in_videos(self):
graph = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
graph = _wired_prompt("a cat")
_write_mp4(
self.output_dir / "clip.mp4", json.dumps({"prompt": json.dumps(graph)})
)
tests/test_api_collections.py
+494
View File
@@ -0,0 +1,494 @@
"""Collections (#331): the routes under /prompt_manager/collections."""
import json
import ntpath
import sqlite3
from unittest.mock import patch
try:
from test_api_images import ImageAPITestCase, make_png
from open_handles import assert_closed
except ImportError:
from tests.test_api_images import ImageAPITestCase, make_png
from tests.open_handles import assert_closed
BASE = "/prompt_manager/collections"
ITEMS = BASE + "/items"
MEMBERSHIP = BASE + "/membership"
SOURCES_CHANGED = "The sources changed. Refresh the images and try again."
class CollectionRoutesTestCase(ImageAPITestCase):
async def tearDownAsync(self):
self.api.db.close_all()
assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def make(self, *names, root=None):
for name in names:
make_png((root or self.output_dir) / name)
def scope(self, index=0):
roots = list(self.api._get_all_output_dirs())
return self.api._gallery_root_scope(roots[index], index)
def item(self, path, index=0, scope=None):
return {
"root_index": index,
"root_scope": scope or self.scope(index),
"path": path,
}
def two_sources(self):
second = self.comfy_dir / "second"
second.mkdir()
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
return second
async def call(self, method, url, body=None, status=200, **params):
response = await self.client.request(method, url, json=body, params=params)
text = await response.text()
self.assertEqual(response.status, status, text)
data = json.loads(text)
self.assertNotIn(self.comfy_dir.as_posix(), text)
self.assertNotIn(str(self.comfy_dir).replace("\\", "\\\\"), text)
return data
async def add(self, status=200, **body):
return await self.call("POST", ITEMS, body, status)
async def listed(self):
return (await self.call("GET", BASE))["collections"]
class TestCollectionRoutes(CollectionRoutesTestCase):
async def test_each_route_succeeds(self):
self.make("a.png", "b.png", "c.png")
self.assertEqual(
await self.call("GET", BASE), {"success": True, "collections": []}
)
made = await self.add(
name=" LoRa testing ", items=[self.item("a.png"), self.item("b.png")]
)
grids = made["collection"]
self.assertEqual(grids["name"], "LoRa testing")
self.assertEqual(
made,
{
"success": True,
"collection": grids,
"created": True,
"added": 2,
"already": 0,
"not_found": 0,
},
)
again = await self.add(
collection_id=grids["id"], items=[self.item("b.png"), self.item("c.png")]
)
self.assertEqual(
(again["created"], again["added"], again["already"]), (False, 1, 1)
)
self.assertEqual(
await self.listed(),
[{"id": grids["id"], "name": "LoRa testing", "count": 3, "missing": 0}],
)
membership = await self.call(
"GET", MEMBERSHIP, root_index="0", root_scope=self.scope(), path="a.png"
)
self.assertEqual(membership, {"success": True, "collection_ids": [grids["id"]]})
removed = await self.call(
"DELETE",
ITEMS,
{"collection_id": grids["id"], "items": [self.item("a.png")]},
)
self.assertEqual(removed, {"success": True, "collection": grids, "removed": 1})
renamed = await self.call(
"PATCH", f"{BASE}/{grids['id']}", {"name": "LoRA grids"}
)
self.assertEqual(
renamed,
{"success": True, "collection": {"id": grids["id"], "name": "LoRA grids"}},
)
deleted = await self.call("DELETE", f"{BASE}/{grids['id']}")
self.assertEqual(deleted, {"success": True, "deleted": True})
self.assertEqual(await self.listed(), [])
# Files are never touched.
self.assertTrue((self.output_dir / "a.png").is_file())
async def test_names_are_unique_ignoring_case(self):
self.make("a.png", "b.png")
first = (await self.add(name="LoRa testing", items=[self.item("a.png")]))[
"collection"
]
second = await self.add(name="lora TESTING", items=[self.item("b.png")])
self.assertEqual(second["collection"], first)
self.assertFalse(second["created"])
self.assertEqual(len(await self.listed()), 1)
async def test_bad_requests_are_400(self):
self.make("a.png")
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
good = [self.item("a.png")]
for body in (
{"items": good},
{"name": "x", "collection_id": cid, "items": good},
{"name": "x"},
{"name": "x", "items": good, "extra": 1},
{"collection_id": "1", "items": good},
{"collection_id": True, "items": good},
{"collection_id": 0, "items": good},
{"collection_id": 10**10, "items": good},
{"name": "", "items": good},
{"name": "x" * 101, "items": good},
{"name": "bell\x07", "items": good},
{"name": 5, "items": good},
{"name": "x", "items": []},
{"name": "x", "items": "a.png"},
{"name": "x", "items": [self.item("a.png")] * 501},
{"name": "x", "items": [{"root_index": 0, "path": "a.png"}]},
{"name": "x", "items": [{**self.item("a.png"), "root_index": "0"}]},
{"name": "x", "items": [{**self.item("a.png"), "root_scope": "xyz"}]},
):
with self.subTest(body=body):
data = await self.add(400, **body)
self.assertFalse(data["success"])
for body in ({"items": good}, {"collection_id": cid}):
with self.subTest(remove=body):
await self.call("DELETE", ITEMS, body, 400)
for raw in ("abc", "0", "1234567890", "-1", "1.0"):
with self.subTest(path_id=raw):
await self.call("PATCH", f"{BASE}/{raw}", {"name": "x"}, 400)
await self.call("DELETE", f"{BASE}/{raw}", None, 400)
await self.call("PATCH", f"{BASE}/{cid}", {"name": ""}, 400)
await self.call("PATCH", f"{BASE}/{cid}", {"label": "x"}, 400)
for params in (
{"root_index": "x", "root_scope": self.scope(), "path": "a.png"},
{"root_index": "0", "path": "a.png"},
{"root_index": "0", "root_scope": "nope", "path": "a.png"},
):
with self.subTest(membership=params):
await self.call("GET", MEMBERSHIP, status=400, **params)
self.assertEqual((await self.listed())[0]["count"], 1)
async def test_unsafe_paths_are_403(self):
self.make("a.png")
for path in (
"../a.png",
"sub/../../a.png",
"C:/a.png",
"a\\b.png",
"/a.png",
"notes.txt",
):
with self.subTest(path=path):
await self.add(403, name="Set", items=[self.item(path)])
await self.call(
"GET",
MEMBERSHIP,
status=403,
root_index="0",
root_scope=self.scope(),
path=path,
)
self.assertEqual(await self.listed(), [])
async def test_a_body_over_the_cap_is_413(self):
response = await self.client.post(
ITEMS, data=b'{"name":"x","items":[' + b" " * (1024 * 1024 + 1) + b"]}"
)
self.assertEqual(response.status, 413)
async def test_missing_collection_and_files_are_404(self):
self.make("a.png")
data = await self.add(404, collection_id=999, items=[self.item("a.png")])
self.assertEqual(data["code"], "collection-not-found")
data = await self.call(
"DELETE", ITEMS, {"collection_id": 999, "items": [self.item("a.png")]}, 404
)
self.assertEqual(data["code"], "collection-not-found")
data = await self.call("PATCH", f"{BASE}/999", {"name": "x"}, 404)
self.assertEqual(data["code"], "collection-not-found")
data = await self.call("DELETE", f"{BASE}/999", None, 404)
self.assertEqual(data["code"], "collection-not-found")
data = await self.add(404, name="New", items=[self.item("never.png")])
self.assertEqual(data["code"], "files-not-found")
self.assertEqual(await self.listed(), [])
async def test_unlisted_files_count_as_not_found(self):
self.make("a.png")
made = await self.add(
name="Set", items=[self.item("a.png"), self.item("later.png")]
)
self.assertEqual((made["added"], made["not_found"]), (1, 1))
again = await self.add(
collection_id=made["collection"]["id"],
items=[self.item("a.png"), self.item("later.png")],
)
self.assertEqual(
(again["added"], again["already"], again["not_found"]), (0, 1, 1)
)
async def test_name_taken_is_409(self):
self.make("a.png")
first = (await self.add(name="First", items=[self.item("a.png")]))["collection"]
second = (await self.add(name="Second", items=[self.item("a.png")]))[
"collection"
]
data = await self.call(
"PATCH", f"{BASE}/{second['id']}", {"name": "FIRST"}, 409
)
self.assertEqual(
(data["code"], data["collection_id"]), ("name-taken", first["id"])
)
renamed = await self.call("PATCH", f"{BASE}/{second['id']}", {"name": "SECOND"})
self.assertEqual(renamed["collection"]["name"], "SECOND")
async def test_stale_scope_is_409_after_reordering_sources(self):
second = self.two_sources()
self.make("a.png")
self.make("b.png", root=second)
stale = self.item("b.png", 1)
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
self.api._get_all_output_dirs = lambda: [second, self.output_dir]
for method, body in (
("POST", {"collection_id": cid, "items": [self.item("b.png", 0), stale]}),
("DELETE", {"collection_id": cid, "items": [stale]}),
):
with self.subTest(method=method):
data = await self.call(method, ITEMS, body, 409)
self.assertEqual(data["code"], "source-changed")
self.assertEqual(data["error"], SOURCES_CHANGED)
data = await self.call(
"GET",
MEMBERSHIP,
status=409,
root_index="1",
root_scope=stale["root_scope"],
path="b.png",
)
self.assertEqual(data["code"], "source-changed")
# Nothing was written by the refused add; the member follows its folder.
listed = await self.listed()
self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0))
membership = await self.call(
"GET", MEMBERSHIP, root_index="1", root_scope=self.scope(1), path="a.png"
)
self.assertEqual(membership["collection_ids"], [cid])
async def test_counts_exclude_a_missing_member(self):
second = self.two_sources()
self.make("a.png", "b.png")
self.make("c.png", root=second)
made = await self.add(
name="Grids",
items=[self.item("a.png"), self.item("b.png"), self.item("c.png", 1)],
)
other = await self.add(name="Other", items=[self.item("c.png", 1)])
(self.output_dir / "b.png").unlink()
self.api.invalidate_gallery_cache()
listed = await self.listed()
self.assertEqual(
listed,
[
{
"id": made["collection"]["id"],
"name": "Grids",
"count": 2,
"missing": 1,
},
{
"id": other["collection"]["id"],
"name": "Other",
"count": 1,
"missing": 0,
},
],
)
# The member is removable although it is no longer listed.
removed = await self.call(
"DELETE",
ITEMS,
{"collection_id": made["collection"]["id"], "items": [self.item("b.png")]},
)
self.assertEqual(removed["removed"], 1)
async def test_counts_are_null_when_a_source_scan_fails(self):
self.make("a.png")
await self.add(name="Set", items=[self.item("a.png")])
async def broken(_root):
raise OSError("scan failed at /secret/place")
with patch.object(self.api, "_get_gallery_files", side_effect=broken):
data = await self.call("GET", BASE)
self.assertEqual(
data["collections"][0] | {"id": 0},
{"id": 0, "name": "Set", "count": None, "missing": None},
)
self.assertNotIn("secret", json.dumps(data))
async def test_windows_case_spellings_are_one_member(self):
self.make("Sub/A.png")
with patch("py.api.images._path_flavor", ntpath):
made = await self.add(
name="Set", items=[self.item("Sub/A.png"), self.item("SUB/a.PNG")]
)
self.assertEqual((made["added"], made["already"]), (1, 0))
(self.output_dir / "Sub" / "A.png").rename(
self.output_dir / "Sub" / "a.png"
)
(self.output_dir / "Sub").rename(self.output_dir / "sub")
self.api.invalidate_gallery_cache()
listed = await self.listed()
self.assertEqual((listed[0]["count"], listed[0]["missing"]), (1, 0))
membership = await self.call(
"GET",
MEMBERSHIP,
root_index="0",
root_scope=self.scope(),
path="sub/a.png",
)
self.assertEqual(membership["collection_ids"], [made["collection"]["id"]])
response = await self.client.get(
"/prompt_manager/images/output",
params={
"mode": "workspace",
"collection": str(made["collection"]["id"]),
},
)
page = await response.json()
self.assertEqual([i["path"] for i in page["images"]], ["sub/a.png"])
async def test_posix_case_distinct_files_stay_two_members(self):
self.make("A.png", "a.png")
if len(list(self.output_dir.iterdir())) != 2:
self.skipTest("case-insensitive filesystem")
made = await self.add(
name="Set", items=[self.item("A.png"), self.item("a.png")]
)
self.assertEqual(made["added"], 2)
async def test_a_500_never_carries_the_exception_text(self):
self.make("a.png")
made = await self.add(name="Set", items=[self.item("a.png")])
cid = made["collection"]["id"]
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("disk I/O error at /secret/prompts.db")
cases = (
(
"list_collections",
"GET",
BASE,
None,
{},
"Collections could not be read.",
),
(
"file_collections",
"GET",
MEMBERSHIP,
None,
{"root_index": "0", "root_scope": self.scope(), "path": "a.png"},
"Collections could not be read.",
),
(
"add_to_collection",
"POST",
ITEMS,
{"collection_id": cid, "items": [self.item("a.png")]},
{},
"The collection could not be changed.",
),
(
"remove_from_collection",
"DELETE",
ITEMS,
{"collection_id": cid, "items": [self.item("a.png")]},
{},
"The collection could not be changed.",
),
(
"rename_collection",
"PATCH",
f"{BASE}/{cid}",
{"name": "x"},
{},
"The collection could not be changed.",
),
(
"delete_collection",
"DELETE",
f"{BASE}/{cid}",
None,
{},
"The collection could not be changed.",
),
)
for method_name, method, url, body, params, message in cases:
with self.subTest(method=method_name):
with patch.object(self.api.db, method_name, side_effect=broken):
data = await self.call(method, url, body, 500, **params)
self.assertEqual(data, {"success": False, "error": message})
class TestCollectionListingFilter(CollectionRoutesTestCase):
"""#392's "all N matching" filter carries ``collection`` like the listing."""
async def test_a_filter_with_a_collection_resolves_only_members(self):
second = self.two_sources()
self.make("a.png", "b.png")
self.make("c.png", root=second)
made = await self.add(
name="Set", items=[self.item("a.png"), self.item("c.png", 1)]
)
cid = str(made["collection"]["id"])
selection = await self.api.resolve_listing_filter({"collection": cid})
self.assertEqual(
sorted((index, path) for index, _scope, path in selection.targets),
[(0, "a.png"), (1, "c.png")],
)
selection = await self.api.resolve_listing_filter(
{"collection": cid, "root": "1", "root_scope": self.scope(1)}
)
self.assertEqual([t[2] for t in selection.targets], ["c.png"])
for bad in ("0", "x", "1234567890"):
with self.subTest(bad=bad):
with self.assertRaises(ValueError):
from py.api.images import parse_listing_filter
parse_listing_filter({"collection": bad})
async def test_bulk_category_acts_on_collection_members_only(self):
self.make("a.png", "b.png")
inside = self._save_prompt("Inside", category="Old")
outside = self._save_prompt("Outside", category="Old")
self.api.db.link_image_to_prompt(inside, str(self.output_dir / "a.png"))
self.api.db.link_image_to_prompt(outside, str(self.output_dir / "b.png"))
made = await self.add(name="Set", items=[self.item("a.png")])
cid = str(made["collection"]["id"])
data = await self.call(
"POST",
"/prompt_manager/bulk/category",
{"images_filter": {"collection": cid}, "category": "New"},
)
self.assertEqual((data["updated_count"], data["prompt_count"]), (1, 1))
self.assertEqual(self.api.db.get_prompt_by_id(inside)["category"], "New")
self.assertEqual(self.api.db.get_prompt_by_id(outside)["category"], "Old")
async def test_a_deleted_collection_in_a_bulk_filter_is_a_404(self):
self.make("a.png")
data = await self.call(
"POST",
"/prompt_manager/bulk/category",
{"images_filter": {"collection": "77"}, "category": "New"},
status=404,
)
self.assertEqual(data["code"], "collection-not-found")
with self.assertRaises(Exception) as caught:
await self.api.resolve_listing_filter({"collection": "77"})
self.assertEqual(type(caught.exception).__name__, "CollectionMissing")
tests/test_api_core.py
-92
View File
@@ -835,7 +835,6 @@ class TestMetadataParsing(unittest.TestCase):
self.assertEqual(parsed["positive_prompt"], "a cat")
self.assertEqual(parsed["negative_prompt"], "dog")
self.assertEqual(parsed["parameters"]["parameters"], meta["parameters"])
self.assertEqual(self.api._extract_readable_prompt(parsed), "a cat")
def test_parse_comfyui_fields(self):
meta = {
@@ -856,97 +855,6 @@ class TestMetadataParsing(unittest.TestCase):
self.assertEqual(parsed["prompt"], "plain text")
self.assertEqual(parsed["workflow"], {"nodes": []})
def test_extract_readable_prompt_variants(self):
extract = self.api._extract_readable_prompt
self.assertEqual(extract({"prompt": "plain"}), "plain")
self.assertEqual(extract({"prompt": ["a", "", "b"]}), "a b")
self.assertEqual(extract({"prompt": 42}), "42")
self.assertEqual(
extract({"prompt": None, "parameters": {"positive": ["x", "y"]}}), "x y"
)
self.assertIsNone(extract({"prompt": {}, "parameters": {}}))
graph = {
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "neg"}},
"3": {
"class_type": "KSampler",
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
},
}
self.assertEqual(extract({"prompt": graph}), "pos")
self.assertEqual(extract({"prompt": {"x": 1}, "workflow": graph}), "pos")
def test_node_inputs_and_text_lookup(self):
self.assertEqual(self.api._get_node_inputs("nope"), {})
self.assertEqual(self.api._get_node_inputs({"inputs": 5}), {})
self.assertEqual(self.api._get_node_inputs({"inputs": {"a": 1}}), {"a": 1})
listed = self.api._get_node_inputs(
{"inputs": [{"name": "text", "link": None}, {"nope": 1}, "junk"]}
)
self.assertEqual(list(listed), ["text"])
self.assertIsNone(self.api._find_text_in_node(None))
self.assertEqual(self.api._find_text_in_node({"inputs": {"text": "hi"}}), "hi")
self.assertEqual(
self.api._find_text_in_node(
{"type": "CLIPTextEncode", "widgets_values": ["w"]}
),
"w",
)
self.assertIsNone(
self.api._find_text_in_node(
{"type": "CLIPTextEncode", "widgets_values": [" "]}
)
)
self.assertIsNone(self.api._find_text_in_node({"type": "KSampler"}))
def test_positive_prompt_from_nodes_array_and_fallbacks(self):
extract = self.api._extract_positive_prompt_from_comfyui_data
self.assertIsNone(extract("x"))
self.assertIsNone(extract({}))
self.assertIsNone(extract({"nodes": []}))
workflow = {
"nodes": [
{"id": 1, "type": "CLIPTextEncode", "widgets_values": ["pos"]},
{
"id": 2,
"type": "CLIPTextEncode",
"widgets_values": ["neg"],
"title": "Negative",
},
{"id": 3, "type": "KSampler", "inputs": [{"name": "seed", "link": 4}]},
"junk",
]
}
self.assertEqual(extract(workflow), "pos")
only_negative = {
"nodes": [
{
"id": 2,
"type": "CLIPTextEncode",
"widgets_values": ["neg"],
"title": "neg",
},
]
}
self.assertEqual(extract(only_negative), "neg")
bad_link = {
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}},
"3": {
"class_type": "KSampler",
"inputs": {"positive": "oops", "negative": []},
},
"4": {
"class_type": "KSampler",
"inputs": {"positive": ["zz", 0], "negative": []},
},
}
self.assertEqual(extract(bad_link), "pos")
if __name__ == "__main__":
unittest.main()
tests/test_api_duplicate_files.py
+36
View File
@@ -687,6 +687,39 @@ class TestCleanupJob(DuplicateFilesTestCase):
"Moved 2 database links to the kept copies", "\n".join(cleanup.log)
)
async def test_duplicate_removal_moves_collection_memberships(self):
self.trio()
db = self.api.db
root_key = duplicate_files.favorite_root_key(self.output_dir)
grids = db.add_to_collection(name="LoRA grids", items=[(root_key, "b.png")])
grids = grids["collection"]["id"]
other = db.add_to_collection(
name="Other", items=[(root_key, "sub/c.png"), (root_key, "d.png")]
)["collection"]["id"]
db.add_to_collection(collection_id=grids, items=[(root_key, "sub/c.png")])
_plan, cleanup = await self.remove_all_but("a.png")
# The union: a.png joins both; c.png's second grids row was a repeat.
self.assertEqual(db.collection_members(grids, root_key), {"a.png"})
self.assertEqual(db.collection_members(other, root_key), {"a.png", "d.png"})
self.assertEqual(sorted(db.file_collections(root_key, "a.png")), [grids, other])
self.assertEqual(db.file_collections(root_key, "b.png"), [])
self.assertEqual(db.file_collections(root_key, "sub/c.png"), [])
self.assertIn(
"Moved 2 collection memberships to the kept copies", "\n".join(cleanup.log)
)
def test_one_collection_membership_is_singular(self):
job = type(
"Job", (), {"log": [], "log_line": lambda s, line: s.log.append(line)}
)()
totals = dict.fromkeys(
("moved", "merged", "favorites", "tags", "review", "collections"), 0
)
duplicate_files.DuplicateFilesMixin._log_moved_records(
job, {**totals, "collections": 1}
)
self.assertEqual(job.log, ["Moved 1 collection membership to the kept copies"])
async def test_a_failed_removal_leaves_its_links_on_the_kept_copy(self):
self.trio()
prompt = self.save("Stays on c")
@@ -843,8 +876,11 @@ class TestLinkQueries(DuplicateFilesTestCase):
with self.subTest(sql=sql):
self.assertNotIn("SCAN generated_images", plan)
self.assertNotIn("SCAN image_favorites", plan)
self.assertNotIn("SCAN image_collection_items", plan)
finally:
conn.close()
statements = duplicate_files.query_plan_sql()
self.assertTrue(any("image_collection_items" in sql for sql in statements))
class TestReviewHelpers(DuplicateFilesTestCase):
tests/test_api_workspace_images.py
+168
View File
@@ -2494,3 +2494,171 @@ class TestListingFilterSelection(ImageAPITestCase):
self.assertEqual(response.status, 409)
self.assertEqual((await response.json())["code"], "source-changed")
class TestWorkspaceCollectionFilter(ImageAPITestCase):
"""#331: ``collection=<id>`` keeps only that collection's listed members."""
media = TestWorkspaceImages.media
get = TestWorkspaceImages.get
_names = TestWorkspaceImages._names
async def tearDownAsync(self):
self.api.db.close_all()
assert_closed(self, self._temp_db.name, self._output_tmp.name)
await super().tearDownAsync()
def keys(self, root, relative):
from py.api.images import _favorite_relative_key, _favorite_root_key
return (_favorite_root_key(root), _favorite_relative_key(relative))
def collect(self, name, *members):
"""A collection holding ``(root, relative)`` members; its id."""
result = self.api.db.add_to_collection(
name=name, items=[self.keys(root, rel) for root, rel in members]
)
return result["collection"]["id"]
def two_sources(self):
second = self.comfy_dir / "second"
second.mkdir()
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
return second
async def test_listing_keeps_only_collection_members(self):
second = self.two_sources()
self.media("a.png", 300)
self.media("b.png", 200)
self.media("sub/c.png", 100, root=second)
self.media("d.png", 50, root=second)
grids = self.collect(
"Grids",
(self.output_dir, "a.png"),
(second, "sub/c.png"),
(self.output_dir, "gone.png"), # stored, never listed
)
self.collect("Other", (self.output_dir, "b.png"))
response, data = await self.get(collection=str(grids))
self.assertEqual(response.status, 200, data)
self.assertEqual(self._names(data), ["a.png", "c.png"])
self.assertEqual(data["total"], 2)
_, data = await self.get(collection=str(grids), root="1")
self.assertEqual(self._names(data), ["c.png"])
# Without the filter the listing is unchanged.
_, data = await self.get()
self.assertEqual(data["total"], 4)
async def test_collection_filter_combines_with_the_other_filters(self):
night = self._save_prompt("Neon zebra", category="Night", tags=["a"])
day = self._save_prompt("Morning field", category="Day", tags=["a"])
files = {
name: self.media(name, 100 + index)
for index, name in enumerate(
("night.png", "day.png", "zebra_plain.png", "clip.mp4", "out.png")
)
}
self.api.db.link_image_to_prompt(night, str(files["night.png"]))
self.api.db.link_image_to_prompt(day, str(files["day.png"]))
members = [(self.output_dir, name) for name in files if name != "out.png"]
grids = str(self.collect("Grids", *members))
cases = (
({}, ["clip.mp4", "day.png", "night.png", "zebra_plain.png"]),
({"q": "zebra"}, ["night.png", "zebra_plain.png"]),
({"category": "Night"}, ["night.png"]),
({"tags": "a"}, ["day.png", "night.png"]),
({"media": "video"}, ["clip.mp4"]),
({"media": "image", "q": "day"}, ["day.png"]),
)
for params, expected in cases:
with self.subTest(**params):
response, data = await self.get(collection=grids, sort="name", **params)
self.assertEqual(response.status, 200, data)
self.assertEqual(self._names(data), expected)
self.assertEqual(data["total"], len(expected))
async def test_invalid_collection_is_a_400(self):
for value in ("x", "0", "-1", "1.5", "1234567890", " 1"):
with self.subTest(value=value):
response, data = await self.get(collection=value)
self.assertEqual(response.status, 400, data)
async def test_unknown_collection_is_a_404(self):
self.media("a.png")
response, data = await self.get(collection="77")
self.assertEqual(response.status, 404)
self.assertEqual(data["code"], "collection-not-found")
self.assertNotIn("images", data)
self.api._get_all_output_dirs = lambda: []
response, data = await self.get(collection="77")
self.assertEqual((response.status, data["code"]), (404, "collection-not-found"))
async def test_a_collection_deleted_during_the_read_is_a_404(self):
self.media("a.png")
grids = self.collect("Grids", (self.output_dir, "a.png"))
real = self.api.db.collection_members
def deleted_meanwhile(collection_id, root_key):
self.api.db.delete_collection(collection_id)
return real(collection_id, root_key)
with patch.object(
self.api.db, "collection_members", side_effect=deleted_meanwhile
):
response, data = await self.get(collection=str(grids))
self.assertEqual((response.status, data["code"]), (404, "collection-not-found"))
async def test_a_failed_members_read_is_a_500_never_an_unfiltered_page(self):
self.media("a.png")
self.media("b.png")
grids = self.collect("Grids", (self.output_dir, "a.png"))
def broken(*_args, **_kwargs):
raise sqlite3.OperationalError("database is locked")
with patch.object(self.api.db, "collection_members", side_effect=broken):
response, data = await self.get(collection=str(grids))
self.assertEqual(response.status, 500)
self.assertNotIn("images", data)
self.assertNotIn("locked", json.dumps(data))
async def test_reordering_sources_keeps_members(self):
second = self.two_sources()
self.media("a.png", root=second)
self.media("b.png")
grids = str(self.collect("Grids", (second, "a.png")))
_, data = await self.get(collection=grids)
self.assertEqual(
[(i["root_index"], i["path"]) for i in data["images"]], [(1, "a.png")]
)
self.api._get_all_output_dirs = lambda: [second, self.output_dir]
_, data = await self.get(collection=grids)
self.assertEqual(
[(i["root_index"], i["path"]) for i in data["images"]], [(0, "a.png")]
)
async def test_cursor_from_another_collection_is_refused(self):
for index, name in enumerate(("a.png", "b.png", "c.png")):
self.media(name, 100 + index)
everything = [(self.output_dir, name) for name in ("a.png", "b.png", "c.png")]
first = str(self.collect("First", *everything))
second = str(self.collect("Second", *everything))
_, page = await self.get(collection=first, limit=1)
token = page["next_cursor"]
self.assertIsInstance(token, str)
response, _ = await self.get(collection=first, limit=1, cursor=token)
self.assertEqual(response.status, 200)
for params in ({"collection": second}, {}, {"favorites": "1"}):
with self.subTest(**params):
response, _ = await self.get(cursor=token, limit=1, **params)
self.assertEqual(response.status, 400)
async def test_collection_wins_over_favorites(self):
self.media("a.png")
self.media("b.png")
grids = str(self.collect("Grids", (self.output_dir, "a.png")))
from py.api.images import _favorite_root_key
self.api.db.set_image_favorite(_favorite_root_key(self.output_dir), "b.png")
_, data = await self.get(collection=grids, favorites="1")
self.assertEqual(self._names(data), ["a.png"])
tests/test_database_backup.py
+44
View File
@@ -614,6 +614,50 @@ class TestRestoreFromFile(BackupTestCase):
)
assert_closed(self, str(backups[0]))
def test_backup_and_restore_keep_collections(self):
with self.db.model.get_connection() as conn:
conn.execute(
"INSERT INTO image_collections (id, name, name_key, created_at) "
"VALUES (4, 'LoRa testing', 'lora testing', '2026-10-27T00:00:00Z')"
)
conn.execute(
"INSERT INTO image_collection_items VALUES "
"(4, 'root', 'grid.png', '2026-10-27T00:00:00Z')"
)
with tempfile.TemporaryDirectory() as directory:
# The rows still sit in the WAL; a backup must carry them.
backup = os.path.join(directory, "backup.db")
self.assertTrue(self.db.model.backup_database(backup))
self.assertEqual(
PromptModel.verify_database_file(backup, require_library_state=True),
(True, "ok"),
)
with self.db.model.get_connection() as conn:
conn.execute("DELETE FROM image_collections")
self.db.model.restore_from_file(backup)
conn = self.db.model.get_connection()
self.assertEqual(
[
tuple(r)
for r in conn.execute("SELECT id, name FROM image_collections")
],
[(4, "LoRa testing")],
)
self.assertEqual(
[
tuple(r)
for r in conn.execute(
"SELECT collection_id, root_key, relative_path "
"FROM image_collection_items"
)
],
[(4, "root", "grid.png")],
)
PromptModel.close_all_instances()
assert_closed(self, directory)
def test_restore_without_a_live_database_returns_empty_backup_path(self):
src = self._make_source_db(["fresh"])
self.db.close_all()
tests/test_image_collections.py
+223
View File
@@ -0,0 +1,223 @@
"""Collections (#331): the database operations."""
import sqlite3
import tempfile
import unittest
from pathlib import Path
try:
from tests.open_handles import assert_closed
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from database.models import PromptModel
from database.operations import PromptDatabase
try: # guarded so a missing module fails assertions, not test loading
from database.image_collections import (
CollectionFilesMissing,
CollectionInvalid,
CollectionMissing,
CollectionNameTaken,
READ_SQL,
)
except ImportError:
CollectionFilesMissing = CollectionMissing = CollectionNameTaken = None
CollectionInvalid = ValueError
READ_SQL = None
ROOT = "/real/output"
OTHER = "/real/other"
class ImageCollectionOperationsTests(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.addCleanup(assert_closed, self, self.temp.name)
self.addCleanup(PromptModel.reset_schema_cache)
self.db = PromptDatabase(str(Path(self.temp.name) / "prompts.db"))
self.addCleanup(self.db.close_all)
def require(self):
self.assertTrue(
hasattr(self.db, "add_to_collection"),
"ImageCollectionOperations is not implemented",
)
def add(self, items, **target):
return self.db.add_to_collection(items=items, **target)
def test_add_creates_a_collection_by_name(self):
self.require()
result = self.add(
[(ROOT, "a.png"), (ROOT, "b.png"), (OTHER, "a.png")], name="LoRa testing"
)
collection = result["collection"]
self.assertEqual(collection["name"], "LoRa testing")
self.assertEqual(
result,
{"collection": collection, "created": True, "added": 3, "already": 0},
)
self.assertEqual(self.db.list_collections(), [collection])
self.assertTrue(self.db.collection_exists(collection["id"]))
self.assertEqual(self.db.get_collection(collection["id"]), collection)
self.assertEqual(
self.db.collection_members(collection["id"], ROOT), {"a.png", "b.png"}
)
self.assertEqual(
self.db.root_collection_members(OTHER), {collection["id"]: {"a.png"}}
)
self.assertEqual(self.db.collection_counts(), {collection["id"]: 3})
self.assertEqual(self.db.file_collections(ROOT, "a.png"), [collection["id"]])
self.assertEqual(self.db.file_collections(ROOT, "zzz.png"), [])
def test_add_by_name_ignores_case_and_unicode_form(self):
self.require()
first = self.add([(ROOT, "a.png")], name="Café grids")["collection"]
# NFD spelling (macOS) in another case joins the same collection.
second = self.add([(ROOT, "b.png")], name="CAFÉ GRIDS")
self.assertEqual(second["collection"], first)
self.assertFalse(second["created"])
self.assertEqual(len(self.db.list_collections()), 1)
self.assertEqual(
self.db.collection_members(first["id"], ROOT), {"a.png", "b.png"}
)
def test_add_counts_already_present(self):
self.require()
made = self.add([(ROOT, "a.png"), (ROOT, "b.png")], name="Set")
result = self.add(
[(ROOT, "b.png"), (ROOT, "c.png"), (ROOT, "c.png")],
collection_id=made["collection"]["id"],
)
self.assertEqual(
(result["created"], result["added"], result["already"]), (False, 1, 1)
)
self.assertEqual(self.db.collection_counts(), {made["collection"]["id"]: 3})
def test_add_with_a_new_name_and_nothing_inserted_creates_nothing(self):
self.require()
with self.assertRaises(CollectionFilesMissing):
self.add([], name="Empty")
self.assertEqual(self.db.list_collections(), [])
# An existing name with nothing to add is not an error.
made = self.add([(ROOT, "a.png")], name="Kept")["collection"]
self.assertEqual(
self.add([], name="kept"),
{"collection": made, "created": False, "added": 0, "already": 0},
)
def test_add_to_a_deleted_collection_raises_missing(self):
self.require()
made = self.add([(ROOT, "a.png")], name="Gone")["collection"]
self.db.delete_collection(made["id"])
with self.assertRaises(CollectionMissing):
self.add([(ROOT, "b.png")], collection_id=made["id"])
with self.assertRaises(CollectionMissing):
self.db.remove_from_collection(made["id"], [(ROOT, "a.png")])
with self.assertRaises(CollectionMissing):
self.db.rename_collection(made["id"], "Back")
self.assertFalse(self.db.collection_exists(made["id"]))
self.assertIsNone(self.db.get_collection(made["id"]))
def test_remove_returns_the_count(self):
self.require()
made = self.add([(ROOT, "a.png"), (ROOT, "b.png")], name="Set")["collection"]
result = self.db.remove_from_collection(
made["id"], [(ROOT, "a.png"), (ROOT, "a.png"), (ROOT, "never.png")]
)
self.assertEqual(result, {"collection": made, "removed": 1})
self.assertEqual(self.db.collection_members(made["id"], ROOT), {"b.png"})
# Removing the last member keeps the (now empty) collection.
self.db.remove_from_collection(made["id"], [(ROOT, "b.png")])
self.assertEqual(self.db.list_collections(), [made])
self.assertEqual(self.db.collection_counts(), {made["id"]: 0})
def test_rename_refuses_a_taken_name(self):
self.require()
first = self.add([(ROOT, "a.png")], name="First")["collection"]
second = self.add([(ROOT, "a.png")], name="Second")["collection"]
with self.assertRaises(CollectionNameTaken) as caught:
self.db.rename_collection(second["id"], " first ")
self.assertEqual(caught.exception.collection_id, first["id"])
self.assertEqual(self.db.get_collection(second["id"]), second)
def test_rename_accepts_a_case_variant_of_its_own_name(self):
self.require()
made = self.add([(ROOT, "a.png")], name="lora grids")["collection"]
renamed = self.db.rename_collection(made["id"], "LoRA grids ")
self.assertEqual(renamed, {"id": made["id"], "name": "LoRA grids"})
self.assertEqual(self.db.list_collections(), [renamed])
other = self.db.rename_collection(made["id"], "Portraits")
self.assertEqual(other["name"], "Portraits")
def test_delete_unknown_raises_missing(self):
self.require()
made = self.add([(ROOT, "a.png")], name="Doomed")["collection"]
self.assertTrue(self.db.delete_collection(made["id"]))
self.assertEqual(self.db.root_collection_members(ROOT), {})
with self.assertRaises(CollectionMissing):
self.db.delete_collection(made["id"])
with self.assertRaises(CollectionMissing):
self.db.delete_collection(999)
def test_list_is_in_name_key_order(self):
self.require()
for name in ("beta", "Alpha", "été", "Zeta", "alpha two"):
self.add([(ROOT, "a.png")], name=name)
self.assertEqual(
[row["name"] for row in self.db.list_collections()],
["Alpha", "alpha two", "beta", "Zeta", "été"],
)
def test_names_are_trimmed_and_collapsed(self):
self.require()
made = self.add([(ROOT, "a.png")], name=" LoRa \t testing\n ")["collection"]
self.assertEqual(made["name"], "LoRa testing")
for bad in ("", " ", "x" * 101, "bell\x07", 5, None):
with self.subTest(name=bad):
with self.assertRaises(CollectionInvalid):
self.add([(ROOT, "a.png")], name=bad)
self.assertEqual(len(self.db.list_collections()), 1)
def test_targets_and_items_are_validated(self):
self.require()
made = self.add([(ROOT, "a.png")], name="Set")["collection"]
for target in (
{},
{"name": "x", "collection_id": made["id"]},
{"collection_id": True},
{"collection_id": 0},
{"collection_id": "1"},
):
with self.subTest(target=target):
with self.assertRaises(CollectionInvalid):
self.add([(ROOT, "a.png")], **target)
for items in ([("root",)], [(ROOT, 5)], "a.png", [(ROOT, "x" * 4097)]):
with self.subTest(items=items):
with self.assertRaises(CollectionInvalid):
self.add(items, collection_id=made["id"])
def test_reads_use_an_index(self):
self.require()
self.add([(ROOT, "a.png")], name="Set")
conn = sqlite3.connect(self.db.model.db_path)
try:
for sql in READ_SQL:
plan = " ".join(
row[3]
for row in conn.execute(
"EXPLAIN QUERY PLAN " + sql, [None] * sql.count("?")
)
)
with self.subTest(sql=sql):
self.assertNotIn("SCAN image_collection_items", plan)
if "?" in sql: # a lookup, not the whole list
self.assertNotIn("SCAN image_collections", plan)
finally:
conn.close()
if __name__ == "__main__":
unittest.main()
tests/test_image_collections_schema.py
+485
View File
@@ -0,0 +1,485 @@
"""Collections schema (#331): install, validation and restore safety."""
import re
import sqlite3
import tempfile
import unittest
import unicodedata
from contextlib import closing
from pathlib import Path
from unittest.mock import patch
try:
from tests.open_handles import assert_closed
from tests.test_annotation_schema import LEGACY_TABLES, table_rows
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from test_annotation_schema import LEGACY_TABLES, table_rows
from database.models import PromptModel
try: # guarded so a missing module fails assertions, not test loading
from database import image_collections
from database.image_collections import (
image_collections_schema_error,
migrate_image_collections,
)
except ImportError:
image_collections = None
image_collections_schema_error = None
migrate_image_collections = None
COLLECTIONS = "image_collections"
ITEMS = "image_collection_items"
INDEX = "idx_image_collection_items_file"
TABLES = (COLLECTIONS, ITEMS)
DDL = {
COLLECTIONS: (
f"CREATE TABLE {COLLECTIONS} (id INTEGER PRIMARY KEY AUTOINCREMENT, "
"name TEXT NOT NULL, name_key TEXT NOT NULL UNIQUE, "
"created_at TEXT NOT NULL)"
),
ITEMS: (
f"CREATE TABLE {ITEMS} (collection_id INTEGER NOT NULL REFERENCES "
f"{COLLECTIONS}(id) ON DELETE CASCADE, root_key TEXT NOT NULL, "
"relative_path TEXT NOT NULL, added_at TEXT NOT NULL, "
"PRIMARY KEY (collection_id, root_key, relative_path))"
),
}
INDEX_DDL = f"CREATE INDEX {INDEX} ON {ITEMS}(root_key, relative_path)"
NOW = "2026-10-27T00:00:00.000000+00:00"
IMAGES = 23_000
def _compact(sql):
return re.sub(r"\s+", "", sql or "").lower()
def expected_catalog():
return {
("table", COLLECTIONS, _compact(DDL[COLLECTIONS])),
("table", ITEMS, _compact(DDL[ITEMS])),
("index", INDEX, _compact(INDEX_DDL)),
("index", f"sqlite_autoindex_{COLLECTIONS}_1", ""),
("index", f"sqlite_autoindex_{ITEMS}_1", ""),
}
def catalog(conn):
return {
(kind, name, _compact(sql))
for kind, name, sql in conn.execute(
"SELECT type,name,sql FROM sqlite_master "
"WHERE name IN (?,?,?) OR tbl_name IN (?,?)",
(*TABLES, INDEX, *TABLES),
)
}
def key(name):
return unicodedata.normalize("NFC", unicodedata.normalize("NFC", name).casefold())
def schema(conn, exclude=()):
return sorted(
tuple(row)
for row in conn.execute(
"SELECT type,name,tbl_name,sql FROM sqlite_master ORDER BY name"
)
if row[1] not in exclude and row[2] not in exclude
)
def row_counts(conn):
tables = [
name
for (name,) in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' ORDER BY name"
)
]
return {
table: conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] # nosec
for table in tables
}
class ImageCollectionsSchemaTests(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.addCleanup(assert_closed, self, self.temp.name)
self.addCleanup(PromptModel.reset_schema_cache)
self.models = []
self.addCleanup(self.close_models)
self.path = str(Path(self.temp.name) / "live.db")
def close_models(self):
for model in self.models:
model.close_all()
def model(self, path=None):
model = PromptModel(path or self.path)
self.models.append(model)
return model
def require_module(self):
self.assertIsNotNone(
image_collections_schema_error,
"database.image_collections is not implemented",
)
def closed_current(self, name="current.db", *statements):
"""A closed current database, altered afterwards by *statements*."""
path = str(Path(self.temp.name) / name)
model = self.model(path)
model.close_all()
PromptModel.reset_schema_cache()
with closing(sqlite3.connect(path)) as raw, raw:
raw.execute("PRAGMA foreign_keys=OFF")
for statement in statements:
if isinstance(statement, tuple):
raw.execute(*statement)
else:
raw.execute(statement)
return path
@staticmethod
def drop_group(path):
with closing(sqlite3.connect(path)) as raw, raw:
raw.execute(f"DROP TABLE {ITEMS}")
raw.execute(f"DROP TABLE {COLLECTIONS}")
def test_migration_creates_tables_and_index(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
self.assertEqual(catalog(conn), expected_catalog())
self.assertIsNone(image_collections_schema_error(conn, allow_missing=False))
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {COLLECTIONS}").fetchone()[0], 0
)
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {ITEMS}").fetchone()[0], 0
)
self.assertEqual(
PromptModel.verify_database_file(self.path, require_library_state=True),
(True, "ok"),
)
def test_migration_is_idempotent(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
conn.execute(
f"INSERT INTO {COLLECTIONS} (name, name_key, created_at) "
"VALUES ('Kept', 'kept', ?)",
(NOW,),
)
conn.execute(f"INSERT INTO {ITEMS} VALUES (1, 'root', 'a.png', ?)", (NOW,))
before = schema(conn)
model.close_all()
for _ in range(2):
PromptModel.reset_schema_cache()
model = self.model()
with model.get_connection() as conn:
migrate_image_collections(conn)
self.assertEqual(schema(conn), before)
self.assertEqual(
[tuple(r) for r in conn.execute(f"SELECT * FROM {ITEMS}")],
[(1, "root", "a.png", NOW)],
)
model.close_all()
def test_migration_reads_only_its_own_tables(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
conn.execute("INSERT INTO prompts(id,text,hash) VALUES(1,'positive','p')")
conn.executemany(
"INSERT INTO generated_images(prompt_id,image_path,filename,"
"file_path) VALUES(1,?,?,?)",
((f"/out/{n}.png", f"{n}.png", f"/out/{n}.png") for n in range(IMAGES)),
)
conn.execute("INSERT INTO image_favorites VALUES('root','a.png',?)", (NOW,))
model.close_all()
self.drop_group(self.path)
with closing(sqlite3.connect(self.path)) as raw:
raw.execute("PRAGMA foreign_keys=ON")
counts_before = row_counts(raw)
schema_before = schema(raw)
legacy_before = table_rows(raw, LEGACY_TABLES)
self.assertEqual(counts_before["generated_images"], IMAGES)
touched = []
def record(action, arg1, arg2, _db, _trigger):
touched.append((action, arg1, arg2))
return sqlite3.SQLITE_OK
raw.set_authorizer(record)
try:
migrate_image_collections(raw)
finally:
raw.set_authorizer(None)
raw.commit()
reads = {
arg1 for action, arg1, _arg2 in touched if action == sqlite3.SQLITE_READ
}
self.assertTrue(reads)
self.assertLessEqual(reads, {"sqlite_master", COLLECTIONS, ITEMS})
writes = {
arg1
for action, arg1, _arg2 in touched
if action
in (sqlite3.SQLITE_INSERT, sqlite3.SQLITE_UPDATE, sqlite3.SQLITE_DELETE)
}
self.assertLessEqual(writes, {"sqlite_master"})
self.assertEqual(catalog(raw), expected_catalog())
counts_after = row_counts(raw)
self.assertEqual(counts_after.pop(COLLECTIONS), 0)
self.assertEqual(counts_after.pop(ITEMS), 0)
self.assertEqual(counts_after, counts_before)
self.assertEqual(schema(raw, exclude=TABLES), schema_before)
self.assertEqual(table_rows(raw, LEGACY_TABLES), legacy_before)
# A second run is a no-op.
after_first = schema(raw)
changes = raw.total_changes
migrate_image_collections(raw)
raw.commit()
self.assertEqual(schema(raw), after_first)
self.assertEqual(raw.total_changes, changes)
def test_migration_refuses_a_foreign_object_of_the_same_name(self):
self.require_module()
for statements in (
(f"CREATE TABLE {COLLECTIONS} (x)",),
(f"CREATE VIEW {COLLECTIONS} AS SELECT 1",),
(f"CREATE INDEX {INDEX} ON prompts(text)",),
):
with self.subTest(statements=statements):
path = str(Path(self.temp.name) / f"foreign{len(self.models)}.db")
self.model(path).close_all()
PromptModel.reset_schema_cache()
self.drop_group(path)
with closing(sqlite3.connect(path)) as raw, raw:
for statement in statements:
raw.execute(statement)
with closing(sqlite3.connect(path)) as raw:
before = schema(raw)
self.assertIsNotNone(image_collections_schema_error(raw))
with self.assertRaises(sqlite3.DatabaseError):
self.model(path).get_connection()
PromptModel.close_all_instances()
PromptModel.reset_schema_cache()
with closing(sqlite3.connect(path)) as raw:
self.assertEqual(schema(raw), before)
self.assertIsNone(
raw.execute(
"SELECT 1 FROM sqlite_master WHERE name=?", (ITEMS,)
).fetchone()
)
self.assertFalse(PromptModel.verify_database_file(path)[0])
def test_migration_rolls_back_on_error(self):
self.require_module()
self.model().close_all()
PromptModel.reset_schema_cache()
self.drop_group(self.path)
with closing(sqlite3.connect(self.path)) as raw:
with patch.object(
image_collections,
"image_collections_schema_error",
side_effect=[None, "Image collections went wrong"],
):
with self.assertRaisesRegex(sqlite3.DatabaseError, "went wrong"):
migrate_image_collections(raw)
raw.commit()
self.assertEqual(catalog(raw), set())
def test_startup_accepts_a_value_saved_under_older_rules(self):
self.require_module()
path = self.closed_current(
"older.db",
(
f"INSERT INTO {COLLECTIONS} (name, name_key, created_at) "
"VALUES (?, 'not the key', ?)",
(" " + "x" * 150 + " ", NOW),
),
)
model = self.model(path)
with model.get_connection() as conn:
self.assertIsNone(image_collections_schema_error(conn, check_values=False))
self.assertIsNotNone(image_collections_schema_error(conn))
self.assertEqual(
conn.execute(f"SELECT COUNT(*) FROM {COLLECTIONS}").fetchone()[0], 1
)
model.close_all()
self.assertFalse(
PromptModel.verify_database_file(path, require_library_state=True)[0]
)
def test_restore_accepts_an_old_backup_without_collections(self):
self.require_module()
path = self.closed_current("old.db")
self.drop_group(path)
with closing(sqlite3.connect(path)) as raw:
self.assertIsNone(image_collections_schema_error(raw))
self.assertEqual(
image_collections_schema_error(raw, allow_missing=False),
"Missing image collections schema",
)
self.assertEqual(PromptModel.verify_database_file(path), (True, "ok"))
model = self.model(path)
with model.get_connection() as conn:
self.assertEqual(catalog(conn), expected_catalog())
model.close_all()
self.assertEqual(
PromptModel.verify_database_file(path, require_library_state=True),
(True, "ok"),
)
def test_restore_requires_collections_on_a_current_stage(self):
self.require_module()
path = self.closed_current("stage.db")
self.drop_group(path)
self.assertEqual(
PromptModel.verify_database_file(path, require_library_state=True),
(False, "Missing image collections schema"),
)
# One table without the other is never acceptable.
partial = self.closed_current("partial.db", f"DROP TABLE {ITEMS}")
self.assertFalse(PromptModel.verify_database_file(partial)[0])
def assert_refused(self, name, *statements):
self.require_module()
path = self.closed_current(
f"{name}.db",
(
f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) "
"VALUES (1, 'Good', 'good', ?)",
(NOW,),
),
(f"INSERT INTO {ITEMS} VALUES (1, 'root', 'a.png', ?)", (NOW,)),
*statements,
)
with closing(sqlite3.connect(path)) as raw:
self.assertIsNone(image_collections_schema_error(raw, check_values=False))
error = image_collections_schema_error(raw)
self.assertIsNotNone(error)
self.assertIn("collections", error.lower())
ok, reason = PromptModel.verify_database_file(path)
self.assertFalse(ok)
self.assertEqual(reason, error)
self.assertFalse(
PromptModel.verify_database_file(path, require_library_state=True)[0]
)
def test_a_valid_current_database_is_accepted(self):
self.require_module()
path = self.closed_current(
"valid.db",
(
f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) "
"VALUES (1, ?, ?, ?)",
("LoRa testing", key("LoRa testing"), NOW),
),
(f"INSERT INTO {ITEMS} VALUES (1, 'root', ?, ?)", ("é" * 4096, NOW)),
)
self.assertEqual(
PromptModel.verify_database_file(path, require_library_state=True),
(True, "ok"),
)
def test_restore_refuses_a_non_text_name(self):
self.assert_refused("int-name", f"UPDATE {COLLECTIONS} SET name = 5")
def test_restore_refuses_an_empty_name(self):
self.assert_refused("empty", f"UPDATE {COLLECTIONS} SET name = '', name_key=''")
def test_restore_refuses_a_name_over_100_characters(self):
name = "é" * 101
self.assert_refused(
"long",
(f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", (name, key(name))),
)
def test_restore_refuses_a_control_character(self):
name = "a\x07b"
self.assert_refused(
"control",
(f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?", (name, key(name))),
)
def test_restore_refuses_a_leading_or_trailing_space(self):
for name in (" Good", "Good "):
with self.subTest(name=name):
self.assert_refused(
f"padded{len(name)}{name[0] == ' '}",
(
f"UPDATE {COLLECTIONS} SET name = ?, name_key = ?",
(name, key(name)),
),
)
def test_restore_refuses_a_wrong_name_key(self):
self.assert_refused("key", f"UPDATE {COLLECTIONS} SET name_key = 'Good'")
def test_restore_refuses_a_non_text_key(self):
# Column affinity turns a number into text; a blob stays a blob.
self.assert_refused("blob-key", f"UPDATE {ITEMS} SET relative_path = x'61'")
def test_restore_refuses_a_key_over_4096_characters(self):
self.assert_refused(
"long-key", (f"UPDATE {ITEMS} SET root_key = ?", ("é" * 4097,))
)
def test_restore_refuses_a_timestamp_over_64_bytes(self):
self.assert_refused(
"stamp", (f"UPDATE {COLLECTIONS} SET created_at = ?", ("t" * 65,))
)
self.assert_refused(
"item-stamp", (f"UPDATE {ITEMS} SET added_at = ?", ("t" * 65,))
)
def test_restore_refuses_an_item_of_a_missing_collection(self):
self.assert_refused("orphan", f"INSERT INTO {ITEMS} VALUES (9, 'r', 'b', 'n')")
def test_restore_refuses_a_non_integer_collection_id(self):
self.assert_refused("text-id", f"UPDATE {ITEMS} SET collection_id = 'one'")
def test_restore_refuses_a_trigger_or_extra_index(self):
self.require_module()
for name, statement in (
(
"trigger",
f"CREATE TRIGGER t AFTER INSERT ON {ITEMS} BEGIN SELECT 1; END",
),
("extra", f"CREATE INDEX extra ON {COLLECTIONS}(name)"),
):
with self.subTest(name=name):
path = self.closed_current(f"{name}.db", statement)
with closing(sqlite3.connect(path)) as raw:
self.assertIsNotNone(image_collections_schema_error(raw))
self.assertFalse(PromptModel.verify_database_file(path)[0])
def test_deleting_a_collection_deletes_its_members(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
conn.execute(
f"INSERT INTO {COLLECTIONS} (id, name, name_key, created_at) "
"VALUES (1, 'A', 'a', ?), (2, 'B', 'b', ?)",
(NOW, NOW),
)
conn.executemany(
f"INSERT INTO {ITEMS} VALUES (?, 'root', ?, ?)",
((1, "a.png", NOW), (1, "b.png", NOW), (2, "a.png", NOW)),
)
conn.execute(f"DELETE FROM {COLLECTIONS} WHERE id = 1")
self.assertEqual(
[tuple(r) for r in conn.execute(f"SELECT collection_id FROM {ITEMS}")],
[(2,)],
)
if __name__ == "__main__":
unittest.main()
tests/test_image_tag_jobs.py
+10
View File
@@ -884,6 +884,16 @@ class TestFilterStart(ImageTagJobTestCase):
self.assertEqual((status, data["code"]), (400, "no-images"))
self.assertEqual(self.api._jobs.list(), [])
async def test_filter_with_a_deleted_collection_is_a_404_without_a_job(self):
self.png("a.png")
status, data = await self.start(
{"mode": "add", "filter": {"collection": "77"}, "tags": ["x"]}
)
self.assertEqual((status, data["code"]), (404, "collection-not-found"))
self.assertEqual(self.api._jobs.list(), [])
async def test_filter_over_the_cap_is_refused_without_a_job(self):
from py.api import image_tag_jobs
tests/test_pair_scan.py
+158 -1
View File
@@ -115,7 +115,12 @@ class TestPairScan(AdminAPITestCase):
self.assertEqual(second["pairs_added"], 0)
def test_legacy_single_prompt_second_pass_not_reported_as_new_link(self):
meta = {"prompt": json.dumps({"1": graph()["1"]})}
# A sampler-less graph: one encoder wired into a guider (#252).
single = {
"1": graph()["1"],
"2": {"class_type": "BasicGuider", "inputs": {"conditioning": ["1", 0]}},
}
meta = {"prompt": json.dumps(single)}
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")
@@ -292,3 +297,155 @@ class TestPairScan(AdminAPITestCase):
self.api.db.get_image_prompt_info(str(path))["workflow_data"], workflow
)
self.assertEqual(self.api._ingest_scanned_file(path, metadata), "unchanged")
def flux_graph(text="a red fox", encoder="CLIPTextEncode"):
"""A sampler-less Flux graph: the prompt reaches a guider's conditioning (#252)."""
return {
"1": {"class_type": "UNETLoader", "inputs": {"unet_name": "flux.safetensors"}},
"2": {"class_type": "DualCLIPLoader", "inputs": {"clip_name1": "t5"}},
"6": {"class_type": encoder, "inputs": {"text": text, "clip": ["2", 0]}},
"22": {
"class_type": "BasicGuider",
"inputs": {"model": ["1", 0], "conditioning": ["6", 0]},
},
"13": {
"class_type": "SamplerCustomAdvanced",
"inputs": {"guider": ["22", 0], "sigmas": ["17", 0]},
},
"17": {"class_type": "BasicScheduler", "inputs": {"steps": 20}},
"9": {"class_type": "SaveImage", "inputs": {"images": ["13", 0]}},
}
class TestSamplerLessScan(AdminAPITestCase):
"""Graphs where no node takes a positive or negative input (#252)."""
def _scanned(self, name, prompt_graph, workflow=None):
meta = {"prompt": json.dumps(prompt_graph)}
if workflow is not None:
meta["workflow"] = json.dumps(workflow)
return self.api._ingest_scanned_file(self.output_dir / name, meta)
def _texts(self):
return [p["text"] for p in self.api.db.get_recent_prompts()["prompts"]]
def test_the_one_wired_prompt_is_linked(self):
self.assertEqual(self._scanned("flux.png", flux_graph()), "added")
self.assertEqual(self._texts(), ["a red fox"])
def test_a_text_node_feeding_the_encoder_is_one_prompt(self):
g = flux_graph(text=["30", 0])
g["30"] = {"class_type": "TextInput", "inputs": {"text": "a red fox"}}
self.assertEqual(self._scanned("textinput.png", g), "added")
self.assertEqual(self._texts(), ["a red fox"])
def test_a_display_node_that_feeds_nothing_is_never_the_prompt(self):
# A real scan stored a "=== GIT CLONE REPORT ===" display text this way.
g = flux_graph()
g["40"] = {
"class_type": "easy showAnything",
"inputs": {"text": "=== GIT CLONE REPORT ===\nTotal repositories: 1"},
}
self.assertEqual(self._scanned("report.png", g), "added")
self.assertEqual(self._texts(), ["a red fox"])
def test_only_unconsumed_texts_store_nothing(self):
g = {
"40": {"class_type": "easy showAnything", "inputs": {"text": "report"}},
"41": {"class_type": "Textbox", "inputs": {"text": "install note"}},
}
self.assertEqual(self._scanned("notes.png", g), "found")
self.assertEqual(self._texts(), [])
def test_two_wired_texts_are_ambiguous_and_stay_unlinked(self):
# A LoRA loader's text input is wired too: the scan stored "<lora:...>".
g = flux_graph()
g["50"] = {
"class_type": "Lora Loader (LoraManager)",
"inputs": {
"model": ["1", 0],
"clip": ["2", 0],
"text": "<lora:style:1.00>",
},
}
g["22"]["inputs"]["model"] = ["50", 0]
g["6"]["inputs"]["clip"] = ["50", 1]
self.assertEqual(self._scanned("lora.png", g), "unresolved")
self.assertEqual(self._texts(), [])
def test_an_untraceable_wired_text_makes_the_graph_ambiguous(self):
# The encoder's text comes from a node no rule reads: never guess the other.
g = flux_graph(text=["60", 0])
g["60"] = {"class_type": "SomeLLMNode", "inputs": {"seed": 1}}
g["50"] = {
"class_type": "Lora Loader (LoraManager)",
"inputs": {"model": ["1", 0], "text": "<lora:style:1.00>"},
}
g["22"]["inputs"]["model"] = ["50", 0]
self.assertEqual(self._scanned("llm.png", g), "unresolved")
self.assertEqual(self._texts(), [])
def test_this_repositorys_prompt_node_is_authoritative(self):
g = flux_graph(text="a red fox", encoder="PromptManager")
g["50"] = {
"class_type": "Lora Loader (LoraManager)",
"inputs": {"model": ["1", 0], "text": "<lora:style:1.00>"},
}
g["22"]["inputs"]["model"] = ["50", 0]
self.assertEqual(self._scanned("own.png", g), "added")
self.assertEqual(self._texts(), ["a red fox"])
def test_a_numeric_text_is_never_stored_as_a_prompt(self):
self.assertEqual(self._scanned("number.png", flux_graph(text="2048")), "found")
self.assertEqual(self._texts(), [])
def test_workflow_widget_values_are_never_read(self):
# The API text is linked from a node no rule reads; the UI workflow
# still holds a stale widget value, which the old scan stored.
g = flux_graph(text=["70", 0])
g["70"] = {"class_type": "PromptBuilder", "inputs": {"prompt": "x", "n": 2}}
workflow = {
"nodes": [
{"id": 6, "type": "CLIPTextEncode", "widgets_values": ["stale text"]}
]
}
self.assertEqual(self._scanned("stale.png", g, workflow), "unresolved")
self.assertEqual(self._texts(), [])
def test_a_lora_loaders_trigger_word_output_is_not_its_text(self):
# The encoder reads the loader's trigger-word output (slot 2), not its
# text input: the prompt is unknown, never the "<lora:...>" syntax.
g = flux_graph(text=["50", 2])
g["50"] = {
"class_type": "Lora Loader (LoraManager)",
"inputs": {"model": ["1", 0], "text": "<lora:style:1.00>"},
}
g["22"]["inputs"]["model"] = ["50", 0]
self.assertEqual(self._scanned("trigger.png", g), "unresolved")
self.assertEqual(self._texts(), [])
def test_a_seeded_text_node_holds_a_template_not_the_prompt(self):
g = flux_graph(text=["80", 0])
g["80"] = {
"class_type": "RandomPrompt",
"inputs": {"text": "a {red|blue} fox", "seed": 5},
}
self.assertEqual(self._scanned("random.png", g), "unresolved")
self.assertEqual(self._texts(), [])
def test_a1111_parameters_beside_an_undecided_graph_are_read(self):
# Metadata savers write both; the old scan read the A1111 positive.
g = flux_graph()
g["50"] = {
"class_type": "Lora Loader (LoraManager)",
"inputs": {"model": ["1", 0], "text": "<lora:style:1.00>"},
}
g["22"]["inputs"]["model"] = ["50", 0]
meta = {
"prompt": json.dumps(g),
"parameters": "a red fox\nNegative prompt: blur\nSteps: 20",
}
path = self.output_dir / "saver.png"
self.assertEqual(self.api._ingest_scanned_file(path, meta), "added")
self.assertEqual(self._texts(), ["a red fox"])
tests/test_restore_plans.py
+1
View File
@@ -49,6 +49,7 @@ SUMMARY_KEYS = {
"image_links",
"runs",
"saved_views",
"collections",
}
tests/test_restore_stage.py
+43
View File
@@ -226,6 +226,7 @@ class TestStageAndPublish(StageTestCase):
for key in ("pairs", "archived_pairs", "image_links", "runs"):
self.assertEqual(staged.summary[key], 0, key)
self.assertEqual(staged.summary["saved_views"], 0)
self.assertEqual(staged.summary["collections"], 0)
self.assertEqual(staged.schema, {"migrated": False, "library_state": True})
self.assertEqual(
staged.digest,
@@ -239,6 +240,48 @@ class TestStageAndPublish(StageTestCase):
staged.discard()
class TestCollectionsSummary(StageTestCase):
"""#331: collections join the review's summary and the live fingerprint."""
def add_collection(self, db, name):
with db.model.get_connection() as conn:
conn.execute(
"INSERT INTO image_collections (name, name_key, created_at) "
"VALUES (?, ?, '2026-10-27T00:00:00+00:00')",
(name, name.casefold()),
)
def test_summary_counts_collections(self):
src_path = os.path.join(self.uploads, "with-collections.db")
src = PromptDatabase(src_path)
try:
self.save(src, "kept")
self.add_collection(src, "First")
self.add_collection(src, "Second")
finally:
src.close_all()
staged = self.db.model.stage_restore(src_path)
try:
self.assertEqual(staged.summary["collections"], 2)
finally:
staged.discard()
self.assertIn("collections", restore_stage.SUMMARY_KEYS)
def test_adding_a_collection_changes_the_fingerprint(self):
conn = self.db.model.get_connection()
summary, before = restore_stage.live_fingerprint(conn)
self.assertEqual(summary["collections"], 0)
self.add_collection(self.db, "After the review")
summary, after = restore_stage.live_fingerprint(conn)
self.assertEqual(summary["collections"], 1)
self.assertNotEqual(after, before)
# A delete and a re-add with the same count still changes it (rowid).
with self.db.model.get_connection() as write:
write.execute("DELETE FROM image_collections")
self.add_collection(self.db, "Replacement")
self.assertNotEqual(restore_stage.live_fingerprint(conn)[1], after)
class TestOlderBackups(StageTestCase):
"""Rollback-journal and older-schema backups restore exactly as before."""
tests/test_sampler_less_prompt.py
+112
View File
@@ -0,0 +1,112 @@
"""``sampler_less_prompt``: the one wired text of a graph with no sampler (#252)."""
import time
import unittest
from utils.prompt_pair_resolver import sampler_less_prompt
def _guided(**nodes):
"""``nodes`` plus a guider consuming node ``"1"``."""
graph = {"9": {"class_type": "BasicGuider", "inputs": {"conditioning": ["1", 0]}}}
graph.update(nodes)
return graph
class TestSamplerLessPrompt(unittest.TestCase):
def test_not_a_graph_or_too_large_is_absent(self):
for graph in (None, [], {}, "text"):
self.assertEqual(sampler_less_prompt(graph)["state"], "absent")
big = {str(i): {"class_type": "X", "inputs": {}} for i in range(3)}
self.assertEqual(sampler_less_prompt(big, max_visits=2)["state"], "absent")
def test_one_literal_text(self):
graph = _guided(**{"1": {"class_type": "Enc", "inputs": {"text": " a fox "}}})
self.assertEqual(
sampler_less_prompt(graph),
{"state": "resolved", "text": "a fox", "texts": ["a fox"]},
)
def test_a_number_literal_is_never_a_prompt(self):
for value in (2048, "2048", "1.5", " 7 "):
graph = _guided(**{"1": {"class_type": "Enc", "inputs": {"text": value}}})
self.assertEqual(sampler_less_prompt(graph)["state"], "absent", value)
def test_a_long_digit_run_with_a_letter_is_text_and_fast(self):
text = "9" * 50_000 + "x"
graph = _guided(**{"1": {"class_type": "Enc", "inputs": {"text": text}}})
started = time.monotonic()
self.assertEqual(sampler_less_prompt(graph)["text"], text)
self.assertLess(time.monotonic() - started, 1.0)
def test_number_like_words_are_prompts(self):
for value in ("nan", "inf", "Infinity", "1_000"):
graph = _guided(**{"1": {"class_type": "Enc", "inputs": {"text": value}}})
self.assertEqual(sampler_less_prompt(graph)["text"], value)
def test_an_own_node_with_unreadable_text_beside_another_text_is_ambiguous(self):
graph = _guided(
**{
"1": {"class_type": "PromptManager", "inputs": {"text": ["5", 0]}},
"5": {"class_type": "LLM", "inputs": {"seed": 3}},
"2": {"class_type": "Enc", "inputs": {"text": "<lora:x:1>"}},
"3": {"class_type": "Sink", "inputs": {"c": ["2", 0]}},
}
)
self.assertEqual(sampler_less_prompt(graph)["state"], "ambiguous")
def test_unhashable_class_type_is_an_unknown_node_not_a_crash(self):
graph = _guided(**{"1": {"class_type": ["x"], "inputs": {"text": "a fox"}}})
self.assertEqual(sampler_less_prompt(graph)["text"], "a fox")
def test_own_negative_node_is_not_authoritative(self):
graph = _guided(
**{
"1": {"class_type": "Enc", "inputs": {"text": "a fox"}},
"2": {
"class_type": "PromptManagerNegative",
"inputs": {"text": "blur"},
},
"3": {"class_type": "Sink", "inputs": {"c": ["2", 0]}},
}
)
found = sampler_less_prompt(graph)
self.assertEqual(found["state"], "ambiguous")
self.assertEqual(found["texts"], ["a fox", "blur"])
def test_own_node_text_wins_over_an_untraceable_text(self):
graph = _guided(
**{
"1": {"class_type": "PromptManagerText", "inputs": {"text": "a fox"}},
"2": {"class_type": "Enc", "inputs": {"text": ["5", 0]}},
"5": {"class_type": "LLM", "inputs": {"seed": 3}},
"6": {"class_type": "Sink", "inputs": {"c": ["2", 0]}},
}
)
self.assertEqual(sampler_less_prompt(graph)["text"], "a fox")
def test_two_own_texts_are_ambiguous(self):
graph = _guided(
**{
"1": {"class_type": "PromptManager", "inputs": {"text": "a fox"}},
"2": {"class_type": "PromptManager", "inputs": {"text": "a cat"}},
"3": {"class_type": "Sink", "inputs": {"c": ["2", 0]}},
}
)
self.assertEqual(sampler_less_prompt(graph)["state"], "ambiguous")
def test_only_an_untraceable_text_is_ambiguous_with_nothing_to_show(self):
graph = _guided(
**{
"1": {"class_type": "Enc", "inputs": {"text": ["5", 0]}},
"5": {"class_type": "LLM", "inputs": {"seed": 3}},
}
)
self.assertEqual(
sampler_less_prompt(graph),
{"state": "ambiguous", "text": None, "texts": []},
)
if __name__ == "__main__":
unittest.main()