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.
1679 lines
72 KiB
Python
1679 lines
72 KiB
Python
"""Opt-in workspace gallery contracts against disposable SQLite and media."""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess # nosec B404
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from PIL import Image, PngImagePlugin
|
|
|
|
from utils.generation_params import extract_from_records
|
|
|
|
try:
|
|
from test_api_images import ImageAPITestCase, make_png
|
|
from open_handles import assert_closed
|
|
from generation_params_cases import CASES, encoder, ksampler
|
|
except ImportError:
|
|
from tests.test_api_images import ImageAPITestCase, make_png
|
|
from tests.open_handles import assert_closed
|
|
from tests.generation_params_cases import CASES, encoder, ksampler
|
|
|
|
# Wall-clock budgets run only on request and untraced (TESTING.md
|
|
# "Performance budgets"); their structure is asserted everywhere.
|
|
PERF_BUDGET = os.environ.get("PROMPTMANAGER_PERF_BUDGET") == "1"
|
|
TRACED = sys.gettrace() is not None or (
|
|
sys.version_info >= (3, 12)
|
|
and any(sys.monitoring.get_tool(tool) for tool in range(6))
|
|
)
|
|
SEED_HARNESS = (
|
|
Path(__file__).resolve().parent / "js" / "support" / "inspection-seed.mjs"
|
|
)
|
|
NODE_TIMEOUT_SECONDS = 120
|
|
|
|
|
|
def embedded_png(path, positive, negative=""):
|
|
"""A PNG whose API prompt graph resolves to *positive* and *negative*.
|
|
|
|
The sampler has a model link, as every ComfyUI KSampler does, so the shared
|
|
parser (which needs one) and the server resolver both read the prompts.
|
|
"""
|
|
graph = {
|
|
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}},
|
|
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"positive": ["1", 0], "negative": ["2", 0], "model": ["4", 0]},
|
|
},
|
|
"4": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "fixture.safetensors"},
|
|
},
|
|
}
|
|
info = PngImagePlugin.PngInfo()
|
|
info.add_text("prompt", json.dumps(graph))
|
|
with Image.new("RGB", (4, 4)) as image:
|
|
image.save(path, pnginfo=info)
|
|
os.utime(path, (100, 100))
|
|
return path
|
|
|
|
|
|
def records_png(path, records, mtime=100):
|
|
"""A PNG carrying the given text records (None values are left out)."""
|
|
info = PngImagePlugin.PngInfo()
|
|
for key, text in records.items():
|
|
if text is not None:
|
|
info.add_text(key, text)
|
|
with Image.new("RGB", (4, 4)) as image:
|
|
image.save(path, pnginfo=info)
|
|
os.utime(path, (mtime, mtime))
|
|
return path
|
|
|
|
|
|
def png_chunk(kind, data):
|
|
"""One PNG chunk: length, type, data and CRC."""
|
|
import struct
|
|
import zlib
|
|
|
|
return (
|
|
struct.pack(">I", len(data))
|
|
+ kind
|
|
+ data
|
|
+ struct.pack(">I", zlib.crc32(kind + data))
|
|
)
|
|
|
|
|
|
def png_with_text_chunks(path, before=(), after=(), mtime=100):
|
|
"""A PNG with ``(keyword, text)`` tEXt chunks before and after its IDAT."""
|
|
import io
|
|
|
|
buffer = io.BytesIO()
|
|
with Image.new("RGB", (4, 4)) as image:
|
|
image.save(buffer, format="PNG")
|
|
data = buffer.getvalue()
|
|
idat = data.index(b"IDAT") - 4
|
|
iend = data.index(b"IEND") - 4
|
|
|
|
def chunks(records):
|
|
return b"".join(
|
|
png_chunk(b"tEXt", key.encode("latin-1") + b"\0" + text.encode("latin-1"))
|
|
for key, text in records
|
|
)
|
|
|
|
path.write_bytes(
|
|
data[:idat] + chunks(before) + data[idat:iend] + chunks(after) + data[iend:]
|
|
)
|
|
os.utime(path, (mtime, mtime))
|
|
return path
|
|
|
|
|
|
def seeded_records(positive, seed):
|
|
"""API prompt records for a one-sampler graph with *positive* and *seed*."""
|
|
graph = {
|
|
"4": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "concept.safetensors"},
|
|
},
|
|
"6": encoder(positive),
|
|
"7": encoder("blurry"),
|
|
"3": ksampler(("6", 0), ("7", 0), seed=seed),
|
|
}
|
|
return {"prompt": json.dumps(graph)}
|
|
|
|
|
|
def header_params_of_images():
|
|
"""The card parser the images module calls (unpatched)."""
|
|
from py.api import images
|
|
|
|
return images.header_params
|
|
|
|
|
|
class TestWorkspaceImages(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 media(self, name, mtime=100, root=None):
|
|
target = (root or self.output_dir) / name
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
if target.suffix == ".mp4":
|
|
target.write_bytes(b"video")
|
|
else:
|
|
make_png(target)
|
|
os.utime(target, (mtime, mtime))
|
|
return target
|
|
|
|
async def get(self, **params):
|
|
response = await self.client.get(
|
|
"/prompt_manager/images/output", params={"mode": "workspace", **params}
|
|
)
|
|
return response, await response.json()
|
|
|
|
async def test_keyset_survives_insert_and_deterministic_ties(self):
|
|
for name, stamp in (("a.png", 300), ("b.png", 200), ("c.png", 200)):
|
|
self.media(name, stamp)
|
|
response, first = await self.get(limit=2)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual([i["filename"] for i in first["images"]], ["a.png", "b.png"])
|
|
self.assertTrue(first["has_more"])
|
|
self.assertNotIn("offset", first)
|
|
self.media("new.png", 400)
|
|
self.api.invalidate_gallery_cache()
|
|
_, second = await self.get(limit=2, cursor=first["next_cursor"])
|
|
self.assertEqual([i["filename"] for i in second["images"]], ["c.png"])
|
|
self.assertFalse(second["has_more"])
|
|
self.assertIsNone(second["next_cursor"])
|
|
self.assertEqual(second["total"], 4)
|
|
|
|
async def test_filters_apply_before_limit_and_sorts_are_complete(self):
|
|
self.media("folder/Zebra.png", 200)
|
|
self.media("folder/apple.png", 300)
|
|
self.media("other/Zebra.mp4", 100)
|
|
_, data = await self.get(q="ZEBRA", media="image", limit=1)
|
|
self.assertEqual(data["total"], 1)
|
|
self.assertEqual(data["images"][0]["filename"], "Zebra.png")
|
|
for order, expected in (
|
|
("newest", ["apple.png", "Zebra.png", "Zebra.mp4"]),
|
|
("oldest", ["Zebra.mp4", "Zebra.png", "apple.png"]),
|
|
("name", ["apple.png", "Zebra.mp4", "Zebra.png"]),
|
|
("name-desc", ["Zebra.png", "Zebra.mp4", "apple.png"]),
|
|
):
|
|
with self.subTest(order=order):
|
|
_, data = await self.get(sort=order)
|
|
self.assertEqual([i["filename"] for i in data["images"]], expected)
|
|
_, data = await self.get(subfolder="folder", media="video")
|
|
self.assertEqual(data["images"], [])
|
|
self.assertIsNone(data["next_cursor"])
|
|
|
|
async def test_cursor_cannot_be_reused_with_changed_filters_or_tampered(self):
|
|
self.media("a.png")
|
|
self.media("b.png")
|
|
_, first = await self.get(limit=1)
|
|
token = first["next_cursor"]
|
|
self.assertIsInstance(token, str)
|
|
for params in (
|
|
{"cursor": token, "q": "a"},
|
|
{"cursor": token, "sort": "oldest"},
|
|
{"cursor": token + "x"},
|
|
{"cursor": "garbage"},
|
|
{"cursor": "x" * 4097},
|
|
):
|
|
response, data = await self.get(**params)
|
|
self.assertEqual(response.status, 400, params)
|
|
self.assertNotIn(str(self.comfy_dir), str(data))
|
|
|
|
async def test_invalid_input_is_safe_and_page_limit_is_capped(self):
|
|
for params in (
|
|
{"root": "-1"},
|
|
{"root": "1"},
|
|
{"root": "no"},
|
|
{"media": "audio"},
|
|
{"sort": "size"},
|
|
{"q": "x" * 501},
|
|
{"limit": "no"},
|
|
{"subfolder": "../secret"},
|
|
):
|
|
response, data = await self.get(**params)
|
|
self.assertEqual(response.status, 400, params)
|
|
self.assertNotIn(str(self.comfy_dir), str(data))
|
|
_, data = await self.get(limit=999999)
|
|
from py.api.prompts import MAX_PAGE_LIMIT
|
|
|
|
self.assertEqual(data["limit"], MAX_PAGE_LIMIT)
|
|
|
|
async def test_root_identity_and_root_specific_folders(self):
|
|
second = self.comfy_dir / "second"
|
|
second.mkdir()
|
|
self.media("first/a.png")
|
|
self.media("nested/b.png", root=second)
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
|
|
_, data = await self.get(root=1)
|
|
self.assertEqual(data["total"], 1)
|
|
self.assertEqual(data["images"][0]["root_index"], 1)
|
|
response = await self.client.get(
|
|
"/prompt_manager/gallery/subfolders?mode=workspace&root=1"
|
|
)
|
|
folders = await response.json()
|
|
self.assertEqual(folders["subfolders"], ["nested"])
|
|
self.assertEqual([r["index"] for r in folders["roots"]], [0, 1])
|
|
self.assertNotIn(str(self.comfy_dir), str(folders))
|
|
response = await self.client.get(
|
|
"/prompt_manager/gallery/subfolders?mode=workspace&root=7"
|
|
)
|
|
self.assertEqual(response.status, 400)
|
|
|
|
async def test_workspace_dimensions_are_header_only_and_video_is_unknown(self):
|
|
self.media("a.png")
|
|
self.media("b.mp4")
|
|
self.media("broken.png").write_bytes(b"invalid")
|
|
with patch(
|
|
"PIL.PngImagePlugin.PngImageFile.load",
|
|
side_effect=AssertionError("Do not decode"),
|
|
):
|
|
_, data = await self.get()
|
|
rows = {row["filename"]: row for row in data["images"]}
|
|
self.assertEqual((rows["a.png"]["width"], rows["a.png"]["height"]), (4, 4))
|
|
self.assertIsNone(rows["b.mp4"]["width"])
|
|
self.assertIsNone(rows["broken.png"]["height"])
|
|
response = await self.client.get("/prompt_manager/images/output")
|
|
self.assertNotIn("width", (await response.json())["images"][0])
|
|
|
|
async def test_invalid_limit_does_not_echo_private_input(self):
|
|
response, data = await self.get(limit=str(self.comfy_dir))
|
|
self.assertEqual(response.status, 400)
|
|
self.assertNotIn(str(self.comfy_dir), str(data))
|
|
|
|
async def test_each_sort_can_traverse_ties_and_limit_can_change(self):
|
|
for name in ("a.png", "b.png", "sub/b.png", "c.png"):
|
|
self.media(name)
|
|
for order in ("newest", "oldest", "name", "name-desc"):
|
|
with self.subTest(order=order):
|
|
_, complete = await self.get(sort=order)
|
|
_, first = await self.get(sort=order, limit=1)
|
|
_, rest = await self.get(
|
|
sort=order, limit=10, cursor=first["next_cursor"]
|
|
)
|
|
self.assertEqual(first["images"] + rest["images"], complete["images"])
|
|
|
|
async def test_root_ties_are_distinct_and_cursor_rejects_root_changes(self):
|
|
second = self.comfy_dir / "second"
|
|
second.mkdir()
|
|
self.media("same.png")
|
|
self.media("same.png", root=second)
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
|
|
_, first = await self.get(limit=1)
|
|
_, second_page = await self.get(limit=1, cursor=first["next_cursor"])
|
|
self.assertEqual(
|
|
[first["images"][0]["root_index"], second_page["images"][0]["root_index"]],
|
|
[0, 1],
|
|
)
|
|
response, _ = await self.get(root=0, cursor=first["next_cursor"])
|
|
self.assertEqual(response.status, 400)
|
|
self.api._get_all_output_dirs = lambda: [second, self.output_dir]
|
|
response, _ = await self.get(cursor=first["next_cursor"])
|
|
self.assertEqual(response.status, 400)
|
|
|
|
async def test_empty_roots_are_successful_empty_workspace(self):
|
|
self.api._get_all_output_dirs = lambda: []
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(data["images"], [])
|
|
self.assertEqual(data["total"], 0)
|
|
self.assertIsNone(data["next_cursor"])
|
|
|
|
async def test_deep_unicode_paths_emit_bounded_usable_cursor(self):
|
|
folder = self.output_dir.joinpath(*(["嵐" * 80] * 12))
|
|
entries = [
|
|
(folder / name, 100, self.output_dir, 0) for name in ("a.png", "b.png")
|
|
]
|
|
|
|
def format_page(page):
|
|
return [{"filename": entry[0].name} for entry in page]
|
|
|
|
with (
|
|
patch.object(
|
|
self.api,
|
|
"_get_gallery_files",
|
|
AsyncMock(return_value=[(entry[0], entry[1]) for entry in entries]),
|
|
),
|
|
patch.object(self.api, "_format_workspace_page", side_effect=format_page),
|
|
):
|
|
for order in ("newest", "oldest", "name", "name-desc"):
|
|
with self.subTest(sort=order):
|
|
response, first = await self.get(sort=order, limit=1)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertLessEqual(len(first["next_cursor"]), 4096)
|
|
response, second = await self.get(
|
|
sort=order, cursor=first["next_cursor"]
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(len(second["images"]), 1)
|
|
self.assertNotEqual(first["images"], second["images"])
|
|
|
|
async def test_filter_and_sort_selection_runs_in_executor(self):
|
|
self.media("a.png")
|
|
import threading
|
|
|
|
event_thread = threading.get_ident()
|
|
from py.api import images as images_module
|
|
|
|
original = images_module._workspace_key
|
|
threads = []
|
|
|
|
def observe(entry, order):
|
|
threads.append(threading.get_ident())
|
|
return original(entry, order)
|
|
|
|
with patch.object(images_module, "_workspace_key", side_effect=observe):
|
|
response, _ = await self.get()
|
|
self.assertEqual(response.status, 200)
|
|
self.assertTrue(threads)
|
|
self.assertTrue(all(thread != event_thread for thread in threads))
|
|
|
|
async def test_workspace_reuses_scan_cache_and_legacy_contract(self):
|
|
self.media("a.png")
|
|
with patch.object(
|
|
self.api,
|
|
"_scan_gallery_files_sync",
|
|
wraps=self.api._scan_gallery_files_sync,
|
|
) as scan:
|
|
await self.get()
|
|
await self.get(q="a")
|
|
self.assertEqual(scan.call_count, 1)
|
|
response = await self.client.get(
|
|
"/prompt_manager/images/output?limit=1&offset=0"
|
|
)
|
|
data = await response.json()
|
|
self.assertEqual(data["offset"], 0)
|
|
self.assertNotIn("next_cursor", data)
|
|
response = await self.client.get("/prompt_manager/gallery/subfolders")
|
|
self.assertNotIn("roots", await response.json())
|
|
|
|
async def test_root_scope_workspace_descriptors_guard_original_and_thumbnail_urls(
|
|
self,
|
|
):
|
|
from urllib.parse import parse_qs, quote, urlsplit
|
|
|
|
roots = [self.output_dir, self.comfy_dir / "input"]
|
|
self.api._get_all_output_dirs = lambda: roots
|
|
relative = "nested/猫 % #.png"
|
|
for index, root in enumerate(roots):
|
|
self.media(relative, root=root)
|
|
self.media("thumbnails/nested/猫 % #_thumb.png", root=root)
|
|
make_png(root / relative, color=(index, 0, 0))
|
|
make_png(root / "thumbnails/nested/猫 % #_thumb.png", color=(index, 0, 0))
|
|
_, result = await self.get(limit=10)
|
|
for entry in result["images"]:
|
|
self.assertRegex(entry.get("root_scope", ""), r"^[0-9a-f]{64}$")
|
|
self.assertIn(quote(relative, safe="/"), entry["url"])
|
|
for url in (entry["url"], entry["thumbnail_url"]):
|
|
query = parse_qs(urlsplit(url).query)
|
|
self.assertEqual(query["root_scope"], [entry["root_scope"]])
|
|
self.assertEqual(query["root"], [str(entry["root_index"])])
|
|
response = await self.client.get(url)
|
|
self.assertEqual(response.status, 200)
|
|
expected = roots[entry["root_index"]] / (
|
|
"thumbnails/nested/猫 % #_thumb.png"
|
|
if url == entry["thumbnail_url"]
|
|
else relative
|
|
)
|
|
self.assertEqual(await response.read(), expected.read_bytes())
|
|
self.assertIn("v", parse_qs(urlsplit(entry["thumbnail_url"]).query))
|
|
response = await self.client.get(
|
|
f"/prompt_manager/images/inspect/{quote(relative, safe='/')}",
|
|
params={"root": entry["root_index"], "root_scope": entry["root_scope"]},
|
|
)
|
|
data = await response.json()
|
|
self.assertEqual(data["root_scope"], entry["root_scope"])
|
|
self.assertEqual(
|
|
data["image"],
|
|
{
|
|
"root_index": entry["root_index"],
|
|
"path": relative,
|
|
"filename": "猫 % #.png",
|
|
},
|
|
)
|
|
|
|
def _link_fixtures(self):
|
|
"""Linked (titled, first-line, pair, untitled pair), ambiguous and none."""
|
|
db = self.api.db
|
|
long_line = "lantern " * 15
|
|
titled = self._save_prompt("Titled positive", category="Night", rating=5)
|
|
with db.model.get_connection() as connection:
|
|
connection.execute(
|
|
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
|
(titled, "Stored library title"),
|
|
)
|
|
plain = self._save_prompt(long_line)
|
|
pair = db.save_prompt_pair("Pair positive", "Noise", title="Pair title")
|
|
bare = db.save_prompt_pair("Bare pair positive\nmore", "Blur")
|
|
first, second = self._save_prompt("Owner one"), self._save_prompt("Owner two")
|
|
files = {
|
|
name: self.media(f"{name}.png")
|
|
for name in ("titled", "pair", "ambiguous", "none")
|
|
}
|
|
# First-line titles come from files whose own prompt evidence agrees.
|
|
files["plain"] = embedded_png(self.output_dir / "plain.png", long_line)
|
|
files["bare"] = embedded_png(
|
|
self.output_dir / "bare.png", "Bare pair positive\nmore", "Blur"
|
|
)
|
|
db.link_image_to_prompt(titled, str(files["titled"]))
|
|
db.link_image_to_prompt(plain, str(files["plain"]))
|
|
for name, pair_id in (("pair", pair), ("bare", bare)):
|
|
positive = db.get_prompt_pair(pair_id)["positive_prompt_id"]
|
|
image = db.link_image_to_prompt(positive, str(files[name]))
|
|
db.link_image_to_pair(pair_id, image, {"source": "fixture"})
|
|
db.link_image_to_prompt(first, str(files["ambiguous"]))
|
|
db.link_image_to_prompt(second, str(files["ambiguous"]))
|
|
return {
|
|
"titled": ("linked", titled, "Stored library title"),
|
|
"plain": ("linked", plain, long_line.strip()[:80] + "\u2026"),
|
|
"pair": (
|
|
"linked",
|
|
db.get_prompt_pair(pair)["positive_prompt_id"],
|
|
"Pair title",
|
|
),
|
|
"bare": (
|
|
"linked",
|
|
db.get_prompt_pair(bare)["positive_prompt_id"],
|
|
"Bare pair positive",
|
|
),
|
|
"ambiguous": ("ambiguous", None, None),
|
|
"none": ("none", None, None),
|
|
}
|
|
|
|
async def test_workspace_items_carry_link_summary_matching_inspection(self):
|
|
expected = self._link_fixtures()
|
|
_, data = await self.get(limit=50)
|
|
items = {Path(item["filename"]).stem: item for item in data["images"]}
|
|
self.assertEqual(set(items), set(expected))
|
|
for name, (status, positive_id, title) in expected.items():
|
|
with self.subTest(file=name):
|
|
linked = items[name]["linked"]
|
|
self.assertEqual(
|
|
set(linked),
|
|
{
|
|
"status",
|
|
"positive_id",
|
|
"title",
|
|
"title_source",
|
|
"category",
|
|
"rating",
|
|
},
|
|
)
|
|
self.assertEqual(
|
|
(linked["status"], linked["positive_id"], linked["title"]),
|
|
(status, positive_id, title),
|
|
)
|
|
response = await self.client.get(
|
|
f"/prompt_manager/images/inspect/{name}.png",
|
|
params={"root": 0, "root_scope": items[name]["root_scope"]},
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
library = (await response.json())["library"]
|
|
self.assertEqual(
|
|
(library["status"], library["positive_id"], library["title"]),
|
|
(linked["status"], linked["positive_id"], linked["title"]),
|
|
)
|
|
self.assertEqual(
|
|
(
|
|
items["titled"]["linked"]["category"],
|
|
items["titled"]["linked"]["rating"],
|
|
),
|
|
("Night", 5),
|
|
)
|
|
|
|
async def test_stored_title_shows_on_card_and_inspector_without_evidence(self):
|
|
# A stored title is user data: shown on both, even when it equals the
|
|
# first line and the file carries no prompt metadata.
|
|
prompt = self._save_prompt("Same words\nrest")
|
|
with self.api.db.model.get_connection() as connection:
|
|
connection.execute(
|
|
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
|
(prompt, "Same words"),
|
|
)
|
|
path = self.media("stored.png")
|
|
self.api.db.link_image_to_prompt(prompt, str(path))
|
|
_, data = await self.get()
|
|
linked = data["images"][0]["linked"]
|
|
self.assertEqual(
|
|
(linked["title"], linked["title_source"]), ("Same words", "stored")
|
|
)
|
|
library = await self.inspected("stored.png", data["images"][0])
|
|
self.assertEqual(library["title"], linked["title"])
|
|
self.assertEqual(library["positive_id"], prompt)
|
|
|
|
async def test_unconfirmed_text_title_is_withheld_on_card_and_inspector(self):
|
|
prompt = self._save_prompt("Secret first line\nrest")
|
|
path = self.media("plain.png")
|
|
self.api.db.link_image_to_prompt(prompt, str(path))
|
|
_, data = await self.get()
|
|
item = data["images"][0]
|
|
response = await self.client.get(
|
|
"/prompt_manager/images/inspect/plain.png",
|
|
params={"root": 0, "root_scope": item["root_scope"]},
|
|
)
|
|
body = await response.json()
|
|
self.assertEqual(body["library"]["positive_id"], prompt)
|
|
self.assertIsNone(body["library"]["title"])
|
|
self.assertNotIn("Secret first line", json.dumps(body))
|
|
self.assertIsNone(item["linked"]["title"])
|
|
self.assertNotIn("Secret first line", json.dumps(data))
|
|
|
|
async def test_summary_failure_keeps_inspector_link(self):
|
|
prompt = self._save_prompt("Readable link")
|
|
path = self.media("a.png")
|
|
self.api.db.link_image_to_prompt(prompt, str(path))
|
|
_, data = await self.get()
|
|
with patch.object(
|
|
self.api.db,
|
|
"get_image_link_summaries",
|
|
side_effect=RuntimeError("summary gone"),
|
|
):
|
|
library = await self.inspected("a.png", data["images"][0])
|
|
self.assertEqual(library["status"], "linked")
|
|
self.assertEqual(library["positive_id"], prompt)
|
|
self.assertIsNone(library["title"])
|
|
|
|
async def inspected(self, name, item):
|
|
response = await self.client.get(
|
|
f"/prompt_manager/images/inspect/{name}",
|
|
params={"root": 0, "root_scope": item["root_scope"]},
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
return (await response.json())["library"]
|
|
|
|
async def test_link_summary_uses_one_query_per_page(self):
|
|
prompt = self._save_prompt("Paged prompt")
|
|
for index in range(12):
|
|
path = self.media(f"{index:02}.png", mtime=100 + index)
|
|
if index % 2:
|
|
self.api.db.link_image_to_prompt(prompt, str(path))
|
|
statements = []
|
|
original = self.api.db.model.get_connection
|
|
|
|
def trace_connection():
|
|
connection = original()
|
|
connection.set_trace_callback(statements.append)
|
|
return connection
|
|
|
|
with patch.object(
|
|
self.api.db.model, "get_connection", side_effect=trace_connection
|
|
):
|
|
response, data = await self.get(limit=12)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(len(data["images"]), 12)
|
|
self.assertEqual(
|
|
sorted(item["linked"]["status"] for item in data["images"]),
|
|
["linked"] * 6 + ["none"] * 6,
|
|
)
|
|
selects = [
|
|
query
|
|
for query in statements
|
|
if query.lstrip().upper().startswith(("SELECT", "WITH"))
|
|
]
|
|
# The page's link summary, the positive count's one link read and
|
|
# the page's one favorites read.
|
|
self.assertEqual(len(selects), 3, selects)
|
|
self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects)
|
|
self.assertEqual(data["positive_count"], 1) # six files, one prompt
|
|
|
|
async def test_link_summary_runs_in_executor_and_survives_database_errors(self):
|
|
import threading
|
|
|
|
self.media("a.png")
|
|
event_thread = threading.get_ident()
|
|
threads = []
|
|
|
|
def broken(paths, allow_filename_fallback=True):
|
|
threads.append(threading.get_ident())
|
|
raise RuntimeError("database is gone")
|
|
|
|
with patch.object(self.api.db, "get_image_link_summaries", side_effect=broken):
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200)
|
|
self.assertTrue(threads)
|
|
self.assertNotEqual(threads[0], event_thread)
|
|
self.assertEqual(data["images"][0]["linked"]["status"], "unavailable")
|
|
self.assertIsNone(data["images"][0]["linked"]["title"])
|
|
self.assertIsNone(data["images"][0]["linked"]["title_source"])
|
|
self.assertNotIn("database is gone", str(data))
|
|
|
|
async def test_items_carry_seed_from_header_records(self):
|
|
records_png(self.output_dir / "seeded.png", seeded_records("a fox", 42001))
|
|
self.media("plain.png")
|
|
self.media("clip.mp4")
|
|
self.media("broken.png").write_bytes(b"invalid")
|
|
_, data = await self.get()
|
|
seeds = {item["filename"]: item["seed"] for item in data["images"]}
|
|
self.assertEqual(
|
|
seeds,
|
|
{
|
|
"seeded.png": "42001",
|
|
"plain.png": None,
|
|
"clip.mp4": None,
|
|
"broken.png": None,
|
|
},
|
|
)
|
|
|
|
async def test_card_reads_each_header_once(self):
|
|
from py.api import images
|
|
|
|
for index in range(3):
|
|
records_png(
|
|
self.output_dir / f"{index}.png", seeded_records("a fox", index)
|
|
)
|
|
self.media("clip.mp4")
|
|
with (
|
|
patch.object(
|
|
images, "read_inspection_file", wraps=images.read_inspection_file
|
|
) as read,
|
|
patch("PIL.Image.open", wraps=Image.open) as opened,
|
|
):
|
|
_, data = await self.get()
|
|
self.assertEqual(len(data["images"]), 4)
|
|
# One inspector read per PNG (dimensions and records); no decoder open.
|
|
self.assertEqual(read.call_count, 3)
|
|
self.assertEqual(opened.call_count, 0)
|
|
|
|
async def test_card_reads_png_records_as_the_inspector_does(self):
|
|
one, two = seeded_records("a fox", 1), seeded_records("a fox", 2)
|
|
# Conflicting duplicate prompt chunks: the inspector refuses the file,
|
|
# so the card shows no seed (Pillow alone would keep the last one).
|
|
png_with_text_chunks(
|
|
self.output_dir / "conflict.png",
|
|
before=[("prompt", one["prompt"]), ("prompt", two["prompt"])],
|
|
mtime=300,
|
|
)
|
|
# Text after IDAT: Pillow's header open misses it; the inspector reads it.
|
|
png_with_text_chunks(
|
|
self.output_dir / "late.png", after=[("prompt", one["prompt"])], mtime=200
|
|
)
|
|
png_with_text_chunks(
|
|
self.output_dir / "repeat.png",
|
|
before=[("prompt", one["prompt"]), ("prompt", one["prompt"])],
|
|
)
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200, data)
|
|
items = {item["filename"]: item for item in data["images"]}
|
|
self.assertIsNone(items["conflict.png"]["seed"])
|
|
self.assertEqual(
|
|
(items["conflict.png"]["width"], items["conflict.png"]["height"]), (4, 4)
|
|
)
|
|
self.assertEqual(items["late.png"]["seed"], "1")
|
|
self.assertEqual(items["repeat.png"]["seed"], "1")
|
|
shown = await self._inspector_seeds(data["images"])
|
|
self.assertNotEqual(shown["conflict.png"]["status"], "known")
|
|
self.assertEqual(shown["late.png"], {"status": "known", "value": "1"})
|
|
|
|
async def test_deeply_nested_record_still_lists_and_inspects(self):
|
|
nested = "[" * 900 + "1" + "]" * 900
|
|
workflow = (
|
|
'{"nodes":[{"id":%s,"type":"KSampler","widgets_values":[1]}],'
|
|
'"links":[[%s,1,0,2,0]]}' % (nested, nested)
|
|
)
|
|
records_png(self.output_dir / "ok.png", seeded_records("a fox", 5), mtime=200)
|
|
records_png(self.output_dir / "evil.png", {"workflow": workflow})
|
|
prompt = self._save_prompt("Evil first line")
|
|
self.api.db.link_image_to_prompt(prompt, str(self.output_dir / "evil.png"))
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200, data)
|
|
items = {item["filename"]: item for item in data["images"]}
|
|
self.assertEqual(items["ok.png"]["seed"], "5")
|
|
response = await self.client.get(
|
|
"/prompt_manager/images/inspect/evil.png",
|
|
params={"root": 0, "root_scope": items["evil.png"]["root_scope"]},
|
|
)
|
|
self.assertEqual(response.status, 200, await response.text())
|
|
library = (await response.json())["library"]
|
|
self.assertEqual(library["title"], items["evil.png"]["linked"]["title"])
|
|
shown = (await self._inspector_seeds([items["evil.png"]]))["evil.png"]
|
|
card = items["evil.png"]["seed"]
|
|
if card is None:
|
|
self.assertNotEqual(shown["status"], "known")
|
|
else:
|
|
self.assertEqual(shown, {"status": "known", "value": card})
|
|
|
|
async def test_warm_page_reuses_card_facts_per_scan(self):
|
|
from py.api import images
|
|
|
|
a = records_png(self.output_dir / "a.png", seeded_records("a fox", 1), 300)
|
|
b = records_png(self.output_dir / "b.png", seeded_records("a cat", 2), 200)
|
|
gone = records_png(self.output_dir / "gone.png", seeded_records("x", 3), 100)
|
|
real_read, real_parse = images.read_inspection_file, images.header_params
|
|
read, parsed = [], []
|
|
|
|
def reading(path):
|
|
read.append(path.name)
|
|
return real_read(path)
|
|
|
|
def parsing(info):
|
|
parsed.append(info)
|
|
return real_parse(info)
|
|
|
|
async def page():
|
|
read.clear()
|
|
parsed.clear()
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200, data)
|
|
return {item["filename"]: item for item in data["images"]}
|
|
|
|
with (
|
|
patch.object(images, "read_inspection_file", side_effect=reading),
|
|
patch.object(images, "header_params", side_effect=parsing),
|
|
):
|
|
cold = await page()
|
|
self.assertEqual(sorted(read), ["a.png", "b.png", "gone.png"])
|
|
self.assertEqual(len(parsed), 3)
|
|
# A warm page reads and parses no header.
|
|
prompt = self._save_prompt("a fox")
|
|
self.api.db.link_image_to_prompt(prompt, str(a))
|
|
warm = await page()
|
|
self.assertEqual((read, parsed), ([], []))
|
|
self.assertEqual(warm["a.png"]["seed"], "1")
|
|
self.assertEqual((warm["a.png"]["width"], warm["a.png"]["height"]), (4, 4))
|
|
# Links are never cached: the new link's confirmed title shows.
|
|
self.assertIsNone(cold["a.png"]["linked"]["title"])
|
|
self.assertEqual(warm["a.png"]["linked"]["title"], "a fox")
|
|
# A monitor invalidation keeps the values of unchanged files.
|
|
records_png(b, seeded_records("a cat", 22), mtime=250)
|
|
gone.unlink()
|
|
self.api.invalidate_gallery_cache()
|
|
changed = await page()
|
|
self.assertEqual(read, ["b.png"])
|
|
self.assertEqual(changed["b.png"]["seed"], "22")
|
|
cards = self.api._workspace_card_cache[(0, str(self.output_dir))]
|
|
self.assertEqual(
|
|
sorted(path.name for path, _mtime in cards), ["a.png", "b.png"]
|
|
)
|
|
# A header that cannot be read (locked, half-written) is not stored.
|
|
records_png(self.output_dir / "c.png", seeded_records("a cow", 4), 400)
|
|
self.api.invalidate_gallery_cache()
|
|
with patch.object(images, "_card_header", side_effect=OSError("busy")):
|
|
locked = await page()
|
|
self.assertIsNone(locked["c.png"]["seed"])
|
|
self.assertEqual(locked["a.png"]["seed"], "1")
|
|
self.assertEqual((await page())["c.png"]["seed"], "4")
|
|
self.assertEqual(read, ["c.png"])
|
|
# A roots change drops them, as it drops the per-file values.
|
|
self.api.reset_gallery_roots_cache()
|
|
self.assertEqual(self.api._workspace_card_cache, {})
|
|
await page()
|
|
self.assertEqual(sorted(read), ["a.png", "b.png", "c.png"])
|
|
|
|
async def test_unbounded_digit_string_seed_still_lists(self):
|
|
records_png(
|
|
self.output_dir / "huge.png", seeded_records("a fox", "9" * 5000), 200
|
|
)
|
|
records_png(self.output_dir / "ok.png", seeded_records("a fox", 5))
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200, data)
|
|
items = {item["filename"]: item for item in data["images"]}
|
|
self.assertIsNone(items["huge.png"]["seed"])
|
|
self.assertEqual(items["ok.png"]["seed"], "5")
|
|
|
|
async def test_one_cards_parse_error_spares_the_page(self):
|
|
prompt = self._save_prompt("a fox")
|
|
for name, mtime in (("bad.png", 200), ("ok.png", 100)):
|
|
path = records_png(
|
|
self.output_dir / name, seeded_records("a fox", 5), mtime
|
|
)
|
|
self.api.db.link_image_to_prompt(prompt, str(path))
|
|
real_parse = header_params_of_images()
|
|
|
|
def parsing(info):
|
|
if parsing.calls == 0:
|
|
parsing.calls += 1
|
|
raise RuntimeError("unexpected parser bug")
|
|
return real_parse(info)
|
|
|
|
parsing.calls = 0
|
|
with (
|
|
patch("py.api.images.header_params", side_effect=parsing),
|
|
self.assertLogs(self.api.logger, "WARNING") as logged,
|
|
):
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200, data)
|
|
items = {item["filename"]: item for item in data["images"]}
|
|
self.assertIsNone(items["bad.png"]["seed"])
|
|
self.assertIsNone(items["bad.png"]["linked"]["title"])
|
|
self.assertEqual(items["bad.png"]["linked"]["status"], "linked")
|
|
self.assertEqual(items["ok.png"]["seed"], "5")
|
|
self.assertEqual(items["ok.png"]["linked"]["title"], "a fox")
|
|
self.assertNotIn("unexpected parser bug", json.dumps(data))
|
|
self.assertIn("unexpected parser bug", "\n".join(logged.output))
|
|
|
|
async def test_library_previews_skip_card_facts(self):
|
|
records_png(self.output_dir / "a.png", seeded_records("a fox", 5))
|
|
page = await self.api._collect_output_images(
|
|
self.api._get_all_output_dirs(), ""
|
|
)
|
|
with patch("py.api.images.header_params") as parsed:
|
|
items = self.api._format_workspace_page(page, linked=False)
|
|
parsed.assert_not_called()
|
|
self.assertNotIn("seed", items[0])
|
|
self.assertEqual((items[0]["width"], items[0]["height"]), (4, 4))
|
|
|
|
async def test_card_facts_parse_after_the_file_is_closed(self):
|
|
records_png(self.output_dir / "a.png", seeded_records("a fox", 5))
|
|
streams, closed = [], []
|
|
real_open, real_parse = Path.open, header_params_of_images()
|
|
|
|
def opening(path, *args, **kwargs):
|
|
stream = real_open(path, *args, **kwargs)
|
|
if path.name == "a.png":
|
|
streams.append(stream)
|
|
return stream
|
|
|
|
def parsing(info):
|
|
closed.append(bool(streams) and all(s.closed for s in streams))
|
|
return real_parse(info)
|
|
|
|
with (
|
|
patch.object(Path, "open", autospec=True, side_effect=opening),
|
|
patch("py.api.images.header_params", side_effect=parsing),
|
|
):
|
|
_, data = await self.get()
|
|
self.assertEqual(data["images"][0]["seed"], "5")
|
|
self.assertEqual(closed, [True])
|
|
|
|
def _seed_files(self):
|
|
"""Concept-style one-sampler files plus every shared parser case."""
|
|
expected = {}
|
|
for i in range(2):
|
|
for j in range(3):
|
|
name = f"concept-{i}-{j}.png"
|
|
seed = 42001 + i * 100 + j
|
|
records_png(
|
|
self.output_dir / name, seeded_records(f"concept {i} {j}", seed)
|
|
)
|
|
expected[name] = str(seed)
|
|
# ComfyUI's randomized seeds mostly exceed 2**53, where JSON.parse rounds.
|
|
big = 2**63 + 12345
|
|
records_png(self.output_dir / "big-seed.png", seeded_records("big", big))
|
|
expected["big-seed.png"] = str(big)
|
|
for case, records in CASES.items():
|
|
records_png(self.output_dir / f"case-{case}.png", records)
|
|
records_png(
|
|
self.output_dir / "a1111.png",
|
|
{"parameters": "a fox\nNegative prompt: blur\nSteps: 20, Seed: 9"},
|
|
)
|
|
self.media("bare.png")
|
|
return expected
|
|
|
|
async def _inspector_seeds(self, items):
|
|
"""The seed the inspector shows for each listed item (Node harness)."""
|
|
return (await self._node_seeds(items))["inspector"]
|
|
|
|
async def _node_seeds(self, items, listing=None):
|
|
"""The harness output: inspector seeds and, from *listing*, card seeds.
|
|
|
|
*listing* is the exact listing response text; the harness JSON.parses
|
|
it as the browser does, so the card seeds are what a page would read.
|
|
"""
|
|
node = shutil.which("node")
|
|
if node is None: # pragma: no cover - CI runners and dev hosts ship Node
|
|
if os.environ.get("CI"):
|
|
self.fail("Node.js is required in CI for the seed parity check")
|
|
self.skipTest("Node.js is not installed")
|
|
bodies = []
|
|
for item in items:
|
|
response = await self.client.get(
|
|
f"/prompt_manager/images/inspect/{item['path']}",
|
|
params={"root": 0, "root_scope": item["root_scope"]},
|
|
)
|
|
self.assertEqual(response.status, 200, item["path"])
|
|
bodies.append({"name": item["path"], "body": await response.text()})
|
|
# Fixed local Node harness, no shell.
|
|
completed = subprocess.run( # nosec B603
|
|
[node, str(SEED_HARNESS)],
|
|
input=json.dumps({"listing": listing, "inspections": bodies}).encode(
|
|
"utf-8"
|
|
),
|
|
capture_output=True,
|
|
timeout=NODE_TIMEOUT_SECONDS,
|
|
check=False,
|
|
)
|
|
self.assertEqual(completed.returncode, 0, completed.stderr.decode("utf-8"))
|
|
return json.loads(completed.stdout)
|
|
|
|
async def test_card_seed_equals_inspector_seed_when_both_present(self):
|
|
expected = self._seed_files()
|
|
response = await self.client.get(
|
|
"/prompt_manager/images/output",
|
|
params={"mode": "workspace", "limit": 200},
|
|
)
|
|
listing = await response.text()
|
|
items = json.loads(listing)["images"]
|
|
self.assertEqual(len(items), len(CASES) + len(expected) + 2)
|
|
seeds = await self._node_seeds(items, listing)
|
|
inspector, browser = seeds["inspector"], seeds["card"]
|
|
both = 0
|
|
for item in items:
|
|
with self.subTest(file=item["path"]):
|
|
card, shown = item["seed"], inspector[item["path"]]
|
|
# Sent as a decimal string, so JSON.parse keeps every digit.
|
|
self.assertTrue(card is None or isinstance(card, str), card)
|
|
self.assertEqual(browser[item["path"]], card)
|
|
if item["path"] in expected:
|
|
self.assertEqual(card, expected[item["path"]])
|
|
if card is None:
|
|
continue
|
|
# The card never shows a seed the inspector does not.
|
|
self.assertEqual(shown, {"status": "known", "value": card})
|
|
both += 1
|
|
self.assertGreaterEqual(both, len(expected) + 10)
|
|
self.assertGreater(int(browser["big-seed.png"]), 2**53)
|
|
|
|
async def test_text_title_needs_the_files_parsed_positive(self):
|
|
prompt = self._save_prompt("Saved words\nrest")
|
|
other = records_png(
|
|
self.output_dir / "other.png", seeded_records("Different words", 1)
|
|
)
|
|
shared = records_png(
|
|
self.output_dir / "shared.png", CASES["shared-prompt-pipeline"]
|
|
)
|
|
clip = self.media("clip.mp4")
|
|
shared_prompt = self._save_prompt(
|
|
extract_from_records(CASES["shared-prompt-pipeline"]["prompt"], None)[
|
|
"positivePrompt"
|
|
]
|
|
)
|
|
for path, owner in ((other, prompt), (shared, shared_prompt), (clip, prompt)):
|
|
self.api.db.link_image_to_prompt(owner, str(path))
|
|
_, data = await self.get()
|
|
items = {item["filename"]: item for item in data["images"]}
|
|
for name in ("other.png", "clip.mp4"):
|
|
with self.subTest(file=name):
|
|
self.assertEqual(items[name]["linked"]["status"], "linked")
|
|
self.assertEqual(items[name]["linked"]["title_source"], "text")
|
|
self.assertIsNone(items[name]["linked"]["title"])
|
|
library = await self.inspected("other.png", items["other.png"])
|
|
self.assertIsNone(library["title"])
|
|
self.assertNotIn("Saved words", json.dumps(data))
|
|
# The resolver cannot decide this graph; the shared parser (and so the
|
|
# inspector's prompt) confirms the saved first line on both.
|
|
title = items["shared.png"]["linked"]["title"]
|
|
self.assertIsNotNone(title)
|
|
self.assertEqual(
|
|
(await self.inspected("shared.png", items["shared.png"]))["title"], title
|
|
)
|
|
|
|
# B3: search, category/tag filters and rating sort over linked positives.
|
|
|
|
def _names(self, data):
|
|
return [item["filename"] for item in data["images"]]
|
|
|
|
def _filter_fixtures(self):
|
|
"""Linked, pair-linked, ambiguous and unlinked files for the B3 filters."""
|
|
db = self.api.db
|
|
night_ab = self._save_prompt(
|
|
"A ZEBRA under neon", category="Night", tags=["a", "b"], rating=5
|
|
)
|
|
night_a = self._save_prompt("Quiet harbor", category="Night", tags=["a"])
|
|
day_ab = self._save_prompt(
|
|
"Morning field", category="Day", tags=["a", "b", "zebra-tag"], rating=3
|
|
)
|
|
titled = self._save_prompt("Plain words", category="Day", rating=1)
|
|
with db.model.get_connection() as connection:
|
|
connection.execute(
|
|
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
|
(titled, "Zebra stored title"),
|
|
)
|
|
percent = self._save_prompt("100% wild_card", category="Day")
|
|
pair_positive = self._save_prompt("Pair positive", category="Night", rating=4)
|
|
pair = db.save_prompt_pair("Pair positive", "Noise", title="Zebra pair")
|
|
owners = (
|
|
self._save_prompt("Zebra one", category="Night"),
|
|
self._save_prompt("Zebra two", category="Night"),
|
|
)
|
|
files = {
|
|
name: self.media(f"{name}.png", mtime=100 + index)
|
|
for index, name in enumerate(
|
|
(
|
|
"night_ab",
|
|
"night_a",
|
|
"day_ab",
|
|
"titled",
|
|
"percent",
|
|
"paired",
|
|
"ambiguous",
|
|
"zebra_unlinked",
|
|
"plain_unlinked",
|
|
)
|
|
)
|
|
}
|
|
db.link_image_to_prompt(night_ab, str(files["night_ab"]))
|
|
# A non-normalized spelling of the same file still compares equal.
|
|
db.link_image_to_prompt(
|
|
night_a, str(self.output_dir / "sub" / ".." / "night_a.png")
|
|
)
|
|
db.link_image_to_prompt(day_ab, str(files["day_ab"]))
|
|
db.link_image_to_prompt(titled, str(files["titled"]))
|
|
db.link_image_to_prompt(percent, str(files["percent"]))
|
|
self.assertEqual(db.get_prompt_pair(pair)["positive_prompt_id"], pair_positive)
|
|
image = db.link_image_to_prompt(pair_positive, str(files["paired"]))
|
|
db.link_image_to_pair(pair, image, {"source": "fixture"})
|
|
for owner in owners:
|
|
db.link_image_to_prompt(owner, str(files["ambiguous"]))
|
|
return files
|
|
|
|
async def test_positive_count_counts_distinct_linked_positives(self):
|
|
files = self._filter_fixtures()
|
|
db = self.api.db
|
|
night_ab = db.get_prompt_id_for_image(str(files["night_ab"]))
|
|
# A second file of one positive still counts that positive once.
|
|
again = self.media("again.png", mtime=50)
|
|
db.link_image_to_prompt(night_ab, str(again))
|
|
cases = (
|
|
# Six unambiguous positives; the ambiguous file and unlinked count none.
|
|
({}, 10, 6),
|
|
# Text, title, tag, pair title (its positive) and the filename match.
|
|
({"q": "ZEBRA"}, 6, 4),
|
|
({"category": "Night"}, 4, 3),
|
|
({"tags": "a,b"}, 3, 2),
|
|
({"q": "plain_unlinked"}, 1, 0),
|
|
({"media": "video"}, 0, 0),
|
|
({"sort": "rating"}, 10, 6),
|
|
({"sort": "name", "category": "Day"}, 3, 3),
|
|
)
|
|
for params, total, positives in cases:
|
|
with self.subTest(**params):
|
|
response, data = await self.get(**params)
|
|
self.assertEqual(response.status, 200, data)
|
|
self.assertEqual(data["total"], total)
|
|
self.assertEqual(data["positive_count"], positives)
|
|
# The count covers the whole matching set. A filter or the rating
|
|
# sort reads the links on every page, so its cursor pages carry it too.
|
|
for params in ({"q": "ZEBRA"}, {"sort": "rating"}):
|
|
with self.subTest(cursor=params):
|
|
response, first = await self.get(limit=2, **params)
|
|
response, second = await self.get(
|
|
limit=2, cursor=first["next_cursor"], **params
|
|
)
|
|
self.assertEqual(response.status, 200, second)
|
|
self.assertEqual(second["positive_count"], first["positive_count"])
|
|
# A plain cursor page skips the read and omits the key (absent, not
|
|
# null): the frontend keeps the first page's value.
|
|
for params in ({}, {"sort": "oldest"}, {"sort": "name"}):
|
|
with self.subTest(plain_cursor=params):
|
|
response, first = await self.get(limit=2, **params)
|
|
self.assertEqual(first["positive_count"], 6)
|
|
response, second = await self.get(
|
|
limit=2, cursor=first["next_cursor"], **params
|
|
)
|
|
self.assertEqual(response.status, 200, second)
|
|
self.assertNotIn("positive_count", second)
|
|
self.assertEqual(second["total"], 10)
|
|
# One SELECT counts the positives; the page's link summary is the other.
|
|
statements = []
|
|
original = db.model.get_connection
|
|
|
|
def trace_connection():
|
|
connection = original()
|
|
connection.set_trace_callback(statements.append)
|
|
return connection
|
|
|
|
with patch.object(db.model, "get_connection", side_effect=trace_connection):
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200)
|
|
selects = [
|
|
query
|
|
for query in statements
|
|
if query.lstrip().upper().startswith(("SELECT", "WITH"))
|
|
]
|
|
self.assertEqual(len(selects), 3, selects)
|
|
self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects)
|
|
_, first = await self.get(limit=2)
|
|
statements.clear()
|
|
with patch.object(db.model, "get_connection", side_effect=trace_connection):
|
|
response, data = await self.get(limit=2, cursor=first["next_cursor"])
|
|
self.assertEqual(response.status, 200)
|
|
selects = [
|
|
query
|
|
for query in statements
|
|
if query.lstrip().upper().startswith(("SELECT", "WITH"))
|
|
]
|
|
# The page's link summary and its favorites read only.
|
|
self.assertEqual(len(selects), 2, selects)
|
|
self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects)
|
|
self.assertNotIn("gi.file_path, p.rating", selects[0])
|
|
# Without filters the listing outlives a database failure; the count
|
|
# is then unknown rather than wrong.
|
|
with patch.object(
|
|
db, "get_image_link_matches", side_effect=RuntimeError("database gone")
|
|
):
|
|
response, data = await self.get()
|
|
self.assertEqual(response.status, 200, data)
|
|
self.assertEqual(data["total"], 10)
|
|
self.assertIsNone(data["positive_count"])
|
|
self.assertNotIn("database gone", json.dumps(data))
|
|
# A linked filter or the rating sort needs the read, so its failure is
|
|
# an error (filename-only search survives it, see below).
|
|
for params in (
|
|
{"category": "Night"},
|
|
{"tags": "a"},
|
|
{"q": "zebra", "category": "Night"},
|
|
{"sort": "rating"},
|
|
):
|
|
with (
|
|
self.subTest(failing=params),
|
|
patch.object(
|
|
db, "get_image_link_matches", side_effect=RuntimeError("gone")
|
|
),
|
|
):
|
|
response, data = await self.get(**params)
|
|
self.assertEqual(response.status, 500, data)
|
|
self.assertNotIn("images", data)
|
|
|
|
async def test_filename_search_survives_a_database_failure(self):
|
|
for name, mtime in (("zebra-1.png", 300), ("zebra-2.png", 200)):
|
|
self.media(name, mtime)
|
|
self.media("other.png", 100)
|
|
with (
|
|
patch.object(
|
|
self.api.db,
|
|
"get_image_link_matches",
|
|
side_effect=RuntimeError("database gone"),
|
|
),
|
|
self.assertLogs(self.api.logger, "WARNING") as logged,
|
|
):
|
|
response, first = await self.get(q="ZEBRA", limit=1)
|
|
self.assertEqual(response.status, 200, first)
|
|
response, second = await self.get(
|
|
q="ZEBRA", limit=1, cursor=first["next_cursor"]
|
|
)
|
|
self.assertEqual(response.status, 200, second)
|
|
# Filename matches only; the count of linked positives is unknown.
|
|
self.assertEqual(self._names(first), ["zebra-1.png"])
|
|
self.assertEqual(self._names(second), ["zebra-2.png"])
|
|
for page in (first, second):
|
|
self.assertEqual(page["total"], 2)
|
|
self.assertIsNone(page["positive_count"])
|
|
self.assertNotIn("database gone", json.dumps(page))
|
|
self.assertTrue(
|
|
all(record.exc_info for record in logged.records), logged.output
|
|
)
|
|
|
|
async def test_search_matches_filename_prompt_text_and_tags(self):
|
|
from database.models import normalize_image_path
|
|
|
|
files = self._filter_fixtures()
|
|
_, data = await self.get(q="ZEBRA", sort="name")
|
|
# Text, stored title, tag, pair title and filename; never an ambiguous link.
|
|
self.assertEqual(
|
|
self._names(data),
|
|
[
|
|
"day_ab.png",
|
|
"night_ab.png",
|
|
"paired.png",
|
|
"titled.png",
|
|
"zebra_unlinked.png",
|
|
],
|
|
)
|
|
self.assertEqual(data["total"], 5)
|
|
_, data = await self.get(q="0% w")
|
|
self.assertEqual(self._names(data), ["percent.png"])
|
|
_, data = await self.get(q="%")
|
|
self.assertEqual(self._names(data), ["percent.png"])
|
|
_, data = await self.get(q="harbor")
|
|
self.assertEqual(self._names(data), ["night_a.png"])
|
|
self.assertEqual(
|
|
self.api.db.get_image_paths_matching("zebra", None, []),
|
|
{
|
|
normalize_image_path(str(files[name]))
|
|
for name in ("night_ab", "day_ab", "titled", "paired")
|
|
},
|
|
)
|
|
|
|
async def test_category_and_tags_filter_linked_images_only(self):
|
|
self._filter_fixtures()
|
|
cases = (
|
|
({"category": "Night"}, ["night_ab.png", "night_a.png", "paired.png"]),
|
|
({"category": "night"}, []),
|
|
({"tags": "a,b"}, ["night_ab.png", "day_ab.png"]),
|
|
({"tags": " a , b ,"}, ["night_ab.png", "day_ab.png"]),
|
|
({"category": "Night", "tags": "a,b"}, ["night_ab.png"]),
|
|
({"category": "Night", "q": "harbor"}, ["night_a.png"]),
|
|
({"category": "Night", "q": "paired"}, ["paired.png"]),
|
|
({"category": "Day", "q": "zebra"}, ["day_ab.png", "titled.png"]),
|
|
({"category": ""}, None),
|
|
({"category": " Night "}, ["night_ab.png", "night_a.png", "paired.png"]),
|
|
({"category": " "}, None),
|
|
)
|
|
_, everything = await self.get(sort="name")
|
|
for params, expected in cases:
|
|
with self.subTest(**params):
|
|
response, data = await self.get(sort="name", **params)
|
|
self.assertEqual(response.status, 200, data)
|
|
if expected is None:
|
|
self.assertEqual(data["images"], everything["images"])
|
|
else:
|
|
self.assertEqual(self._names(data), sorted(expected))
|
|
self.assertEqual(data["total"], len(expected))
|
|
for params in (
|
|
{"category": "x" * 201},
|
|
{"tags": ",".join(f"t{i}" for i in range(31))},
|
|
{"tags": "x" * 201},
|
|
):
|
|
response, data = await self.get(**params)
|
|
self.assertEqual(response.status, 400, params)
|
|
# The filters read one SELECT besides the page's link summary.
|
|
statements = []
|
|
original = self.api.db.model.get_connection
|
|
|
|
def trace_connection():
|
|
connection = original()
|
|
connection.set_trace_callback(statements.append)
|
|
return connection
|
|
|
|
with patch.object(
|
|
self.api.db.model, "get_connection", side_effect=trace_connection
|
|
):
|
|
response, _ = await self.get(
|
|
category="Night", tags="a", q="zebra", sort="rating"
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
selects = [
|
|
query
|
|
for query in statements
|
|
if query.lstrip().upper().startswith(("SELECT", "WITH"))
|
|
]
|
|
self.assertEqual(len(selects), 3, selects)
|
|
self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects)
|
|
|
|
async def test_rating_sort_orders_linked_then_unlinked(self):
|
|
self._filter_fixtures()
|
|
response, data = await self.get(sort="rating")
|
|
self.assertEqual(response.status, 200, data)
|
|
self.assertEqual(
|
|
self._names(data),
|
|
[
|
|
"night_ab.png", # 5
|
|
"paired.png", # 4, the pair's positive
|
|
"day_ab.png", # 3
|
|
"titled.png", # 1
|
|
# linked without a rating, newest first
|
|
"percent.png",
|
|
"night_a.png",
|
|
# unlinked (and ambiguous) last, newest first
|
|
"plain_unlinked.png",
|
|
"zebra_unlinked.png",
|
|
"ambiguous.png",
|
|
],
|
|
)
|
|
_, first = await self.get(sort="rating", limit=4)
|
|
_, rest = await self.get(sort="rating", limit=10, cursor=first["next_cursor"])
|
|
self.assertEqual(first["images"] + rest["images"], data["images"])
|
|
_, night = await self.get(sort="rating", category="Night")
|
|
self.assertEqual(
|
|
self._names(night), ["night_ab.png", "paired.png", "night_a.png"]
|
|
)
|
|
|
|
async def test_cursor_rejects_changed_filters(self):
|
|
self._filter_fixtures()
|
|
_, first = await self.get(category="Night", limit=1)
|
|
token = first["next_cursor"]
|
|
self.assertIsInstance(token, str)
|
|
response, _ = await self.get(category="Night", limit=1, cursor=token)
|
|
self.assertEqual(response.status, 200)
|
|
for params in (
|
|
{"category": "Day"},
|
|
{},
|
|
{"category": "Night", "tags": "a"},
|
|
{"category": "Night", "q": "zebra"},
|
|
{"category": "Night", "sort": "rating"},
|
|
):
|
|
with self.subTest(**params):
|
|
response, _ = await self.get(cursor=token, **params)
|
|
self.assertEqual(response.status, 400)
|
|
_, rated = await self.get(sort="rating", limit=1)
|
|
response, _ = await self.get(sort="rating", cursor=rated["next_cursor"])
|
|
self.assertEqual(response.status, 200)
|
|
response, _ = await self.get(sort="newest", cursor=rated["next_cursor"])
|
|
self.assertEqual(response.status, 400)
|
|
_, tagged = await self.get(tags="a,b", limit=1)
|
|
response, _ = await self.get(tags="b,a", cursor=tagged["next_cursor"])
|
|
self.assertEqual(response.status, 200)
|
|
response, _ = await self.get(tags="a", cursor=tagged["next_cursor"])
|
|
self.assertEqual(response.status, 400)
|
|
for params in ({"sort": "rated"}, {"sort": "rating-desc"}, {"media": "x"}):
|
|
response, _ = await self.get(**params)
|
|
self.assertEqual(response.status, 400, params)
|
|
|
|
async def test_windows_style_link_paths_compare_normalized(self):
|
|
import ntpath
|
|
|
|
def windows_normalize(path):
|
|
return ntpath.normcase(ntpath.normpath(path))
|
|
|
|
prompt = self._save_prompt("Case words", category="Night", rating=2)
|
|
listed = self.media("Sub/Night.PNG")
|
|
self.media("Sub/other.png")
|
|
stored = str(listed).upper().replace("/", "\\")
|
|
with (
|
|
patch("database.operations.normalize_image_path", windows_normalize),
|
|
patch("py.api.images.normalize_image_path", windows_normalize),
|
|
):
|
|
self.api.db.link_image_to_prompt(prompt, stored)
|
|
_, data = await self.get(category="Night")
|
|
self.assertEqual(self._names(data), ["Night.PNG"])
|
|
_, data = await self.get(q="case words")
|
|
self.assertEqual(self._names(data), ["Night.PNG"])
|
|
_, data = await self.get(sort="rating")
|
|
self.assertEqual(self._names(data), ["Night.PNG", "other.png"])
|
|
|
|
def _unlisted_link_rows(self, count):
|
|
"""*count* linked rows for files outside the listing, written in bulk."""
|
|
prompt = self._save_prompt("Elsewhere", category="Night", rating=5)
|
|
with self.api.db.model.get_connection() as conn:
|
|
conn.executemany(
|
|
"INSERT INTO generated_images(prompt_id,image_path,filename,file_path)"
|
|
" VALUES(?,?,?,?)",
|
|
(
|
|
(prompt, f"/elsewhere/{i}.png", f"{i}.png", f"/elsewhere/{i}.png")
|
|
for i in range(count)
|
|
),
|
|
)
|
|
|
|
async def test_rating_sort_reads_only_listed_paths_past_the_old_cap(self):
|
|
from database.models import normalize_image_path
|
|
|
|
files = self._filter_fixtures()
|
|
# More rows than the former 500,000-row cap, none of them listed.
|
|
await asyncio.get_running_loop().run_in_executor(
|
|
None, self._unlisted_link_rows, 500_001
|
|
)
|
|
response, data = await self.get(sort="rating", limit=3)
|
|
self.assertEqual(response.status, 200, data)
|
|
self.assertEqual(
|
|
self._names(data), ["night_ab.png", "paired.png", "day_ab.png"]
|
|
)
|
|
response, data = await self.get(category="Night", sort="name")
|
|
self.assertEqual(response.status, 200, data)
|
|
self.assertEqual(
|
|
self._names(data), ["night_a.png", "night_ab.png", "paired.png"]
|
|
)
|
|
wanted = {normalize_image_path(str(files["night_ab"])), "/elsewhere/7.png"}
|
|
elsewhere = self.api.db.get_prompt_id_for_image("/elsewhere/7.png")
|
|
self.assertEqual(
|
|
self.api.db.get_image_link_matches(
|
|
"", "Night", [], every_link=True, wanted=wanted
|
|
),
|
|
{
|
|
normalize_image_path(str(files["night_ab"])): (
|
|
True,
|
|
5,
|
|
self.api.db.get_prompt_id_for_image(str(files["night_ab"])),
|
|
),
|
|
"/elsewhere/7.png": (True, 5, elsewhere),
|
|
},
|
|
)
|
|
|
|
def _budget_fixture(self, files, links):
|
|
"""*files* listed paths and *links* linked rows, written in bulk."""
|
|
from database.models import normalize_image_path
|
|
|
|
listed = [
|
|
(self.output_dir / f"day-{i % 50:02}" / f"img_{i:05}.png", 1_000.0 + i)
|
|
for i in range(files)
|
|
]
|
|
prompts = min(2_000, links)
|
|
with self.api.db.model.get_connection() as conn:
|
|
conn.executemany(
|
|
"INSERT INTO tags(id,name) VALUES(?,?)",
|
|
[(i + 1, f"tag-{i}") for i in range(20)],
|
|
)
|
|
conn.executemany(
|
|
"INSERT INTO prompts(id,text,category,rating,hash) VALUES(?,?,?,?,?)",
|
|
[
|
|
(
|
|
i + 1,
|
|
f"prompt {i} " + "lantern " * (i % 7) + "harbor " * 20,
|
|
f"cat-{i % 10}",
|
|
(i % 6) or None,
|
|
f"hash-{i}",
|
|
)
|
|
for i in range(prompts)
|
|
],
|
|
)
|
|
conn.executemany(
|
|
"INSERT OR IGNORE INTO prompt_tags(prompt_id,tag_id) VALUES(?,?)",
|
|
[(i + 1, t) for i in range(prompts) for t in (i % 20 + 1, 1)],
|
|
)
|
|
conn.executemany(
|
|
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
|
[(i + 1, f"Title {i}") for i in range(0, prompts, 3)],
|
|
)
|
|
conn.executemany(
|
|
"INSERT INTO generated_images(prompt_id,image_path,filename,file_path)"
|
|
" VALUES(?,?,?,?)",
|
|
[
|
|
(
|
|
i % prompts + 1,
|
|
str(listed[i * 2][0]),
|
|
listed[i * 2][0].name,
|
|
normalize_image_path(str(listed[i * 2][0])),
|
|
)
|
|
for i in range(links)
|
|
],
|
|
)
|
|
return listed
|
|
|
|
_BUDGET_REQUESTS = (
|
|
# First, so the timed loop also measures a cursor page of each.
|
|
{},
|
|
{"sort": "oldest"},
|
|
{"q": "lantern"},
|
|
{"q": "img_0123"},
|
|
{"q": "Title 1"},
|
|
{"category": "cat-3"},
|
|
{"tags": "tag-0,tag-2"},
|
|
{"sort": "rating"},
|
|
{"sort": "rating", "category": "cat-4"},
|
|
{"q": "harbor", "category": "cat-5", "tags": "tag-5"},
|
|
{"sort": "name", "q": "tag-7"},
|
|
{"subfolder": "day-07", "q": "lantern"},
|
|
)
|
|
|
|
@staticmethod
|
|
def _stub_page(page):
|
|
return [{"filename": entry[0].name} for entry in page]
|
|
|
|
async def test_filter_work_is_bounded(self):
|
|
"""The budget's structure, counted rather than timed (runs everywhere)."""
|
|
import math
|
|
|
|
from py.api import images as images_module
|
|
from database.operations import link_filter_query
|
|
|
|
listed = await asyncio.get_running_loop().run_in_executor(
|
|
None, self._budget_fixture, 5_000, 2_000
|
|
)
|
|
scans = [listed]
|
|
built, keyed = [], []
|
|
real_listed, real_key = images_module._listed, images_module._workspace_key
|
|
|
|
def counting_listed(*args):
|
|
built.append(args[0])
|
|
return real_listed(*args)
|
|
|
|
def counting_key(entry, order):
|
|
keyed.append(order)
|
|
return real_key(entry, order)
|
|
|
|
statements = []
|
|
original = self.api.db.model.get_connection
|
|
|
|
def trace_connection():
|
|
connection = original()
|
|
connection.set_trace_callback(statements.append)
|
|
return connection
|
|
|
|
async def scan(_root):
|
|
return scans[-1]
|
|
|
|
def selects():
|
|
found = [
|
|
query
|
|
for query in statements
|
|
if query.lstrip().upper().startswith(("SELECT", "WITH"))
|
|
]
|
|
statements.clear()
|
|
return found
|
|
|
|
with (
|
|
patch.object(self.api, "_get_gallery_files", side_effect=scan),
|
|
patch.object(self.api, "_format_workspace_page", self._stub_page),
|
|
patch.object(images_module, "_listed", counting_listed),
|
|
patch.object(images_module, "_workspace_key", counting_key),
|
|
patch.object(
|
|
self.api.db.model, "get_connection", side_effect=trace_connection
|
|
),
|
|
):
|
|
response, _ = await self.get(limit=10)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(len(built), len(listed))
|
|
# No filter still reads the linked positives once, for the count.
|
|
self.assertEqual(len(selects()), 1)
|
|
newest = self.api._workspace_order_cache[((0, str(self.output_dir)),)]
|
|
newest = newest[1]["newest"]
|
|
for params in self._BUDGET_REQUESTS:
|
|
with self.subTest(**params):
|
|
built.clear()
|
|
response, data = await self.get(limit=10, **params)
|
|
self.assertEqual(response.status, 200, data)
|
|
self.assertEqual(built, [])
|
|
self.assertEqual(len(selects()), 1)
|
|
self.assertIs(
|
|
self.api._workspace_order_cache[((0, str(self.output_dir)),)][1][
|
|
"newest"
|
|
],
|
|
newest,
|
|
)
|
|
bound = math.ceil(math.log2(len(listed))) + 3
|
|
for order in ("newest", "name", "name-desc", "oldest", "rating"):
|
|
with self.subTest(cursor=order):
|
|
await self.get(sort=order, limit=10) # builds a lazy order
|
|
_, first = await self.get(sort=order, limit=10)
|
|
keyed.clear()
|
|
selects()
|
|
response, page = await self.get(
|
|
sort=order, limit=10, cursor=first["next_cursor"]
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(len(page["images"]), 10)
|
|
self.assertLessEqual(len(keyed), bound)
|
|
# A plain cursor page reads no links; rating needs them.
|
|
plain = order != "rating"
|
|
self.assertEqual(len(selects()), 0 if plain else 1)
|
|
self.assertEqual("positive_count" in page, not plain)
|
|
key = ((0, str(self.output_dir)),)
|
|
# A TTL rescan that found the same files rebuilds nothing.
|
|
built.clear()
|
|
scans.append(list(listed))
|
|
response, data = await self.get(limit=10)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(built, [])
|
|
self.assertIs(self.api._workspace_order_cache[key][1]["newest"], newest)
|
|
# One changed file is the only one rebuilt.
|
|
added = (self.output_dir / "day-00" / "new.png", 9_999_999.0)
|
|
scans.append([added, *listed[1:]])
|
|
response, data = await self.get(limit=10)
|
|
self.assertEqual(built, [added[0]])
|
|
self.assertEqual(data["total"], len(listed))
|
|
self.assertEqual(data["images"][0]["filename"], "new.png")
|
|
# The monitor invalidates the gallery when it sees a new file: the
|
|
# next request rebuilds only that file's values, and search finds it.
|
|
self.api.invalidate_gallery_cache()
|
|
self.assertNotEqual(self.api._workspace_listed_cache, {})
|
|
monitored = (self.output_dir / "day-01" / "monitored.png", 9_999_999.5)
|
|
scans.append([monitored, *scans[-1]])
|
|
built.clear()
|
|
response, data = await self.get(limit=10, q="monitored")
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(built, [monitored[0]])
|
|
self.assertEqual(self._names(data), ["monitored.png"])
|
|
# A roots change drops every derived value.
|
|
with patch("py.config.GalleryConfig.MONITORING_DIRECTORIES", []):
|
|
self.assertTrue(self.api._apply_gallery_roots([str(self.output_dir)]))
|
|
self.assertEqual(self.api._workspace_listed_cache, {})
|
|
self.assertEqual(self.api._workspace_order_cache, {})
|
|
# The per-row ambiguity probe is an index search, never a scan.
|
|
sql, params = link_filter_query("lantern", "cat-3", ["tag-1"], every_link=True)
|
|
with self.api.db.model.get_connection() as conn:
|
|
plan = [row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, params)]
|
|
self.assertIn(
|
|
"SEARCH other USING COVERING INDEX idx_image_file_path (file_path=?)", plan
|
|
)
|
|
self.assertFalse([step for step in plan if step.startswith("SCAN other")])
|
|
|
|
@unittest.skipUnless(
|
|
PERF_BUDGET and not TRACED,
|
|
"wall-clock budget: set PROMPTMANAGER_PERF_BUDGET=1 and run untraced",
|
|
)
|
|
async def test_filter_budget(self):
|
|
"""p95 < 150 ms at 50,000 files / 20,000 links (TESTING.md budgets)."""
|
|
import math
|
|
import time
|
|
|
|
# Built off the event loop: the bulk inserts take a few hundred ms.
|
|
listed = await asyncio.get_running_loop().run_in_executor(
|
|
None, self._budget_fixture, 50_000, 20_000
|
|
)
|
|
rounds, plain = [], {"first": [], "cursor": []}
|
|
with (
|
|
patch.object(
|
|
self.api, "_get_gallery_files", AsyncMock(return_value=listed)
|
|
),
|
|
patch.object(self.api, "_format_workspace_page", self._stub_page),
|
|
):
|
|
# The first request after a scan builds the cached sort orders.
|
|
started = time.perf_counter()
|
|
response, _ = await self.get(sort="rating")
|
|
cold = time.perf_counter() - started
|
|
self.assertEqual(response.status, 200)
|
|
for _round in range(3):
|
|
timings = []
|
|
for index in range(20):
|
|
requests = self._BUDGET_REQUESTS
|
|
shape = requests[index % len(requests)]
|
|
params = dict(shape, limit=100)
|
|
if index >= len(requests):
|
|
_, first = await self.get(**params)
|
|
if first["next_cursor"]:
|
|
params["cursor"] = first["next_cursor"]
|
|
started = time.perf_counter()
|
|
response, data = await self.get(**params)
|
|
timings.append(time.perf_counter() - started)
|
|
if shape.get("sort", "newest") in ("newest", "oldest") and not (
|
|
set(shape) - {"sort"}
|
|
):
|
|
# Plain first pages read the links for the count;
|
|
# plain cursor pages read none.
|
|
plain["cursor" if "cursor" in params else "first"].append(
|
|
timings[-1]
|
|
)
|
|
self.assertEqual(response.status, 200, (params, data))
|
|
timings.sort()
|
|
rounds.append(timings)
|
|
p95s = [timings[math.ceil(0.95 * len(timings)) - 1] for timings in rounds]
|
|
best = min(p95s)
|
|
plain_ms = {
|
|
kind: sorted(round(t * 1000, 1) for t in found)
|
|
for kind, found in plain.items()
|
|
}
|
|
print(
|
|
"\nworkspace filter budget: best p95="
|
|
f"{best * 1000:.1f} ms, rounds="
|
|
f"{[round(p * 1000, 1) for p in p95s]} ms, cold={cold * 1000:.1f} ms, "
|
|
f"plain listing ms={plain_ms}",
|
|
file=sys.stderr,
|
|
)
|
|
self.assertLess(
|
|
best, 0.150, [[round(t * 1000, 1) for t in timings] for timings in rounds]
|
|
)
|
|
|
|
|
|
class TestWorkspaceCursorEncoding(unittest.TestCase):
|
|
def test_ntfs_unicode_basename_cursor_is_bounded_and_round_trips(self):
|
|
from py.api.images import (
|
|
_workspace_cursor,
|
|
_workspace_cursor_key,
|
|
_workspace_key,
|
|
)
|
|
|
|
root = Path("gallery")
|
|
entry = (root / ("\u0390" * 240 + ".png"), 100, root, 0)
|
|
for order in ("newest", "oldest", "name", "name-desc"):
|
|
with self.subTest(sort=order):
|
|
key = _workspace_key(entry, order)
|
|
token = _workspace_cursor(key, "scope")
|
|
self.assertLessEqual(len(token), 4096)
|
|
self.assertEqual(_workspace_cursor_key(token, "scope", order), key)
|