Vito
6c82098207
* feat(images): link summary on workspace image items
* fix(images): carry title source and cut first-line titles like the Library
Summaries gain title_source (stored or text). The inspector always shows a
stored title and confirms only a first-line title against the file, so a
stored title equal to the first line is no longer withheld. A link-title read
failure no longer marks a readable link unavailable. First-line titles over
80 characters end with an ellipsis, as the Library's title() does. The card
gate for unconfirmed first-line titles is pinned as an expected failure for B2.
* feat(images): unambiguous seed on workspace items
Workspace items gain `seed`, read with the shared generation-parameter
parser from the prompt/workflow records the header open already loaded
(each at most 2 MiB; no extra file reads, videos give None). A
placeholder or non-integer seed gives None. JSON integers beyond 2**53
are recovered exactly, so the card prints the digits the inspector shows.
The same parse confirms first-line card titles: a `text` title shows only
when the file's parsed positive gives the same title. One helper applies
that rule on the card and in the inspector, which now reads the positive
from its own inspection records with the same parser.
* fix(images): a deeply nested record no longer fails the workspace listing
The Python parser recurses over nested arrays where JavaScript String()
does not, so a crafted 2 KB PNG turned every workspace page holding it
(and its inspect route) into a 500. Card facts now treat RecursionError
like any record the parser cannot use: no seed and no confirmed title.
Also: Library previews skip the card parse they never show; the parse
runs after the file handle closes; the exact-seed read uses the one JSON
reader (parse_js_json(..., exact_integers=True)) and the public
number_string; record size is decided from the text length unless only
the UTF-8 byte count can tell.
* feat(images): search and filter workspace images by linked prompts
The workspace listing now takes category (exact), tags (comma list, all
required) and sort=rating, and q also matches the linked positive's text,
Library title, pair title and tag names. One bounded, parameterized SELECT
(get_image_link_matches / get_image_paths_matching) returns normalized
paths of unambiguously linked images; the listing is intersected with it.
Unlinked and ambiguous files are excluded by a category or tag filter and
sort last by rating, with a newest-first tie-break. The cursor scope covers
the new filters and sort; unknown values are still rejected.
The derived listing values and sort orders are cached per gallery scan so a
50,000-file listing with 20,000 links pages well under 150 ms at p95.
* test(images): count the filter budget everywhere, time it in one job
The 150 ms wall-clock budget failed under the CI coverage tracer, so it
is split. test_filter_work_is_bounded runs on every job and under
coverage: one filter SELECT per request, the ambiguity probe is an index
search (EXPLAIN QUERY PLAN), the per-scan values are built on the cold
request only, and a cursor page calls the sort key at most
ceil(log2 n)+3 times. test_filter_budget runs only with
PROMPTMANAGER_PERF_BUDGET=1 and no tracer, asserts the best of three
rounds' p95, and has its own non-required "perf budget" job. TESTING.md
records the convention.
* fix(images): rating sort works however many link rows the database holds
The filter SELECT counted every generated_images row against a 500,000
row cap and materialized all of them first, so a large history made
sort=rating answer 400 with nothing to narrow. The rows are now streamed
and only paths in the listing's normalized set are kept, so memory is
bounded by the listing; the cap and LinkFilterTooBroad are gone.
* fix(images): rebuild workspace listing values only where the scan changed
A TTL rescan that finds the same files reuses the previous values and
sort orders; one that finds a few changes rebuilds only those files.
Gallery invalidation and a gallery-roots change drop the derived values,
so old roots no longer hold memory. category is trimmed like tags.
* fix(images): a monitor invalidation rebuilds only the new files' values
The image monitor invalidates the gallery whenever it sees a new file,
and that dropped every derived workspace value, so the next request paid
a full cold build. Invalidation now keeps the per-file values and drops
only the sort orders; the rescan reuses every unchanged file and builds
only new or changed ones. A gallery-roots change still drops everything.
* feat(library): expose run count, last use and variant totals
* fix(library): an out-of-range last use reads as unknown, not a failed list
* perf(images): read positive counts only on a plain listing's first page
* feat(images): persist image favorites
Favorites are stored in image_favorites keyed by the gallery root label
and the relative path, set and cleared through PUT/DELETE
/prompt_manager/images/favorite with a {root_index, path} body validated
against the rooted listing. Workspace image cards carry favorite: bool.
The table is created by migration, travels with backups, is optional in
old backups and required after a restore's migration.
* fix(images): key favorites by the real root and case-normalized path
Folder names are not unique across gallery roots, so two roots shown as
"output" shared each other's stars. Favorites are now keyed by the
root's real, case-normalized path (server-side only) and a
case-normalized relative path, so a Windows case-only rename keeps its
star while distinct Linux files stay distinct. Toggles look the file up
in a per-scan dictionary instead of walking the listing.
* fix(images): favorites lookup follows monitor invalidation and the stored key bound
A monitor invalidation now also drops the favorites lookup, which held the
previous scan list until the next toggle. PUT refuses a favorite whose root
or relative key would exceed the schema's 4096-character bound, so the live
database keeps passing startup and restore validation.
* docs(testing): workspace data fields for the concept pages
* fix(images): card facts read the inspector's projected graphs
The inspector hands its browser parser only graphs the inspection
projection parsed (strict JSON, depth 64, node ceiling); a deeper record
is "limited" there, so the legacy fill never runs and no seed shows. The
card parsed the raw records instead and showed a seed the inspector did
not, e.g. for a 900-deep class_type beside a normal sampler.
project_graph() is now the one projection step, used by the inspection
response and by the card. The card runs the shared parser on the
projected graphs, the same values the browser receives: big integers
arrive as their exact digits, so the separate exact-integer re-parse and
the parse_js_json flag it needed are gone.
* fix(images): send the card seed as a decimal string
JSON.parse rounds a number past 2**53, so a listing seed such as
9223372036854788153 reached the page as 9223372036854788000 while the
inspector showed the file's digits. Every non-null card seed is now its
decimal string.
The Node parity harness now JSON.parses the real listing response and
compares String(seed) with the seed the inspector shows, including a
seed past 2**53. TESTING.md describes the string contract and drops the
removed exact re-read.
* fix(images): bound digit-string seeds and isolate each card's parse
A seed recorded as a 5000-digit string reached int(), which raises
ValueError past Python's integer string limit, and the listing answered
500. A digit-string seed now has at most 20 digits (2**64 fits), and the
conversion falls back to no seed on ValueError.
Each card's parse is also guarded in the page builder: an unexpected
error in one file's records logs a warning and gives that card no seed
and no first-line title, and the rest of the page still lists.
* fix(images): filename search survives a failed prompt-link read
On main a workspace q matched only file paths, so search worked with a
broken database. Since search also matches linked prompts, any q turned
a failed link read into a 500. A q without category, tags or the rating
sort now falls back to filename matches with a null positive_count and
logs the failure with its traceback; those linked filters still fail.
* fix(images): card reads PNG records with the inspector's reader
The card read PNG text through Pillow's header info, where a later
duplicate prompt chunk silently wins and text after IDAT is missed. The
inspector's bounded reader refuses conflicting duplicates and reads text
after IDAT, so a file with two different prompt chunks showed a card
seed the inspector refused.
PNG cards now read their records with read_inspection_file, the reader
the inspector uses. When it refuses a file the card keeps its
dimensions from Pillow and shows no seed. Other formats still use
Pillow's header info, as the inspector's header read does.
* perf(images): read each card's header facts once per gallery scan
Every page parsed each card's prompt and workflow records again, about
3.4 ms per 50 KB record set and 25 ms per 400 KB workflow, inside
ComfyUI's server process, on every scroll-back and refresh.
Dimensions, the seed and the parsed positive's display title are now
cached per root by (path, mtime), beside the scan's per-file values and
invalidated the same way: a rescan keeps unchanged files and prunes the
rest, a roots change drops them all, and only files of the current scan
are stored. An unreadable header is not stored. Links and favorites
are still read on every page.
* ci: run the perf budget on main pushes and on demand only
The wall-clock budget has under 2x margin on shared ubuntu runners, and
a red exact-head check blocks a merge that must never be rerun to green.
The job now runs on each push to main and on workflow_dispatch, not on
pull requests. TESTING.md says a red run on main means fix the code.
* fix(library): cut display titles at 80 UTF-16 units, never inside a pair
The server counted code points while the Library's JavaScript counts UTF-16
units, so a first line of emoji kept 80 emoji on cards but showed 40 in the
Library. Both now cut at 80 UTF-16 units and drop a character that would
straddle the limit instead of leaving half a surrogate pair.
1599 lines
64 KiB
Python
1599 lines
64 KiB
Python
"""
|
|
Comprehensive database layer tests for PromptManager.
|
|
|
|
Tests all CRUD operations, tag junction tables, pagination,
|
|
search, statistics, image linking, and edge cases using
|
|
an in-memory SQLite database.
|
|
"""
|
|
|
|
import csv
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
from contextlib import closing
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from database.models import PromptModel # noqa: E402
|
|
|
|
try:
|
|
from tests.open_handles import assert_closed # noqa: E402
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from open_handles import assert_closed # noqa: E402
|
|
|
|
from database.operations import (
|
|
ORPHANED_PROMPTS_SQL,
|
|
PromptDatabase,
|
|
_resolve_db_path,
|
|
)
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
|
|
class DatabaseTestCase(unittest.TestCase):
|
|
"""Base class with temp database setup/teardown."""
|
|
|
|
def setUp(self):
|
|
self.temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
|
|
self.temp_db.close()
|
|
self.db = PromptDatabase(self.temp_db.name)
|
|
|
|
def tearDown(self):
|
|
PromptModel.close_all_instances()
|
|
assert_closed(self, self.temp_db.name)
|
|
for suffix in ("", "-wal", "-shm"):
|
|
path = self.temp_db.name + suffix
|
|
if os.path.exists(path):
|
|
os.unlink(path)
|
|
|
|
def _save(
|
|
self, text="Test prompt", category=None, tags=None, rating=None, notes=None
|
|
):
|
|
"""Helper to save a prompt and return its ID."""
|
|
return self.db.save_prompt(
|
|
text=text,
|
|
category=category,
|
|
tags=tags or [],
|
|
rating=rating,
|
|
notes=notes,
|
|
prompt_hash=generate_prompt_hash(text),
|
|
)
|
|
|
|
|
|
class TestPromptCRUD(DatabaseTestCase):
|
|
"""Test basic create, read, update, delete operations."""
|
|
|
|
def test_save_existing_prompt_returns_its_id_when_check_raced(self):
|
|
# Two callers check the hash, both see nothing, both insert: the second
|
|
# insert hits the UNIQUE(hash) constraint and must yield the first id.
|
|
text = "raced prompt"
|
|
prompt_hash = generate_prompt_hash(text)
|
|
original = self.db.get_prompt_by_hash
|
|
calls = []
|
|
|
|
def racing_lookup(h):
|
|
calls.append(h)
|
|
return None if len(calls) == 1 else original(h)
|
|
|
|
first_id = self._save(text)
|
|
with patch.object(self.db, "get_prompt_by_hash", racing_lookup):
|
|
self.assertIsNone(self.db.get_prompt_by_hash(prompt_hash))
|
|
second_id = self.db.save_prompt(text=text, prompt_hash=prompt_hash)
|
|
|
|
self.assertEqual(second_id, first_id)
|
|
self.assertEqual(len(self.db.search_prompts(text=text)), 1)
|
|
|
|
def test_save_without_hash_never_collides(self):
|
|
first = self.db.save_prompt(text="no hash a")
|
|
second = self.db.save_prompt(text="no hash b")
|
|
self.assertNotEqual(first, second)
|
|
|
|
def test_save_and_retrieve(self):
|
|
pid = self._save(
|
|
"A beautiful sunset", category="nature", tags=["sunset", "sky"], rating=5
|
|
)
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(prompt["text"], "A beautiful sunset")
|
|
self.assertEqual(prompt["category"], "nature")
|
|
self.assertIn("sunset", prompt["tags"])
|
|
self.assertIn("sky", prompt["tags"])
|
|
self.assertEqual(prompt["rating"], 5)
|
|
|
|
def test_save_minimal(self):
|
|
pid = self._save("Minimal prompt")
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(prompt["text"], "Minimal prompt")
|
|
self.assertIsNone(prompt["category"])
|
|
self.assertEqual(prompt["tags"], [])
|
|
self.assertIsNone(prompt["rating"])
|
|
|
|
def test_get_nonexistent_prompt(self):
|
|
result = self.db.get_prompt_by_id(99999)
|
|
self.assertIsNone(result)
|
|
|
|
def test_get_by_hash(self):
|
|
text = "Hash test prompt"
|
|
h = generate_prompt_hash(text)
|
|
self._save(text)
|
|
found = self.db.get_prompt_by_hash(h)
|
|
self.assertIsNotNone(found)
|
|
self.assertEqual(found["text"], text)
|
|
|
|
def test_get_by_hash_nonexistent(self):
|
|
result = self.db.get_prompt_by_hash("nonexistent_hash_value")
|
|
self.assertIsNone(result)
|
|
|
|
def test_update_metadata(self):
|
|
pid = self._save("Updatable prompt", category="old", tags=["old_tag"], rating=2)
|
|
self.db.update_prompt_metadata(
|
|
pid, category="new", tags=["new_tag"], rating=5, notes="updated"
|
|
)
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(prompt["category"], "new")
|
|
self.assertIn("new_tag", prompt["tags"])
|
|
self.assertNotIn("old_tag", prompt["tags"])
|
|
self.assertEqual(prompt["rating"], 5)
|
|
self.assertEqual(prompt["notes"], "updated")
|
|
|
|
def test_update_partial_metadata(self):
|
|
pid = self._save("Partial update", category="keep", tags=["keep_tag"], rating=3)
|
|
self.db.update_prompt_metadata(pid, rating=1)
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(prompt["rating"], 1)
|
|
# Category and tags should be unchanged
|
|
self.assertEqual(prompt["category"], "keep")
|
|
self.assertIn("keep_tag", prompt["tags"])
|
|
|
|
def test_delete_prompt(self):
|
|
pid = self._save("To be deleted")
|
|
result = self.db.delete_prompt(pid)
|
|
self.assertTrue(result)
|
|
self.assertIsNone(self.db.get_prompt_by_id(pid))
|
|
|
|
def test_delete_nonexistent(self):
|
|
result = self.db.delete_prompt(99999)
|
|
self.assertFalse(result)
|
|
|
|
|
|
class TestTagSuggestions(DatabaseTestCase):
|
|
"""suggest_tags(prefix, with_tags, limit): autocomplete narrowed by co-occurrence.
|
|
|
|
Context tags restrict suggestions to tags seen on the same prompts.
|
|
"""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
self._save("P1", tags=["asian", "portrait", "smile"])
|
|
self._save("P2", tags=["asian", "portrait"])
|
|
self._save("P3", tags=["portrait", "dog"])
|
|
self._save("P4", tags=["asian", "dog"])
|
|
self._save("P5", tags=["snow_day"])
|
|
|
|
def _names(self, *args, **kwargs):
|
|
return [(t["name"], t["count"]) for t in self.db.suggest_tags(*args, **kwargs)]
|
|
|
|
def test_no_context_lists_every_tag_by_frequency_then_name(self):
|
|
self.assertEqual(
|
|
self._names("", []),
|
|
[("asian", 3), ("portrait", 3), ("dog", 2), ("smile", 1), ("snow_day", 1)],
|
|
)
|
|
|
|
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 test_no_context_appends_unused_names_at_zero_after_used(self):
|
|
self._add_unused("zz_unused", "asian_unused")
|
|
used = [
|
|
("asian", 3),
|
|
("portrait", 3),
|
|
("dog", 2),
|
|
("smile", 1),
|
|
("snow_day", 1),
|
|
]
|
|
self.assertEqual(
|
|
self._names("", []),
|
|
used + [("asian_unused", 0), ("zz_unused", 0)],
|
|
)
|
|
self.assertEqual(self._names("asian", []), [("asian", 3), ("asian_unused", 0)])
|
|
|
|
def test_context_suggestions_never_include_unused_names(self):
|
|
self._add_unused("zz_unused", "asian_unused")
|
|
self.assertEqual(
|
|
self._names("", ["asian"]), [("portrait", 2), ("dog", 1), ("smile", 1)]
|
|
)
|
|
|
|
def test_limit_counts_unused_rows_last(self):
|
|
self._add_unused("zz_unused", "asian_unused")
|
|
self.assertEqual(
|
|
self._names("", [], limit=6),
|
|
[
|
|
("asian", 3),
|
|
("portrait", 3),
|
|
("dog", 2),
|
|
("smile", 1),
|
|
("snow_day", 1),
|
|
("asian_unused", 0),
|
|
],
|
|
)
|
|
|
|
def test_prefix_is_case_insensitive_and_anchored_at_the_start(self):
|
|
self.assertEqual(self._names("P", []), [("portrait", 3)])
|
|
self.assertEqual(self._names("s", []), [("smile", 1), ("snow_day", 1)])
|
|
self.assertEqual(self._names("ort", []), [])
|
|
|
|
def test_context_tags_narrow_to_co_occurring_tags_and_exclude_themselves(self):
|
|
self.assertEqual(
|
|
self._names("", ["asian"]), [("portrait", 2), ("dog", 1), ("smile", 1)]
|
|
)
|
|
self.assertEqual(self._names("", ["asian", "portrait"]), [("smile", 1)])
|
|
self.assertEqual(self._names("S", ["asian"]), [("smile", 1)])
|
|
self.assertEqual(self._names("", ["ASIAN", "dog"]), [])
|
|
|
|
def test_unknown_context_tag_yields_nothing(self):
|
|
self.assertEqual(self._names("", ["unicorn"]), [])
|
|
|
|
def test_limit_and_like_wildcards(self):
|
|
self.assertEqual(len(self._names("", [], limit=2)), 2)
|
|
self.assertEqual(self._names("%", []), [])
|
|
self.assertEqual(self._names("snow_", []), [("snow_day", 1)])
|
|
self.assertEqual(self._names("snow%", []), [])
|
|
|
|
def test_like_wildcards_never_match_unused_names(self):
|
|
self._add_unused("snowXday", "snowfall", "plain")
|
|
self.assertEqual(self._names("_", []), [])
|
|
self.assertEqual(self._names("%", []), [])
|
|
self.assertEqual(self._names("snow_", []), [("snow_day", 1)])
|
|
self.assertEqual(self._names("snow%", []), [])
|
|
self.assertEqual(
|
|
self._names("snow", []),
|
|
[("snow_day", 1), ("snowfall", 0), ("snowXday", 0)],
|
|
)
|
|
|
|
def test_context_limit_ignores_unused_rows(self):
|
|
self._add_unused("aaa_unused", "portrait_unused")
|
|
self.assertEqual(
|
|
self._names("", ["asian"], limit=2), [("portrait", 2), ("dog", 1)]
|
|
)
|
|
self.assertEqual(self._names("p", ["asian"], limit=1), [("portrait", 2)])
|
|
|
|
|
|
class TestTagJunctionTables(DatabaseTestCase):
|
|
"""Test normalized tag storage via junction tables."""
|
|
|
|
def test_tags_stored_in_junction_table(self):
|
|
pid = self._save("Tagged prompt", tags=["alpha", "beta"])
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(sorted(prompt["tags"]), ["alpha", "beta"])
|
|
|
|
def test_get_all_tags(self):
|
|
self._save("P1", tags=["a", "b"])
|
|
self._save("P2", tags=["b", "c"])
|
|
all_tags = self.db.get_all_tags()
|
|
self.assertEqual(sorted(all_tags), ["a", "b", "c"])
|
|
|
|
def test_get_all_tags_includes_unused_registry_rows(self):
|
|
pid = self._save("P1", tags=["a", "b"])
|
|
self.db.set_prompt_tags(pid, ["b"])
|
|
self.assertEqual(self.db.get_all_tags(), ["a", "b"])
|
|
|
|
def test_get_tags_with_counts(self):
|
|
self._save("P1", tags=["common", "rare"])
|
|
self._save("P2", tags=["common"])
|
|
self._save("P3", tags=["common", "other"])
|
|
result = self.db.get_tags_with_counts()
|
|
counts_dict = {t["name"]: t["count"] for t in result["tags"]}
|
|
self.assertEqual(counts_dict["common"], 3)
|
|
self.assertEqual(counts_dict["rare"], 1)
|
|
self.assertEqual(counts_dict["other"], 1)
|
|
|
|
def test_set_prompt_tags(self):
|
|
pid = self._save("Retaggable", tags=["old"])
|
|
self.db.set_prompt_tags(pid, ["new1", "new2"])
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(sorted(prompt["tags"]), ["new1", "new2"])
|
|
|
|
def test_set_empty_tags(self):
|
|
pid = self._save("Clear tags", tags=["remove_me"])
|
|
self.db.set_prompt_tags(pid, [])
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(prompt["tags"], [])
|
|
|
|
def test_rename_tag(self):
|
|
self._save("P1", tags=["old_name"])
|
|
self._save("P2", tags=["old_name", "other"])
|
|
self.db.rename_tag_all_prompts("old_name", "new_name")
|
|
all_tags = self.db.get_all_tags()
|
|
self.assertIn("new_name", all_tags)
|
|
self.assertNotIn("old_name", all_tags)
|
|
|
|
def test_delete_tag(self):
|
|
self._save("P1", tags=["keep", "remove"])
|
|
self._save("P2", tags=["remove"])
|
|
self.db.delete_tag_all_prompts("remove")
|
|
all_tags = self.db.get_all_tags()
|
|
self.assertIn("keep", all_tags)
|
|
self.assertNotIn("remove", all_tags)
|
|
|
|
def test_merge_tags(self):
|
|
self._save("P1", tags=["target"])
|
|
self._save("P2", tags=["source1"])
|
|
self._save("P3", tags=["source2", "target"])
|
|
self.db.merge_tags(["source1", "source2"], "target")
|
|
all_tags = self.db.get_all_tags()
|
|
self.assertIn("target", all_tags)
|
|
self.assertNotIn("source1", all_tags)
|
|
self.assertNotIn("source2", all_tags)
|
|
# All prompts should have the target tag
|
|
result = self.db.get_tags_with_counts()
|
|
counts = {t["name"]: t["count"] for t in result["tags"]}
|
|
self.assertEqual(counts["target"], 3)
|
|
|
|
def test_bulk_add_tags(self):
|
|
p1 = self._save("P1", tags=["existing"])
|
|
p2 = self._save("P2")
|
|
self.db.bulk_add_tags([p1, p2], ["bulk1", "bulk2"])
|
|
prompt1 = self.db.get_prompt_by_id(p1)
|
|
prompt2 = self.db.get_prompt_by_id(p2)
|
|
self.assertIn("bulk1", prompt1["tags"])
|
|
self.assertIn("bulk2", prompt1["tags"])
|
|
self.assertIn("existing", prompt1["tags"])
|
|
self.assertIn("bulk1", prompt2["tags"])
|
|
|
|
def test_untagged_prompts(self):
|
|
self._save("Tagged", tags=["has_tag"])
|
|
self._save("Untagged1")
|
|
self._save("Untagged2")
|
|
count = self.db.get_untagged_prompts_count()
|
|
self.assertEqual(count, 2)
|
|
result = self.db.get_untagged_prompts()
|
|
self.assertEqual(len(result["prompts"]), 2)
|
|
|
|
|
|
class TestSearch(DatabaseTestCase):
|
|
"""Test search and filter operations."""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
self._save(
|
|
"Beautiful mountain landscape",
|
|
category="nature",
|
|
tags=["mountain", "landscape"],
|
|
rating=5,
|
|
)
|
|
self._save(
|
|
"City skyline at night", category="urban", tags=["city", "night"], rating=4
|
|
)
|
|
self._save(
|
|
"Portrait of an artist",
|
|
category="portrait",
|
|
tags=["person", "art"],
|
|
rating=3,
|
|
)
|
|
self._save(
|
|
"Abstract geometric shapes",
|
|
category="abstract",
|
|
tags=["art", "geometric"],
|
|
rating=2,
|
|
)
|
|
|
|
def test_search_by_text(self):
|
|
results = self.db.search_prompts(text="mountain")
|
|
self.assertEqual(len(results), 1)
|
|
self.assertIn("mountain", results[0]["text"])
|
|
|
|
def test_search_by_text_case_insensitive(self):
|
|
results = self.db.search_prompts(text="MOUNTAIN")
|
|
self.assertEqual(len(results), 1)
|
|
|
|
def test_search_by_category(self):
|
|
results = self.db.search_prompts(category="urban")
|
|
self.assertEqual(len(results), 1)
|
|
self.assertEqual(results[0]["category"], "urban")
|
|
|
|
def test_search_by_tag(self):
|
|
results = self.db.search_prompts(tags=["art"])
|
|
self.assertEqual(len(results), 2)
|
|
|
|
def test_search_by_multiple_tags(self):
|
|
results = self.db.search_prompts(tags=["art", "geometric"])
|
|
self.assertEqual(len(results), 1)
|
|
self.assertIn("geometric", results[0]["text"].lower())
|
|
|
|
def test_search_by_rating_min(self):
|
|
results = self.db.search_prompts(rating_min=4)
|
|
self.assertEqual(len(results), 2)
|
|
|
|
def test_search_by_rating_range(self):
|
|
results = self.db.search_prompts(rating_min=3, rating_max=4)
|
|
self.assertEqual(len(results), 2)
|
|
|
|
def test_search_combined_filters(self):
|
|
results = self.db.search_prompts(text="landscape", rating_min=4)
|
|
self.assertEqual(len(results), 1)
|
|
|
|
def test_search_no_results(self):
|
|
results = self.db.search_prompts(text="nonexistent_query_xyz")
|
|
self.assertEqual(len(results), 0)
|
|
|
|
def test_search_with_limit(self):
|
|
results = self.db.search_prompts(limit=2)
|
|
self.assertEqual(len(results), 2)
|
|
|
|
|
|
class TestPagination(DatabaseTestCase):
|
|
"""Test pagination in get_recent_prompts."""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
for i in range(15):
|
|
self._save(f"Prompt number {i:02d}")
|
|
|
|
def test_first_page(self):
|
|
result = self.db.get_recent_prompts(limit=5, offset=0)
|
|
self.assertEqual(len(result["prompts"]), 5)
|
|
self.assertEqual(result["total"], 15)
|
|
self.assertTrue(result["has_more"])
|
|
self.assertEqual(result["page"], 1)
|
|
self.assertEqual(result["total_pages"], 3)
|
|
|
|
def test_middle_page(self):
|
|
result = self.db.get_recent_prompts(limit=5, offset=5)
|
|
self.assertEqual(len(result["prompts"]), 5)
|
|
self.assertTrue(result["has_more"])
|
|
self.assertEqual(result["page"], 2)
|
|
|
|
def test_last_page(self):
|
|
result = self.db.get_recent_prompts(limit=5, offset=10)
|
|
self.assertEqual(len(result["prompts"]), 5)
|
|
self.assertFalse(result["has_more"])
|
|
self.assertEqual(result["page"], 3)
|
|
|
|
def test_beyond_last_page(self):
|
|
result = self.db.get_recent_prompts(limit=5, offset=20)
|
|
self.assertEqual(len(result["prompts"]), 0)
|
|
self.assertFalse(result["has_more"])
|
|
|
|
def test_empty_database(self):
|
|
# Use a fresh empty DB
|
|
empty_db_file = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
|
|
empty_db_file.close()
|
|
empty_db = PromptDatabase(empty_db_file.name)
|
|
try:
|
|
result = empty_db.get_recent_prompts(limit=10, offset=0)
|
|
self.assertEqual(result["total"], 0)
|
|
self.assertEqual(len(result["prompts"]), 0)
|
|
self.assertFalse(result["has_more"])
|
|
finally:
|
|
empty_db.close_all()
|
|
for suffix in ("", "-wal", "-shm"):
|
|
if os.path.exists(empty_db_file.name + suffix):
|
|
os.unlink(empty_db_file.name + suffix)
|
|
|
|
def test_total_count_is_integer(self):
|
|
result = self.db.get_recent_prompts(limit=5)
|
|
self.assertIsInstance(result["total"], int)
|
|
self.assertIsInstance(result["has_more"], bool)
|
|
|
|
|
|
class TestStatistics(DatabaseTestCase):
|
|
"""Test get_statistics method."""
|
|
|
|
def test_empty_database_statistics(self):
|
|
stats = self.db.get_statistics()
|
|
self.assertEqual(stats["total_prompts"], 0)
|
|
self.assertEqual(stats["total_categories"], 0)
|
|
self.assertIsNone(stats.get("average_rating") or stats.get("avg_rating"))
|
|
self.assertEqual(stats["total_tags"], 0)
|
|
|
|
def test_populated_statistics(self):
|
|
self._save("P1", category="cat_a", tags=["t1", "t2"], rating=4)
|
|
self._save("P2", category="cat_b", tags=["t2", "t3"], rating=2)
|
|
self._save("P3", category="cat_a", tags=["t1"])
|
|
stats = self.db.get_statistics()
|
|
self.assertEqual(stats["total_prompts"], 3)
|
|
self.assertEqual(stats["total_categories"], 2)
|
|
self.assertEqual(stats["total_tags"], 3)
|
|
|
|
|
|
class TestCategories(DatabaseTestCase):
|
|
"""Test category operations."""
|
|
|
|
def test_get_prompts_by_category(self):
|
|
self._save("P1", category="nature")
|
|
self._save("P2", category="nature")
|
|
self._save("P3", category="urban")
|
|
results = self.db.get_prompts_by_category("nature")
|
|
self.assertEqual(len(results), 2)
|
|
|
|
def test_get_all_categories(self):
|
|
self._save("P1", category="nature")
|
|
self._save("P2", category="urban")
|
|
self._save("P3", category="nature")
|
|
categories = self.db.get_all_categories()
|
|
self.assertEqual(sorted(categories), ["nature", "urban"])
|
|
|
|
|
|
class TestTopRated(DatabaseTestCase):
|
|
"""Test top-rated prompt retrieval."""
|
|
|
|
def test_get_top_rated(self):
|
|
self._save("Low", rating=1)
|
|
self._save("High", rating=5)
|
|
self._save("Mid", rating=3)
|
|
self._save("Unrated")
|
|
results = self.db.get_top_rated_prompts(limit=2)
|
|
self.assertEqual(len(results), 2)
|
|
self.assertEqual(results[0]["rating"], 5)
|
|
self.assertEqual(results[1]["rating"], 3)
|
|
|
|
|
|
class TestDuplicateDetection(DatabaseTestCase):
|
|
"""Test duplicate handling."""
|
|
|
|
def test_same_hash_detected(self):
|
|
text = "Duplicate content"
|
|
h = generate_prompt_hash(text)
|
|
pid1 = self.db.save_prompt(text=text, prompt_hash=h)
|
|
existing = self.db.get_prompt_by_hash(h)
|
|
self.assertIsNotNone(existing)
|
|
self.assertEqual(existing["id"], pid1)
|
|
|
|
def test_cleanup_duplicates(self):
|
|
# Save multiple prompts first
|
|
self._save("Unique 1")
|
|
self._save("Unique 2")
|
|
removed = self.db.cleanup_duplicates()
|
|
self.assertEqual(removed, 0)
|
|
|
|
|
|
class TestImageOperations(DatabaseTestCase):
|
|
"""Test image linking and retrieval."""
|
|
|
|
def _link_image(self, prompt_id, path="/fake/path/image.png"):
|
|
return self.db.link_image_to_prompt(
|
|
prompt_id=str(prompt_id),
|
|
image_path=path,
|
|
)
|
|
|
|
def test_save_and_get_image(self):
|
|
pid = self._save("Prompt with image")
|
|
img_id = self._link_image(pid, "/fake/path/image.png")
|
|
self.assertIsNotNone(img_id)
|
|
images = self.db.get_prompt_images(str(pid))
|
|
self.assertEqual(len(images), 1)
|
|
self.assertEqual(images[0]["filename"], "image.png")
|
|
|
|
def test_image_count(self):
|
|
pid = self._save("Multi-image prompt")
|
|
for i in range(3):
|
|
self._link_image(pid, f"/fake/path/img{i}.png")
|
|
images = self.db.get_prompt_images(str(pid))
|
|
self.assertEqual(len(images), 3)
|
|
|
|
def test_delete_prompt_cascades_images(self):
|
|
pid = self._save("Cascade test")
|
|
self._link_image(pid, "/fake/path/img.png")
|
|
self.db.delete_prompt(pid)
|
|
images = self.db.get_prompt_images(str(pid))
|
|
self.assertEqual(len(images), 0)
|
|
|
|
|
|
class TestImageUniquenessByPath(DatabaseTestCase):
|
|
"""ComfyUI writes same-named files into different date folders."""
|
|
|
|
def _link(self, prompt_id, path):
|
|
return self.db.link_image_to_prompt(prompt_id=prompt_id, image_path=path)
|
|
|
|
def test_same_filename_in_different_folders_both_link(self):
|
|
pid = self._save("two folders")
|
|
first = self._link(pid, os.path.join("out", "2026-01-01", "a.png"))
|
|
second = self._link(pid, os.path.join("out", "2026-01-02", "a.png"))
|
|
|
|
self.assertTrue(first)
|
|
self.assertTrue(second)
|
|
self.assertNotEqual(first, second)
|
|
self.assertEqual(len(self.db.get_prompt_images(pid)), 2)
|
|
|
|
def test_same_path_spelled_differently_links_once(self):
|
|
pid = self._save("one file")
|
|
first = self._link(pid, os.path.join("out", "x", "..", "a.png"))
|
|
second = self._link(pid, os.path.join(".", "out", "a.png"))
|
|
|
|
self.assertTrue(first)
|
|
self.assertEqual(second, 0)
|
|
images = self.db.get_prompt_images(pid)
|
|
self.assertEqual(len(images), 1)
|
|
expected = os.path.normcase(
|
|
os.path.normpath(os.path.abspath(os.path.join("out", "a.png")))
|
|
)
|
|
self.assertEqual(images[0]["file_path"], expected)
|
|
|
|
def test_same_file_can_link_to_two_prompts(self):
|
|
a = self._save("prompt a")
|
|
b = self._save("prompt b")
|
|
path = os.path.join("out", "shared.png")
|
|
self.assertTrue(self._link(a, path))
|
|
self.assertTrue(self._link(b, path))
|
|
|
|
|
|
class TestLinkImageMetadataCoercion(DatabaseTestCase):
|
|
"""file_info from the request body is coerced; junk never reaches a column."""
|
|
|
|
XSS = '"><img src=x>'
|
|
|
|
def _link(self, file_info):
|
|
pid = self._save("coerce")
|
|
path = f"/out/c{len(self.db.get_prompt_images(pid))}.png"
|
|
image_id = self.db.link_image_to_prompt(pid, path, {"file_info": file_info})
|
|
self.assertGreater(image_id, 0, "image must still link")
|
|
return self.db.get_image_by_id(image_id)
|
|
|
|
def test_string_payloads_are_stored_as_null(self):
|
|
row = self._link({"size": self.XSS, "dimensions": self.XSS, "format": self.XSS})
|
|
|
|
self.assertIsNone(row["file_size"])
|
|
self.assertIsNone(row["width"])
|
|
self.assertIsNone(row["height"])
|
|
self.assertIsNone(row["format"])
|
|
|
|
def test_numbers_and_known_formats_are_kept(self):
|
|
row = self._link({"size": 1234, "dimensions": [640, 480], "format": "png"})
|
|
|
|
self.assertEqual(row["file_size"], 1234)
|
|
self.assertEqual(row["width"], 640)
|
|
self.assertEqual(row["height"], 480)
|
|
self.assertEqual(row["format"], "PNG")
|
|
|
|
def test_numeric_strings_are_accepted(self):
|
|
row = self._link({"size": "99", "dimensions": ["8", "16"], "format": "webp"})
|
|
|
|
self.assertEqual((row["file_size"], row["width"], row["height"]), (99, 8, 16))
|
|
self.assertEqual(row["format"], "WEBP")
|
|
|
|
def test_out_of_range_values_are_dropped(self):
|
|
too_big = self._link({"size": 10**12 + 1, "dimensions": [65536, 65535]})
|
|
negative = self._link({"size": -1, "dimensions": [-1, 0]})
|
|
|
|
self.assertIsNone(too_big["file_size"])
|
|
self.assertIsNone(too_big["width"])
|
|
self.assertEqual(too_big["height"], 65535)
|
|
self.assertIsNone(negative["file_size"])
|
|
self.assertIsNone(negative["width"])
|
|
self.assertEqual(negative["height"], 0)
|
|
|
|
def test_unknown_format_and_short_dimensions_are_dropped(self):
|
|
row = self._link({"dimensions": [640], "format": "svg"})
|
|
|
|
self.assertIsNone(row["width"])
|
|
self.assertIsNone(row["height"])
|
|
self.assertIsNone(row["format"])
|
|
|
|
def test_non_dict_metadata_still_links(self):
|
|
pid = self._save("junk metadata")
|
|
|
|
first = self.db.link_image_to_prompt(pid, "/out/a.png", "not a dict")
|
|
second = self.db.link_image_to_prompt(
|
|
pid, "/out/b.png", {"file_info": ["not", "a", "dict"]}
|
|
)
|
|
|
|
self.assertGreater(first, 0)
|
|
self.assertGreater(second, 0)
|
|
for row in self.db.get_prompt_images(pid):
|
|
self.assertIsNone(row["width"])
|
|
self.assertIsNone(row["format"])
|
|
|
|
|
|
class TestSearchEscapesLikeWildcards(DatabaseTestCase):
|
|
def test_percent_is_literal(self):
|
|
self._save("100% sure")
|
|
self._save("100 percent sure")
|
|
texts = [p["text"] for p in self.db.search_prompts(text="100%")]
|
|
self.assertEqual(texts, ["100% sure"])
|
|
|
|
def test_underscore_is_literal(self):
|
|
self._save("snake_case name")
|
|
self._save("snakeXcase name")
|
|
texts = [p["text"] for p in self.db.search_prompts(text="snake_case")]
|
|
self.assertEqual(texts, ["snake_case name"])
|
|
|
|
def test_backslash_is_literal(self):
|
|
self._save("path C:\\out\\img")
|
|
self._save("path C:out img")
|
|
texts = [p["text"] for p in self.db.search_prompts(text="C:\\out")]
|
|
self.assertEqual(texts, ["path C:\\out\\img"])
|
|
|
|
def test_partial_tag_filter_is_literal(self):
|
|
self._save("tagged a", tags=["50%_off"])
|
|
self._save("tagged b", tags=["50x_off"])
|
|
texts = [
|
|
p["text"] for p in self.db.search_prompts(tags=["50%"], tag_partial=True)
|
|
]
|
|
self.assertEqual(texts, ["tagged a"])
|
|
|
|
|
|
class TestEdgeCases(DatabaseTestCase):
|
|
"""Test edge cases and boundary conditions."""
|
|
|
|
def test_special_characters_in_text(self):
|
|
pid = self._save("Prompt with 'quotes' and \"double quotes\" and <html>")
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertIn("quotes", prompt["text"])
|
|
|
|
def test_unicode_text(self):
|
|
pid = self._save("日本語テスト prompt with émojis 🎨")
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertIn("日本語", prompt["text"])
|
|
|
|
def test_very_long_text(self):
|
|
long_text = "word " * 1000
|
|
pid = self._save(long_text.strip())
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(prompt["text"], long_text.strip())
|
|
|
|
def test_tag_with_special_characters(self):
|
|
pid = self._save(
|
|
"Special tags",
|
|
tags=["tag-with-dash", "tag_with_underscore", "tag.with.dots"],
|
|
)
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual(len(prompt["tags"]), 3)
|
|
|
|
def test_empty_category_string(self):
|
|
pid = self._save("Empty cat", category="")
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
# Empty string category is stored as-is
|
|
self.assertIn(prompt["category"], ["", None])
|
|
|
|
def test_rating_boundary_values(self):
|
|
p1 = self._save("Rating 1", rating=1)
|
|
p5 = self._save("Rating 5", rating=5)
|
|
self.assertEqual(self.db.get_prompt_by_id(p1)["rating"], 1)
|
|
self.assertEqual(self.db.get_prompt_by_id(p5)["rating"], 5)
|
|
|
|
|
|
class TestPreviewImages(DatabaseTestCase):
|
|
"""Test _attach_preview_images functionality."""
|
|
|
|
def test_preview_images_attached(self):
|
|
pid = self._save("Preview test")
|
|
for i in range(5):
|
|
self.db.link_image_to_prompt(
|
|
prompt_id=str(pid),
|
|
image_path=f"/fake/img{i}.png",
|
|
)
|
|
result = self.db.get_recent_prompts(limit=10)
|
|
prompt = result["prompts"][0]
|
|
# Should have preview images (max 3) and total count
|
|
self.assertIn("images", prompt)
|
|
self.assertLessEqual(len(prompt["images"]), 3)
|
|
self.assertEqual(prompt["image_count"], 5)
|
|
|
|
|
|
class TestExport(DatabaseTestCase):
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.out_dir = tempfile.mkdtemp()
|
|
self.addCleanup(self._remove_out_dir)
|
|
|
|
def _remove_out_dir(self):
|
|
for name in os.listdir(self.out_dir):
|
|
os.unlink(os.path.join(self.out_dir, name))
|
|
os.rmdir(self.out_dir)
|
|
|
|
def test_json_export_writes_every_prompt_with_tags(self):
|
|
self._save("first", tags=["a", "b"], category="cat")
|
|
self._save("second")
|
|
path = os.path.join(self.out_dir, "prompts.json")
|
|
|
|
self.assertTrue(self.db.export_prompts(path, "json"))
|
|
|
|
with open(path, encoding="utf-8") as fh:
|
|
data = json.load(fh)
|
|
self.assertEqual(sorted(p["text"] for p in data), ["first", "second"])
|
|
self.assertEqual(
|
|
next(p for p in data if p["text"] == "first")["tags"], ["a", "b"]
|
|
)
|
|
|
|
def test_csv_export_joins_tags(self):
|
|
self._save("first", tags=["a", "b"])
|
|
path = os.path.join(self.out_dir, "prompts.csv")
|
|
|
|
self.assertTrue(self.db.export_prompts(path, "CSV"))
|
|
|
|
with open(path, newline="", encoding="utf-8") as fh:
|
|
rows = list(csv.DictReader(fh))
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertEqual(rows[0]["text"], "first")
|
|
self.assertEqual(rows[0]["tags"], "a, b")
|
|
|
|
def test_csv_export_of_empty_database_writes_nothing(self):
|
|
path = os.path.join(self.out_dir, "empty.csv")
|
|
self.assertTrue(self.db.export_prompts(path, "csv"))
|
|
self.assertFalse(os.path.exists(path))
|
|
|
|
def test_unsupported_format_returns_false(self):
|
|
self._save("x")
|
|
path = os.path.join(self.out_dir, "prompts.xml")
|
|
self.assertFalse(self.db.export_prompts(path, "xml"))
|
|
self.assertFalse(os.path.exists(path))
|
|
|
|
def test_unwritable_destination_returns_false(self):
|
|
self._save("x")
|
|
path = os.path.join(self.out_dir, "missing", "dir", "prompts.json")
|
|
self.assertFalse(self.db.export_prompts(path, "json"))
|
|
|
|
|
|
class TestDuplicateMerging(DatabaseTestCase):
|
|
def _save_dup(self, text, created_at, **kwargs):
|
|
pid = self.db.save_prompt(text=text, prompt_hash=None, **kwargs)
|
|
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
|
|
conn.execute(
|
|
"UPDATE prompts SET created_at = ? WHERE id = ?", (created_at, pid)
|
|
)
|
|
return pid
|
|
|
|
def test_cleanup_keeps_oldest_merges_metadata_and_moves_images(self):
|
|
keep = self._save_dup(
|
|
"dup", "2026-01-01T00:00:00", tags=["a"], notes="first", rating=2
|
|
)
|
|
drop = self._save_dup(
|
|
"DUP",
|
|
"2026-02-01T00:00:00",
|
|
tags=["b"],
|
|
category="cat",
|
|
notes="second",
|
|
rating=4,
|
|
)
|
|
self.db.link_image_to_prompt(drop, "/out/moved.png")
|
|
self.db.link_image_to_prompt(keep, "/out/kept.png")
|
|
|
|
removed = self.db.cleanup_duplicates()
|
|
|
|
self.assertEqual(removed, 1)
|
|
self.assertIsNone(self.db.get_prompt_by_id(drop))
|
|
merged = self.db.get_prompt_by_id(keep)
|
|
self.assertEqual(merged["category"], "cat")
|
|
self.assertEqual(merged["rating"], 4)
|
|
self.assertEqual(merged["notes"], "first | second")
|
|
self.assertEqual(sorted(merged["tags"]), ["a", "b"])
|
|
self.assertEqual(len(self.db.get_prompt_images(keep)), 2)
|
|
|
|
def test_cleanup_takes_notes_from_duplicate_when_primary_has_none(self):
|
|
keep = self._save_dup("note dup", "2026-01-01T00:00:00")
|
|
self._save_dup("note dup", "2026-02-01T00:00:00", notes="only here")
|
|
self.db.cleanup_duplicates()
|
|
self.assertEqual(self.db.get_prompt_by_id(keep)["notes"], "only here")
|
|
|
|
def test_duplicate_scans_return_empty_on_database_error(self):
|
|
with patch.object(
|
|
self.db.model, "get_connection", side_effect=sqlite3.OperationalError
|
|
):
|
|
self.assertEqual(self.db.cleanup_duplicates(), 0)
|
|
|
|
def test_hash_duplicates_are_impossible_under_unique_hash(self):
|
|
self._save("one")
|
|
self.assertEqual(self.db.check_hash_duplicates(), [])
|
|
|
|
def test_merge_helpers_survive_database_errors(self):
|
|
pid = self._save("helper")
|
|
broken = MagicMock()
|
|
broken.execute.side_effect = sqlite3.OperationalError("disk I/O error")
|
|
self.assertEqual(self.db._merge_duplicate_metadata(broken, pid, []), {})
|
|
self.assertEqual(self.db._transfer_images_to_primary(broken, pid, [pid]), 0)
|
|
self.db._update_primary_with_merged_metadata(broken, pid, {"tags": []})
|
|
self.assertEqual(self.db.get_prompt_by_id(pid)["text"], "helper")
|
|
|
|
|
|
class TestImageQueries(DatabaseTestCase):
|
|
METADATA = {
|
|
"file_info": {"size": 123, "dimensions": [640, 480], "format": "PNG"},
|
|
"workflow": {"nodes": [1, 2]},
|
|
"prompt": {"3": {"inputs": {"cfg": float("nan")}}},
|
|
"parameters": {"steps": 20},
|
|
}
|
|
|
|
def test_link_stores_metadata_and_cleans_nan(self):
|
|
pid = self._save("with metadata")
|
|
image_id = self.db.link_image_to_prompt(pid, "/out/meta.png", self.METADATA)
|
|
|
|
image = self.db.get_image_by_id(image_id)
|
|
self.assertEqual(image["file_size"], 123)
|
|
self.assertEqual((image["width"], image["height"]), (640, 480))
|
|
self.assertEqual(image["format"], "PNG")
|
|
self.assertEqual(image["workflow_data"], {"nodes": [1, 2]})
|
|
self.assertIsNone(image["prompt_metadata"]["3"]["inputs"]["cfg"])
|
|
self.assertEqual(image["parameters"], {"steps": 20})
|
|
|
|
def test_link_rejects_temporary_invalid_and_unknown_prompt_ids(self):
|
|
self.assertEqual(self.db.link_image_to_prompt("temp_123", "/out/a.png"), 0)
|
|
self.assertEqual(self.db.link_image_to_prompt("not-an-id", "/out/a.png"), 0)
|
|
self.assertEqual(self.db.link_image_to_prompt(99999, "/out/a.png"), 0)
|
|
self.assertEqual(self.db.link_image_to_prompt(None, "/out/a.png"), 0)
|
|
|
|
def test_link_returns_zero_on_database_error(self):
|
|
pid = self._save("err")
|
|
with patch.object(
|
|
self.db.model, "get_connection", side_effect=sqlite3.OperationalError
|
|
):
|
|
self.assertEqual(self.db.link_image_to_prompt(pid, "/out/a.png"), 0)
|
|
|
|
def test_recent_all_and_search_images_carry_prompt_text(self):
|
|
a = self._save("alpha prompt", tags=["x"])
|
|
b = self._save("beta prompt")
|
|
self.db.link_image_to_prompt(a, "/out/a1.png")
|
|
self.db.link_image_to_prompt(b, "/out/b1.png")
|
|
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
|
|
conn.execute(
|
|
"UPDATE generated_images SET generation_time = '2026-01-01 00:00:00'"
|
|
" WHERE filename = 'a1.png'"
|
|
)
|
|
conn.execute(
|
|
"UPDATE generated_images SET generation_time = '2026-02-01 00:00:00'"
|
|
" WHERE filename = 'b1.png'"
|
|
)
|
|
|
|
recent = self.db.get_recent_images(limit=10)
|
|
self.assertEqual(
|
|
[i["prompt_text"] for i in recent], ["beta prompt", "alpha prompt"]
|
|
)
|
|
|
|
everything = self.db.get_all_images()
|
|
self.assertEqual(len(everything), 2)
|
|
by_text = {i["prompt_text"]: i for i in everything}
|
|
self.assertEqual(by_text["alpha prompt"]["prompt_tags"], ["x"])
|
|
self.assertEqual(by_text["beta prompt"]["prompt_tags"], [])
|
|
|
|
page = self.db.get_all_images(limit=1, offset=1)
|
|
self.assertEqual([i["prompt_text"] for i in page], ["alpha prompt"])
|
|
|
|
found = self.db.search_images_by_prompt("alpha")
|
|
self.assertEqual([i["filename"] for i in found], ["a1.png"])
|
|
|
|
def test_image_search_treats_like_wildcards_literally(self):
|
|
literal = self._save("50% off")
|
|
other = self._save("50 percent off")
|
|
self.db.link_image_to_prompt(literal, "/out/literal.png")
|
|
self.db.link_image_to_prompt(other, "/out/other.png")
|
|
found = self.db.search_images_by_prompt("50%")
|
|
self.assertEqual([i["filename"] for i in found], ["literal.png"])
|
|
|
|
def test_delete_image(self):
|
|
pid = self._save("del")
|
|
image_id = self.db.link_image_to_prompt(pid, "/out/del.png")
|
|
self.assertTrue(self.db.delete_image(image_id))
|
|
self.assertFalse(self.db.delete_image(image_id))
|
|
self.assertIsNone(self.db.get_image_by_id(image_id))
|
|
|
|
def test_cleanup_missing_images_keeps_files_that_exist(self):
|
|
pid = self._save("files")
|
|
fd, existing = tempfile.mkstemp(suffix=".png")
|
|
os.close(fd)
|
|
self.addCleanup(os.unlink, existing)
|
|
self.db.link_image_to_prompt(pid, existing)
|
|
self.db.link_image_to_prompt(
|
|
pid, os.path.join(tempfile.gettempdir(), "gone.png")
|
|
)
|
|
|
|
self.assertEqual(self.db.cleanup_missing_images(), 1)
|
|
images = self.db.get_prompt_images(pid)
|
|
self.assertEqual([i["image_path"] for i in images], [existing])
|
|
|
|
def test_image_prompt_info_by_exact_normalised_or_basename_path(self):
|
|
pid = self._save("lookup", tags=["t"], category="c", rating=3, notes="n")
|
|
stored = os.path.join("out", "sub", "look.png")
|
|
self.db.link_image_to_prompt(pid, stored, self.METADATA)
|
|
|
|
exact = self.db.get_image_prompt_info(stored)
|
|
self.assertEqual(exact["prompt_id"], pid)
|
|
self.assertEqual(exact["tags"], ["t"])
|
|
self.assertEqual(exact["workflow_data"], {"nodes": [1, 2]})
|
|
|
|
normalised = self.db.get_image_prompt_info(
|
|
os.path.join("out", "x", "..", "sub", "look.png")
|
|
)
|
|
self.assertEqual(normalised["prompt_id"], pid)
|
|
basename = self.db.get_image_prompt_info(os.path.join("elsewhere", "look.png"))
|
|
self.assertEqual(basename["prompt_id"], pid)
|
|
self.assertIsNone(self.db.get_image_prompt_info("nothing.png"))
|
|
|
|
self.assertEqual(self.db.get_prompt_id_for_image(stored), pid)
|
|
self.assertEqual(
|
|
self.db.get_prompt_id_for_image(
|
|
os.path.join(".", "out", "sub", "look.png")
|
|
),
|
|
pid,
|
|
)
|
|
self.assertIsNone(self.db.get_prompt_id_for_image("nothing.png"))
|
|
|
|
def test_image_row_with_broken_json_yields_empty_dicts(self):
|
|
pid = self._save("broken")
|
|
image_id = self.db.link_image_to_prompt(pid, "/out/broken.png")
|
|
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
|
|
conn.execute(
|
|
"UPDATE generated_images SET workflow_data = '{not json', "
|
|
"prompt_metadata = '{\"k\": 1}', parameters = NULL WHERE id = ?",
|
|
(image_id,),
|
|
)
|
|
image = self.db.get_image_by_id(image_id)
|
|
self.assertEqual(image["workflow_data"], {})
|
|
self.assertEqual(image["prompt_metadata"], {"k": 1})
|
|
self.assertEqual(image["parameters"], {})
|
|
info = self.db.get_image_prompt_info("/out/broken.png")
|
|
self.assertIsNone(info["workflow_data"])
|
|
self.assertEqual(info["prompt_metadata"], {"k": 1})
|
|
|
|
def test_image_prompt_info_reports_ambiguity_only_when_requested(self):
|
|
stored = os.path.join("out", "exact.png")
|
|
self.assertIsNone(self.db.get_image_prompt_info(stored, False))
|
|
self.assertIsNone(
|
|
self.db.get_image_prompt_info(stored, False, report_ambiguity=True)
|
|
)
|
|
|
|
first = self._save("First owner", category="First", tags=["first"])
|
|
self.db.link_image_to_prompt(first, stored, self.METADATA)
|
|
original = self.db.get_image_prompt_info(stored, False)
|
|
self.assertEqual(original["prompt_id"], first)
|
|
self.assertEqual(original["category"], "First")
|
|
self.assertEqual(original["tags"], ["first"])
|
|
self.assertEqual(
|
|
self.db.get_image_prompt_info(stored, False, report_ambiguity=True),
|
|
original,
|
|
)
|
|
|
|
second = self._save("Second owner", category="Second", tags=["second"])
|
|
self.db.link_image_to_prompt(second, stored, self.METADATA)
|
|
self.assertIsNone(self.db.get_image_prompt_info(stored))
|
|
self.assertIsNone(self.db.get_image_prompt_info(stored, False))
|
|
statements = []
|
|
with self.db.model.get_connection() as conn:
|
|
conn.set_trace_callback(statements.append)
|
|
try:
|
|
result = self.db.get_image_prompt_info(
|
|
stored, False, report_ambiguity=True
|
|
)
|
|
finally:
|
|
conn.set_trace_callback(None)
|
|
self.assertEqual(result, {"association_status": "ambiguous"})
|
|
self.assertEqual(
|
|
len(
|
|
[sql for sql in statements if sql.lstrip().upper().startswith("SELECT")]
|
|
),
|
|
1,
|
|
)
|
|
self.assertIsNone(
|
|
self.db.get_image_prompt_info("missing.png", False, report_ambiguity=True)
|
|
)
|
|
|
|
def test_image_prompt_info_ambiguity_opt_in_keeps_exact_path_precedence(self):
|
|
first = self._save("Exact owner")
|
|
second = self._save("Same basename")
|
|
exact = os.path.join("out", "one", "same.png")
|
|
other = os.path.join("out", "two", "same.png")
|
|
self.db.link_image_to_prompt(first, exact)
|
|
self.db.link_image_to_prompt(second, other)
|
|
original = self.db.get_image_prompt_info(exact)
|
|
self.assertEqual(original["prompt_id"], first)
|
|
self.assertEqual(
|
|
self.db.get_image_prompt_info(exact, report_ambiguity=True), original
|
|
)
|
|
absent = os.path.join("elsewhere", "same.png")
|
|
self.assertIsNone(self.db.get_image_prompt_info(absent))
|
|
self.assertIsNone(
|
|
self.db.get_image_prompt_info(absent, False, report_ambiguity=True)
|
|
)
|
|
self.assertEqual(
|
|
self.db.get_image_prompt_info(absent, report_ambiguity=True),
|
|
{"association_status": "ambiguous"},
|
|
)
|
|
|
|
|
|
class TestMaintenanceOperations(DatabaseTestCase):
|
|
def test_statistics_count_everything(self):
|
|
a = self._save("a", category="cat", tags=["t1", "t2"], rating=4)
|
|
self._save("b", category=" cat ", rating=2)
|
|
self._save("c")
|
|
self.db.link_image_to_prompt(a, "/out/a.png")
|
|
self.db.link_image_to_prompt(a, "/out/b.png")
|
|
|
|
stats = self.db.get_statistics()
|
|
|
|
self.assertEqual(stats["total_prompts"], 3)
|
|
self.assertEqual(stats["total_categories"], 1)
|
|
self.assertEqual(stats["average_rating"], 3.0)
|
|
self.assertEqual(stats["total_tags"], 2)
|
|
self.assertEqual(stats["total_images"], 2)
|
|
self.assertEqual(stats["images_with_prompts"], 1)
|
|
|
|
def test_update_text_and_rating(self):
|
|
pid = self._save("old text", rating=1)
|
|
self.assertTrue(self.db.update_prompt_text(pid, "new text"))
|
|
self.assertTrue(self.db.update_prompt_rating(pid, 5))
|
|
prompt = self.db.get_prompt_by_id(pid)
|
|
self.assertEqual((prompt["text"], prompt["rating"]), ("new text", 5))
|
|
self.assertFalse(self.db.update_prompt_text(99999, "x"))
|
|
self.assertFalse(self.db.update_prompt_rating(99999, 3))
|
|
|
|
def test_update_metadata_validation(self):
|
|
pid = self._save("meta")
|
|
self.assertFalse(self.db.update_prompt_metadata(pid))
|
|
with self.assertRaises(ValueError):
|
|
self.db.update_prompt_metadata(pid, rating=7)
|
|
self.assertTrue(self.db.update_prompt_metadata(pid, tags=["only tags"]))
|
|
self.assertEqual(self.db.get_prompt_by_id(pid)["tags"], ["only tags"])
|
|
|
|
def test_save_prompt_validation(self):
|
|
with self.assertRaises(ValueError):
|
|
self.db.save_prompt(text=" ")
|
|
with self.assertRaises(ValueError):
|
|
self.db.save_prompt(text="rated", rating=0)
|
|
|
|
def test_bulk_operations(self):
|
|
a = self._save("a", tags=["keep"])
|
|
b = self._save("b")
|
|
c = self._save("c")
|
|
self.db.link_image_to_prompt(c, "/out/c.png")
|
|
|
|
self.assertEqual(self.db.bulk_add_tags([a, b, 99999], ["keep", "new"]), 2)
|
|
self.assertEqual(self.db.bulk_add_tags([a], ["keep"]), 0)
|
|
self.assertEqual(sorted(self.db.get_prompt_by_id(a)["tags"]), ["keep", "new"])
|
|
|
|
self.assertEqual(self.db.bulk_set_category([a, b, 99999], "bulk"), 2)
|
|
self.assertEqual(self.db.get_prompt_by_id(b)["category"], "bulk")
|
|
|
|
self.assertEqual(self.db.bulk_delete_prompts([b, c, 99999]), 2)
|
|
self.assertIsNone(self.db.get_prompt_by_id(c))
|
|
self.assertEqual(self.db.get_prompt_images(c), [])
|
|
|
|
def test_prune_orphaned_prompts_spares_protected_and_illustrated(self):
|
|
orphan = self._save("orphan")
|
|
protected = self._save("protected", tags=["__protected__"])
|
|
illustrated = self._save("illustrated")
|
|
self.db.link_image_to_prompt(illustrated, "/out/i.png")
|
|
|
|
self.assertEqual(self.db.prune_orphaned_prompts(), 1)
|
|
self.assertIsNone(self.db.get_prompt_by_id(orphan))
|
|
self.assertIsNotNone(self.db.get_prompt_by_id(protected))
|
|
self.assertIsNotNone(self.db.get_prompt_by_id(illustrated))
|
|
self.assertEqual(self.db.prune_orphaned_prompts(), 0)
|
|
|
|
def test_prune_orphaned_prompts_spares_library_parents(self):
|
|
titled = self._save("titled parent")
|
|
untitled = self._save("untitled parent")
|
|
orphan = self._save("orphan")
|
|
with self.db.model.get_connection() as conn:
|
|
conn.execute(
|
|
"INSERT INTO prompt_library_state(prompt_id,title) "
|
|
"VALUES(?,?),(?,NULL)",
|
|
(titled, "My recipe", untitled),
|
|
)
|
|
conn.commit()
|
|
|
|
self.assertEqual(self.db.prune_orphaned_prompts(), 1)
|
|
self.assertIsNone(self.db.get_prompt_by_id(orphan))
|
|
self.assertIsNotNone(self.db.get_prompt_by_id(titled))
|
|
self.assertIsNotNone(self.db.get_prompt_by_id(untitled))
|
|
with self.db.model.get_connection() as conn:
|
|
titles = conn.execute(
|
|
"SELECT prompt_id,title FROM prompt_library_state ORDER BY prompt_id"
|
|
).fetchall()
|
|
self.assertEqual(
|
|
[tuple(row) for row in titles], [(titled, "My recipe"), (untitled, None)]
|
|
)
|
|
self.assertEqual(self.db.prune_orphaned_prompts(), 0)
|
|
|
|
def test_prune_query_has_no_scan_inside_correlated_subquery(self):
|
|
conn = self.db.model.get_connection()
|
|
plan = conn.execute("EXPLAIN QUERY PLAN " + ORPHANED_PROMPTS_SQL).fetchall()
|
|
parents = {row[0]: row[1] for row in plan}
|
|
details = {row[0]: row[3] for row in plan}
|
|
|
|
def correlated(node):
|
|
while node in parents:
|
|
node = parents[node]
|
|
if details.get(node, "").startswith("CORRELATED"):
|
|
return True
|
|
return False
|
|
|
|
scans = [
|
|
row[3] for row in plan if row[3].startswith("SCAN ") and correlated(row[0])
|
|
]
|
|
self.assertEqual(scans, [])
|
|
self.assertIn("execution_run_contexts", ORPHANED_PROMPTS_SQL)
|
|
|
|
def test_consistency_check_reports_orphaned_rows(self):
|
|
self.assertEqual(self.db.check_consistency(), [])
|
|
pid = self._save("to orphan", tags=["t"])
|
|
self.db.link_image_to_prompt(pid, "/out/o.png")
|
|
self.db.close()
|
|
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
|
|
conn.execute("PRAGMA foreign_keys = OFF")
|
|
conn.execute("DELETE FROM prompts WHERE id = ?", (pid,))
|
|
|
|
issues = self.db.check_consistency()
|
|
|
|
self.assertEqual(len(issues), 2)
|
|
self.assertTrue(any("prompt_tags" in i for i in issues))
|
|
self.assertTrue(any(i.startswith("Image ") for i in issues))
|
|
|
|
def test_subfolders_relative_to_roots_with_ancestors(self):
|
|
pid = self._save("folders")
|
|
root = os.path.join(os.sep, "gallery")
|
|
self.db.link_image_to_prompt(
|
|
pid, os.path.join(root, "2026", "08-Aug", "2026-08-06", "a.png")
|
|
)
|
|
self.db.link_image_to_prompt(pid, os.path.join(root, "b.png"))
|
|
self.db.link_image_to_prompt(
|
|
pid, os.path.join(os.sep, "elsewhere", "deep", "c.png")
|
|
)
|
|
self.db.link_image_to_prompt(pid, "loose.png")
|
|
|
|
flat = self.db.get_prompt_subfolders(root_dirs=[root])
|
|
self.assertIn("2026/08-Aug/2026-08-06", flat)
|
|
self.assertIn(
|
|
os.path.join(os.sep, "elsewhere", "deep").replace(os.sep, "/"), flat
|
|
)
|
|
self.assertNotIn(".", flat)
|
|
|
|
nested = self.db.get_prompt_subfolders(root_dirs=[root], include_ancestors=True)
|
|
self.assertIn("2026", nested)
|
|
self.assertIn("2026/08-Aug", nested)
|
|
|
|
without_roots = self.db.get_prompt_subfolders()
|
|
self.assertIn(
|
|
os.path.join(root, "2026", "08-Aug", "2026-08-06").replace(os.sep, "/"),
|
|
without_roots,
|
|
)
|
|
|
|
def test_prompts_by_tags_and_or_modes(self):
|
|
both = self._save("both", tags=["x", "y"])
|
|
only_x = self._save("only x", tags=["x"])
|
|
self._save("none")
|
|
|
|
self.assertEqual(self.db.get_prompts_by_tags([])["total"], 0)
|
|
either = self.db.get_prompts_by_tags(["x", "y"], mode="or")
|
|
self.assertEqual(sorted(p["id"] for p in either["prompts"]), [both, only_x])
|
|
self.assertEqual(either["total"], 2)
|
|
all_of = self.db.get_prompts_by_tags(["x", "y"], mode="and", limit=1)
|
|
self.assertEqual([p["id"] for p in all_of["prompts"]], [both])
|
|
self.assertFalse(all_of["has_more"])
|
|
|
|
def test_tag_counts_with_search_and_sorts(self):
|
|
self._save("1", tags=["beta", "alpha"])
|
|
self._save("2", tags=["beta"])
|
|
|
|
by_count = self.db.get_tags_with_counts(sort="count_desc")
|
|
self.assertEqual([t["name"] for t in by_count["tags"]], ["beta", "alpha"])
|
|
ascending = self.db.get_tags_with_counts(sort="count_asc")
|
|
self.assertEqual([t["name"] for t in ascending["tags"]], ["alpha", "beta"])
|
|
reverse = self.db.get_tags_with_counts(sort="alpha_desc")
|
|
self.assertEqual([t["name"] for t in reverse["tags"]], ["beta", "alpha"])
|
|
searched = self.db.get_tags_with_counts(search="alp")
|
|
self.assertEqual(searched["total"], 1)
|
|
self.assertEqual(searched["tags"][0]["count"], 1)
|
|
|
|
def test_tag_management_edge_cases(self):
|
|
a = self._save("a", tags=["old", "target"])
|
|
b = self._save("b", tags=["old"])
|
|
|
|
merged = self.db.rename_tag_all_prompts("old", "target")
|
|
self.assertEqual(merged["affected_count"], 2)
|
|
self.assertEqual(self.db.get_prompt_by_id(a)["tags"], ["target"])
|
|
self.assertEqual(self.db.get_prompt_by_id(b)["tags"], ["target"])
|
|
self.assertEqual(
|
|
self.db.rename_tag_all_prompts("missing", "x")["affected_count"], 0
|
|
)
|
|
self.assertEqual(self.db.delete_tag_all_prompts("missing")["affected_count"], 0)
|
|
with self.assertRaises(ValueError):
|
|
self.db.rename_tag_all_prompts(" ", "x")
|
|
with self.assertRaises(ValueError):
|
|
self.db.rename_tag_all_prompts("x", "")
|
|
with self.assertRaises(ValueError):
|
|
self.db.delete_tag_all_prompts("")
|
|
with self.assertRaises(ValueError):
|
|
self.db.merge_tags([], "t")
|
|
with self.assertRaises(ValueError):
|
|
self.db.merge_tags(["a"], " ")
|
|
|
|
result = self.db.merge_tags(["missing", "target"], "final")
|
|
self.assertEqual(result["tags_merged"], 1)
|
|
self.assertEqual(result["affected_count"], 2)
|
|
self.assertEqual(self.db.get_all_tags(), ["final"])
|
|
|
|
def test_untagged_and_category_helpers(self):
|
|
self._save("tagged", tags=["t"], category=" ")
|
|
untagged = self._save("plain", category="shown")
|
|
|
|
self.assertEqual(self.db.get_untagged_prompts_count(), 1)
|
|
page = self.db.get_untagged_prompts(limit=5)
|
|
self.assertEqual([p["id"] for p in page["prompts"]], [untagged])
|
|
self.assertEqual(self.db.get_all_categories(), ["shown"])
|
|
self.assertEqual(self.db.delete_prompts_by_category("nothing"), 0)
|
|
|
|
|
|
class TestRowConversion(DatabaseTestCase):
|
|
def test_legacy_json_tags_column_is_parsed(self):
|
|
convert = self.db._row_to_dict
|
|
self.assertEqual(convert({"tags": '["a", "b"]'})["tags"], ["a", "b"])
|
|
self.assertEqual(convert({"tags": '"a, b"'})["tags"], ["a", "b"])
|
|
self.assertEqual(convert({"tags": "5"})["tags"], [])
|
|
self.assertEqual(convert({"tags": "{not json"})["tags"], [])
|
|
self.assertEqual(convert({"tags": None})["tags"], [])
|
|
self.assertEqual(convert({"_tag_list": None})["tags"], [])
|
|
|
|
def test_search_date_filters(self):
|
|
pid = self._save("dated")
|
|
with closing(sqlite3.connect(self.temp_db.name)) as conn, conn:
|
|
conn.execute(
|
|
"UPDATE prompts SET created_at = '2026-05-01T00:00:00' WHERE id = ?",
|
|
(pid,),
|
|
)
|
|
self.assertEqual(
|
|
len(self.db.search_prompts(date_from="2026-04-01", date_to="2026-06-01")), 1
|
|
)
|
|
self.assertEqual(len(self.db.search_prompts(date_from="2026-06-01")), 0)
|
|
self.assertEqual(len(self.db.search_prompts(date_to="2026-04-01")), 0)
|
|
self.assertEqual(len(self.db.search_prompts(rating_max=3)), 0)
|
|
|
|
|
|
class TestResolveDbPathFromConfig(unittest.TestCase):
|
|
def test_default_path_comes_from_config_when_importable(self):
|
|
tmpdir = tempfile.mkdtemp()
|
|
self.addCleanup(os.rmdir, tmpdir)
|
|
fake = types.ModuleType("py.config")
|
|
fake.PromptManagerConfig = types.SimpleNamespace(
|
|
DEFAULT_DB_PATH=os.path.join(tmpdir, "configured.db")
|
|
)
|
|
with patch.dict(sys.modules, {"py.config": fake}):
|
|
self.assertEqual(
|
|
_resolve_db_path(None), os.path.join(tmpdir, "configured.db")
|
|
)
|
|
|
|
def test_falls_back_to_prompts_db_when_config_cannot_import(self):
|
|
# Outside ComfyUI py.config fails to import (needs the server module)
|
|
with patch.dict(sys.modules, {"py.config": None}):
|
|
resolved = _resolve_db_path(None)
|
|
root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
self.assertEqual(resolved, os.path.join(root, "prompts.db"))
|
|
|
|
|
|
class TestModelHousekeeping(DatabaseTestCase):
|
|
def test_database_info_reports_counts_and_handles_errors(self):
|
|
self._save("rated", rating=4, category="c")
|
|
info = self.db.model.get_database_info()
|
|
self.assertEqual(info["total_prompts"], 1)
|
|
self.assertEqual(info["unique_categories"], 1)
|
|
self.assertEqual(info["average_rating"], 4.0)
|
|
self.assertGreater(info["database_size_bytes"], 0)
|
|
with patch.object(
|
|
self.db.model, "get_connection", side_effect=sqlite3.OperationalError
|
|
):
|
|
self.assertEqual(self.db.model.get_database_info(), {})
|
|
|
|
def test_vacuum_logs_errors_instead_of_raising(self):
|
|
self._save("x")
|
|
self.db.model.vacuum_database()
|
|
with patch(
|
|
"database.models.sqlite3.connect", side_effect=sqlite3.OperationalError
|
|
):
|
|
self.db.model.vacuum_database() # logged, not raised
|
|
self.assertEqual(len(self.db.search_prompts()), 1)
|
|
|
|
def test_close_tolerates_a_connection_that_fails_to_close(self):
|
|
real = self.db.model.get_connection()
|
|
real.close()
|
|
failing = MagicMock()
|
|
failing.close.side_effect = sqlite3.ProgrammingError("already closed")
|
|
self.db.model._local.conn = failing
|
|
self.db.close() # logged, not raised
|
|
self.assertTrue(failing.close.called)
|
|
self.assertIsNot(self.db.model.get_connection(), failing)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|
|
|
|
class TestImagePaging(DatabaseTestCase):
|
|
"""Image listings page in the database, not by over-fetching in the API."""
|
|
|
|
def _link_many(self, count):
|
|
pid = self._save("paged prompt")
|
|
image_ids = [
|
|
self.db.link_image_to_prompt(pid, f"/out/img{i:03d}.png")
|
|
for i in range(count)
|
|
]
|
|
# Stamp distinct times afterwards through the model's own connection,
|
|
# so no second connection holds a write lock while linking.
|
|
with self.db.model.get_connection() as conn:
|
|
for i, image_id in enumerate(image_ids):
|
|
conn.execute(
|
|
"UPDATE generated_images SET generation_time = ? WHERE id = ?",
|
|
(f"2026-01-01T00:00:{i:02d}", image_id),
|
|
)
|
|
return pid
|
|
|
|
def test_recent_images_offset_skips_newest_rows(self):
|
|
self._link_many(5)
|
|
newest_first = [img["filename"] for img in self.db.get_recent_images(limit=5)]
|
|
page = self.db.get_recent_images(limit=2, offset=2)
|
|
self.assertEqual([img["filename"] for img in page], newest_first[2:4])
|
|
self.assertEqual(self.db.get_recent_images(limit=2, offset=10), [])
|
|
|
|
def test_search_images_by_prompt_is_bounded_and_pages(self):
|
|
self._link_many(5)
|
|
first = self.db.search_images_by_prompt("paged", limit=2)
|
|
second = self.db.search_images_by_prompt("paged", limit=2, offset=2)
|
|
self.assertEqual(len(first), 2)
|
|
self.assertEqual(len(second), 2)
|
|
self.assertNotEqual({img["id"] for img in first}, {img["id"] for img in second})
|
|
self.assertEqual(len(self.db.search_images_by_prompt("paged")), 5)
|
|
|
|
|
|
class TestImageLinkSummaries(DatabaseTestCase):
|
|
"""Batched card summaries follow get_image_prompt_info's matching rules."""
|
|
|
|
def _statements(self, call):
|
|
statements = []
|
|
original = self.db.model.get_connection
|
|
|
|
def traced():
|
|
connection = original()
|
|
connection.set_trace_callback(statements.append)
|
|
return connection
|
|
|
|
try:
|
|
with patch.object(self.db.model, "get_connection", side_effect=traced):
|
|
result = call()
|
|
finally:
|
|
original().set_trace_callback(None)
|
|
selects = [s for s in statements if s.lstrip().upper().startswith("SELECT")]
|
|
return result, selects
|
|
|
|
def test_summary_status_matches_image_prompt_info_for_every_rule(self):
|
|
owner = self._save("Owner prompt\nsecond line", category="Cat", rating=4)
|
|
rival = self._save("Rival prompt")
|
|
exact = os.path.join("out", "one", "same.png")
|
|
other = os.path.join("out", "two", "same.png")
|
|
unique = os.path.join("out", "three", "unique.png")
|
|
shared = os.path.join("out", "shared.png")
|
|
self.db.link_image_to_prompt(owner, exact)
|
|
self.db.link_image_to_prompt(rival, other)
|
|
self.db.link_image_to_prompt(owner, unique)
|
|
self.db.link_image_to_prompt(owner, shared)
|
|
self.db.link_image_to_prompt(rival, shared)
|
|
probes = [
|
|
exact,
|
|
os.path.join("out", "x", "..", "one", "same.png"),
|
|
os.path.join("elsewhere", "same.png"),
|
|
os.path.join("elsewhere", "unique.png"),
|
|
shared,
|
|
"missing.png",
|
|
]
|
|
for fallback in (True, False):
|
|
summaries = self.db.get_image_link_summaries(probes, fallback)
|
|
self.assertEqual(set(summaries), set(probes))
|
|
for probe in probes:
|
|
with self.subTest(probe=probe, fallback=fallback):
|
|
info = self.db.get_image_prompt_info(
|
|
probe, fallback, report_ambiguity=True
|
|
)
|
|
summary = summaries[probe]
|
|
if info is None:
|
|
expected = "none"
|
|
elif info.get("association_status") == "ambiguous":
|
|
expected = "ambiguous"
|
|
else:
|
|
expected = "linked"
|
|
self.assertEqual(summary["status"], expected)
|
|
if expected == "linked":
|
|
self.assertEqual(summary["positive_id"], info["prompt_id"])
|
|
else:
|
|
self.assertEqual(
|
|
summary,
|
|
{
|
|
"status": expected,
|
|
"positive_id": None,
|
|
"title": None,
|
|
"title_source": None,
|
|
"category": None,
|
|
"rating": None,
|
|
},
|
|
)
|
|
linked = self.db.get_image_link_summaries([unique])[unique]
|
|
self.assertEqual(
|
|
linked,
|
|
{
|
|
"status": "linked",
|
|
"positive_id": owner,
|
|
"title": "Owner prompt",
|
|
"title_source": "text",
|
|
"category": "Cat",
|
|
"rating": 4,
|
|
},
|
|
)
|
|
|
|
def test_summary_title_prefers_stored_title_then_first_line(self):
|
|
long_line = "word " * 30
|
|
titled = self._save("Untitled text")
|
|
plain = self._save(f" {long_line}\nsecond")
|
|
blank = self._save("Blank title text")
|
|
with self.db.model.get_connection() as conn:
|
|
conn.executemany(
|
|
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
|
[(titled, "Stored title"), (blank, " \u3000 ")],
|
|
)
|
|
paths = {}
|
|
for name, prompt in (("titled", titled), ("plain", plain), ("blank", blank)):
|
|
paths[name] = os.path.join("out", f"{name}.png")
|
|
self.db.link_image_to_prompt(prompt, paths[name])
|
|
summaries = self.db.get_image_link_summaries(list(paths.values()))
|
|
sources = {
|
|
name: (summaries[path]["title"], summaries[path]["title_source"])
|
|
for name, path in paths.items()
|
|
}
|
|
self.assertEqual(sources["titled"], ("Stored title", "stored"))
|
|
# Cut like the Library's title(): 80 characters, then an ellipsis.
|
|
self.assertEqual(sources["plain"], (long_line.strip()[:80] + "\u2026", "text"))
|
|
self.assertEqual(sources["blank"], ("Blank title text", "text"))
|
|
|
|
def test_display_title_matches_library_title_cut(self):
|
|
from database.library import display_title
|
|
|
|
exact = "x" * 80
|
|
self.assertEqual(display_title((), exact), exact)
|
|
self.assertEqual(display_title((), exact + "y"), exact + "\u2026")
|
|
self.assertEqual(display_title((), " short \nnext"), "short")
|
|
self.assertIsNone(display_title((" ",), " \n"))
|
|
self.assertEqual(display_title(("Kept ",), "text"), "Kept ")
|
|
|
|
def test_summary_uses_one_select_per_200_paths(self):
|
|
prompt = self._save("Batch")
|
|
paths = [os.path.join("out", f"{index}.png") for index in range(201)]
|
|
for path in paths[:5]:
|
|
self.db.link_image_to_prompt(prompt, path)
|
|
result, selects = self._statements(
|
|
lambda: self.db.get_image_link_summaries(paths[:200])
|
|
)
|
|
self.assertEqual(len(selects), 1)
|
|
self.assertEqual(len(result), 200)
|
|
self.assertEqual(
|
|
[result[path]["status"] for path in paths[:6]], ["linked"] * 5 + ["none"]
|
|
)
|
|
_, selects = self._statements(lambda: self.db.get_image_link_summaries(paths))
|
|
self.assertEqual(len(selects), 2)
|
|
_, selects = self._statements(lambda: self.db.get_image_link_summaries([]))
|
|
self.assertEqual(selects, [])
|
|
|
|
def test_windows_style_paths_normalize(self):
|
|
import ntpath
|
|
|
|
def windows_normalize(path):
|
|
return ntpath.normcase(ntpath.normpath(path))
|
|
|
|
prompt = self._save("Windows prompt")
|
|
stored = r"C:\ComfyUI\Output\Sub\Image.PNG"
|
|
with patch("database.operations.normalize_image_path", windows_normalize):
|
|
self.db.link_image_to_prompt(prompt, stored)
|
|
probes = [
|
|
stored,
|
|
r"c:\comfyui\output\SUB\image.png",
|
|
r"C:\ComfyUI\Output\x\..\Sub\Image.PNG",
|
|
"C:/ComfyUI/Output/Sub/Image.PNG",
|
|
r"D:\ComfyUI\Output\Sub\Other.png",
|
|
]
|
|
summaries = self.db.get_image_link_summaries(probes, False)
|
|
for probe in probes:
|
|
with self.subTest(probe=probe):
|
|
info = self.db.get_image_prompt_info(probe, False)
|
|
expected = "linked" if info else "none"
|
|
self.assertEqual(summaries[probe]["status"], expected)
|
|
self.assertEqual(
|
|
[summaries[probe]["status"] for probe in probes],
|
|
["linked"] * 4 + ["none"],
|
|
)
|