Vito
c7ba481d3c
* perf(images): refresh the listing when the monitor sees a new output (#348, PR 2) The image monitor now reports every settled output (linked or not, also when linking raised) with its resolved path through set_output_listener. The API's notify_output_settled hops from the monitor's worker thread onto the server's loop (call_soon_threadsafe) and does nothing before the loop is known or after cleanup. mark_root_dirty matches the path against each walk's resolved_root (normcase(realpath), computed in the executor by walk_root, now its fourth value) and marks the root dirty. A dirty root gets one background walk, no sooner than rewalk_interval = min(15, max(5, 5 x last walk's duration)) after the previous walk started (last_started_at / last_duration are recorded on every walk). _schedule_walk is shared with PR 1's delayed stale refresh: one deferred walk per root, never armed while a walk runs (the walk's end re-checks dirty: the trailing walk), at least 0.5 s away, the earlier deadline wins, and every walk start cancels it. A failed walk clears dirty. Rescan semantics are unchanged. * fix(images): keep removed roots and cancelled walks from walking again (#348, PR 2) Review findings on the monitor-triggered walks: - reset_gallery_roots_cache now retires every listing record: its pending timer is cancelled, dirty is cleared, and no mark, timer or trailing walk starts a walk of it until a reader or the warm-up asks for that root again. Before, a timer or trailing walk pending when the roots changed stored a listing for the removed root, and the monitor (which watches the old roots until restart) kept it walking on every new output. - A walk task cancelled outside cleanup_gallery_listings no longer starts a trailing walk while the loop shuts down. - root_contains ignores a Windows extended-length prefix (\\?\ and \\?\UNC\) on either side, which realpath can keep on a file past MAX_PATH. - The REWALK comment states the ~20% GIL share only for walks up to 3 s. * test(images): prove a cancelled walk starts no trailing walk (#348, PR 2) The previous assertions held even with the trailing walk, because that walk finished before they ran; assert the walk number and the scan count. * perf(images): add a settled output to the listing without a walk (#348, PR 2) Live test: each landing in a 20,000-file folder cost a full walk (~0.4 s of GIL) plus a full workspace rebuild, so Images reads during a run took up to 2.1 s. The monitor's report now patches the stored listing instead. - output_settled (the monitor hop's target) slices the file's part below the root in its own case (relative_part; Windows compares folded), and _run_adds places one root's pending files in one executor batch: _settled_entry applies the walk's filters (suffix, thumbnails, depth, symlink by lstat; macOS on-disk spelling), insert_entry puts the (Path, mtime) entry where the newest-first sort does, add_folder copies the folder sets. Stored only when no walk started and no drop happened meanwhile; walked_at is untouched, so the 30 s refresh still walks. - Anything inexact marks the root dirty and walks as before: a running walk, an unsliceable path, a failed lstat, a tie in mtime, the cap, an unexpected error. Rescan, the refresh, the warm-up and invalidation stay full walks; a walk start forgets the pending adds. - The workspace caches follow a patched listing by entry identity (listing_delta): only added values are built, the newest order places them by bisection, the favorites lookup is copied and extended. Synthetic 20,000 files: per landing 252 ms (walk + full build) -> 18 ms (3 ms add + 15 ms warm build). * fix(images): review fixes for adds without a walk (#348, PR 2) - _newest_step steps the cached newest order only when the per-file cache still holds exactly the values this call built. A stale page read racing a newer warm build (both run in executor threads) stored an order labelled with older values than it held, and the next read inserted the new file a second time. It also sorts in full when a roots reset dropped the per-file values meanwhile. - The stored spelling is checked on Windows too (nt._getfinalpathname: realpath keeps the given case when it cannot open the file), and before the name filters, which compare the walk's stored spelling. macOS opens non-blocking, so a FIFO named like an image never holds the executor. - An output in a folder the walk cannot list (no read permission) is not added, as the walk skips it. - Windows-portable unit tests: POSIX paths in the insert tests, and a str.lower fold for the length guard (Windows' normcase keeps lengths). * test(images): reproduce the newest-order race and fix a Windows clock flake (#348, PR 2) - The race test now has the newer build step from the values the cached order was built from, which is the interleaving that inserted a file twice; it fails without the parts check in _newest_step. - test_a_failed_walk_records_its_timing_and_clears_dirty compared two monotonic readings that Windows' ~15.6 ms clock can make equal; the previous start is moved back instead. - TESTING.md: os.access ignores ACLs on Windows.
2171 lines
95 KiB
Python
2171 lines
95 KiB
Python
"""Opt-in workspace gallery contracts against disposable SQLite and media."""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import shutil
|
|
import sqlite3
|
|
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_an_unchanged_refresh_keeps_the_listing_and_derived_caches(self):
|
|
from py.api import gallery_listing
|
|
|
|
self.media("a.png", 300)
|
|
self.media("b.png", 200)
|
|
_, first = await self.get()
|
|
response = await self.client.put(
|
|
"/prompt_manager/images/favorite",
|
|
json={"root_index": 0, "path": "a.png"},
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
record = self.api._root_listings[str(self.output_dir)]
|
|
files = record.files
|
|
key = ((0, str(self.output_dir)),)
|
|
newest = self.api._workspace_order_cache[key][1]["newest"]
|
|
favorites = self.api._favorite_listing_cache[str(self.output_dir)]
|
|
record.walked_at -= gallery_listing.REFRESH_AFTER + 1
|
|
record.last_started_at -= gallery_listing.REFRESH_AFTER + 1
|
|
with patch.object(
|
|
self.api,
|
|
"_scan_gallery_files_sync",
|
|
wraps=self.api._scan_gallery_files_sync,
|
|
) as scan:
|
|
with patch.object(gallery_listing, "REFRESH_DELAY", 0):
|
|
await self.get() # served at once; one refresh walk follows
|
|
while record.rewalk_handle is not None:
|
|
await asyncio.sleep(0.01)
|
|
while self.api._listing_tasks:
|
|
await asyncio.gather(*list(self.api._listing_tasks))
|
|
# A finished walk leaves the set in a done callback, which needs
|
|
# a loop turn; gather over finished tasks does not yield one.
|
|
await asyncio.sleep(0)
|
|
_, again = await self.get()
|
|
self.assertEqual(scan.call_count, 1)
|
|
self.assertIs(record.files, files)
|
|
self.assertIs(self.api._workspace_order_cache[key][1]["newest"], newest)
|
|
self.assertIs(self.api._favorite_listing_cache[str(self.output_dir)], favorites)
|
|
self.assertEqual(again["total"], first["total"])
|
|
|
|
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_path_filter_lists_only_that_exact_file(self):
|
|
# Activity (D57): a scan row opens its file by root and exact relative path,
|
|
# never by a substring that also matches the same name in other folders.
|
|
self.media("ComfyUI_00001_.png")
|
|
for day in range(3):
|
|
self.media(f"2026-10-0{day + 1}/ComfyUI_00001_.png")
|
|
self.media("2026-10-01/ComfyUI_00001_.png.bak.png")
|
|
_, data = await self.get(path="ComfyUI_00001_.png")
|
|
self.assertEqual([i["path"] for i in data["images"]], ["ComfyUI_00001_.png"])
|
|
_, data = await self.get(path="2026-10-02/ComfyUI_00001_.png")
|
|
self.assertEqual(
|
|
[i["path"] for i in data["images"]], ["2026-10-02/ComfyUI_00001_.png"]
|
|
)
|
|
_, data = await self.get(path="missing.png")
|
|
self.assertEqual(data["images"], [])
|
|
for bad in ("../outside.png", "a" * 1025):
|
|
with self.subTest(bad=bad[:12]):
|
|
response, _ = await self.get(path=bad)
|
|
self.assertEqual(response.status, 400)
|
|
|
|
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_roots_list_their_own_folders(self):
|
|
# "Refine this view" offers only the chosen Source's folders (#249 m1).
|
|
second = self.comfy_dir / "second"
|
|
second.mkdir()
|
|
self.media("first/deep/a.png")
|
|
self.media("nested/b.png", root=second)
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
|
|
deep = os.path.join("first", "deep")
|
|
response = await self.client.get(
|
|
"/prompt_manager/gallery/subfolders?mode=workspace&include_ancestors=true"
|
|
)
|
|
folders = await response.json()
|
|
self.assertEqual(folders["subfolders"], sorted(["first", deep, "nested"]))
|
|
self.assertEqual(
|
|
[(r["index"], r["subfolders"]) for r in folders["roots"]],
|
|
[(0, sorted(["first", deep])), (1, ["nested"])],
|
|
)
|
|
# A chosen root narrows the union; every root still lists its own.
|
|
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["subfolders"] for r in folders["roots"]], [[deep], ["nested"]]
|
|
)
|
|
|
|
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, variant-titled 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)
|
|
# A pair's own title names its negative variant; the card shows the
|
|
# linked positive's title (D10/D20), never the variant name.
|
|
pair = db.save_prompt_pair("Pair positive", "Noise", title="Pair title")
|
|
with db.model.get_connection() as connection:
|
|
connection.execute(
|
|
"INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)",
|
|
(db.get_prompt_pair(pair)["positive_prompt_id"], "Pair positive title"),
|
|
)
|
|
variant = db.save_prompt_pair("Variant positive", "Grain", title="Variant name")
|
|
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", "variant", "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), ("variant", variant), ("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 positive title",
|
|
),
|
|
# No stored title and no file evidence: the card shows the filename.
|
|
"variant": (
|
|
"linked",
|
|
db.get_prompt_pair(variant)["positive_prompt_id"],
|
|
None,
|
|
),
|
|
"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"]
|
|
|
|
def test_link_summary_reads_no_pair_title(self):
|
|
# The pair's own title names its negative variant and never titles a
|
|
# card (D10), so the summary query does not read it (#249 m4).
|
|
from database.operations import _LINK_SUMMARY_SQL
|
|
|
|
self.assertNotIn("pp.title", _LINK_SUMMARY_SQL)
|
|
|
|
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_bulk_category_change_shows_on_the_next_listing(self):
|
|
files = self._filter_fixtures()
|
|
db = self.api.db
|
|
night_a = db.get_prompt_id_for_image(str(files["night_a"]))
|
|
# Warm every listing cache (scan, orders, card facts, link reads).
|
|
for params in ({}, {"category": "Landscapes"}, {"category": "Night"}):
|
|
response, data = await self.get(**params)
|
|
self.assertEqual(response.status, 200, data)
|
|
self.assertEqual(data["total"], 3)
|
|
response = await self.client.post(
|
|
"/prompt_manager/bulk/category",
|
|
json={"prompt_ids": [night_a], "category": "Landscapes"},
|
|
)
|
|
self.assertEqual(response.status, 200, await response.json())
|
|
_, moved = await self.get(category="Landscapes")
|
|
self.assertEqual(self._names(moved), ["night_a.png"])
|
|
self.assertEqual(moved["images"][0]["linked"]["category"], "Landscapes")
|
|
_, night = await self.get(category="Night")
|
|
self.assertEqual(night["total"], 2)
|
|
self.assertNotIn("night_a.png", self._names(night))
|
|
|
|
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, stored title, tag and the filename match (never a pair's
|
|
# own variant title, D12).
|
|
({"q": "ZEBRA"}, 5, 3),
|
|
({"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"},
|
|
):
|
|
leak = 'near "gi": syntax error FROM generation_images'
|
|
with (
|
|
self.subTest(failing=params),
|
|
patch.object(
|
|
db, "get_image_link_matches", side_effect=RuntimeError(leak)
|
|
),
|
|
self.assertLogs("prompt_manager", level="ERROR") as logged,
|
|
):
|
|
response, data = await self.get(**params)
|
|
self.assertEqual(response.status, 500, data)
|
|
self.assertNotIn("images", data)
|
|
self.assertEqual(data["code"], "listing-failed")
|
|
self.assertNotIn(leak, json.dumps(data))
|
|
self.assertNotIn("syntax error", json.dumps(data))
|
|
self.assertTrue(any(r.exc_info for r in logged.records))
|
|
|
|
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 and filename; never an ambiguous link.
|
|
self.assertEqual(
|
|
self._names(data),
|
|
[
|
|
"day_ab.png",
|
|
"night_ab.png",
|
|
"titled.png",
|
|
"zebra_unlinked.png",
|
|
],
|
|
)
|
|
self.assertEqual(data["total"], 4)
|
|
# D12: a pair's own title names its negative variant, so the variant
|
|
# name alone matches nothing; the linked positive's text still does.
|
|
_, data = await self.get(q="Zebra pair")
|
|
self.assertEqual(self._names(data), [])
|
|
self.assertEqual(data["total"], 0)
|
|
_, data = await self.get(q="pair positive")
|
|
self.assertEqual(self._names(data), ["paired.png"])
|
|
_, 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")
|
|
},
|
|
)
|
|
|
|
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"))
|
|
]
|
|
# Plus the one stream of tagged image records q and tags read (B3).
|
|
self.assertEqual(len(selects), 4, selects)
|
|
self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects)
|
|
self.assertEqual(sum("media_image_tags" 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, [])
|
|
found = selects()
|
|
# q and tags also stream the tagged image records once
|
|
# (#249 B3); none are tagged here, so nothing is stat'ed.
|
|
streamed = "q" in params or "tags" in params
|
|
self.assertEqual(len(found), 2 if streamed else 1, found)
|
|
self.assertEqual(
|
|
sum("media_image_tags" in query for query in found),
|
|
int(streamed),
|
|
)
|
|
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)
|
|
|
|
|
|
class TestWorkspaceImageTags(ImageAPITestCase):
|
|
"""B3: image tags in the Tags filter and search (spec amendment S1)."""
|
|
|
|
media = TestWorkspaceImages.media
|
|
get = TestWorkspaceImages.get
|
|
|
|
async def tearDownAsync(self):
|
|
self.api.db.close_all()
|
|
assert_closed(self, self._temp_db.name, self._output_tmp.name)
|
|
await super().tearDownAsync()
|
|
|
|
def tag(self, path, tags):
|
|
"""Save *tags* for the current bytes of *path*, as the Tag dialog does."""
|
|
import hashlib
|
|
|
|
stat = path.stat()
|
|
return self.api.db.replace_image_tags(
|
|
locator=os.path.normcase(os.path.realpath(path)),
|
|
sha256=hashlib.sha256(path.read_bytes()).hexdigest(),
|
|
size=stat.st_size,
|
|
mtime_ns=stat.st_mtime_ns,
|
|
expected_epoch=self.api.db.get_execution_epoch(),
|
|
expected_tag_revision=0,
|
|
tags=tags,
|
|
)
|
|
|
|
def names(self, data):
|
|
return [item["path"] for item in data["images"]]
|
|
|
|
async def test_tags_filter_admits_unlinked_image_with_current_image_tags(self):
|
|
tagged = self.media("unlinked_tagged.png", mtime=103)
|
|
self.media("plain.png", mtime=102)
|
|
linked = self.media("linked.png", mtime=101)
|
|
self.api.db.link_image_to_prompt(
|
|
self._save_prompt("Linked", tags=["other"]), str(linked)
|
|
)
|
|
self.tag(tagged, ["dunes", "warm_light"])
|
|
for tags in ("dunes", "dunes,warm_light", "warm_light"):
|
|
with self.subTest(tags=tags):
|
|
response, body = await self.get(tags=tags)
|
|
self.assertEqual(response.status, 200, body)
|
|
self.assertEqual(
|
|
[i["path"] for i in body["images"]], ["unlinked_tagged.png"]
|
|
)
|
|
self.assertEqual(body["total"], 1)
|
|
self.assertEqual(body["positive_count"], 0)
|
|
# Exact names, as for prompt tags; a missing tag excludes the file.
|
|
for tags in ("Dunes", "dunes,missing", "dune"):
|
|
with self.subTest(tags=tags):
|
|
_, body = await self.get(tags=tags)
|
|
self.assertEqual(self.names(body), [])
|
|
# A category filter reads only the linked positive.
|
|
_, body = await self.get(tags="dunes", category="Night")
|
|
self.assertEqual(self.names(body), [])
|
|
|
|
async def test_tags_filter_combines_image_and_positive_tags(self):
|
|
both = self._save_prompt("Both", category="Night", tags=["a", "b"], rating=2)
|
|
only_a = self._save_prompt("Only a", category="Night", tags=["a"], rating=5)
|
|
day_a = self._save_prompt("Day a", category="Day", tags=["a"], rating=4)
|
|
files = {
|
|
name: self.media(f"{name}.png", mtime=100 + index)
|
|
for index, name in enumerate(
|
|
("positive_ab", "positive_a_image_b", "image_ab", "image_a", "day")
|
|
)
|
|
}
|
|
self.api.db.link_image_to_prompt(both, str(files["positive_ab"]))
|
|
self.api.db.link_image_to_prompt(only_a, str(files["positive_a_image_b"]))
|
|
self.api.db.link_image_to_prompt(day_a, str(files["day"]))
|
|
self.tag(files["positive_a_image_b"], ["b"])
|
|
self.tag(files["image_ab"], ["a", "b"])
|
|
self.tag(files["image_a"], ["a"])
|
|
self.tag(files["day"], ["b"])
|
|
_, body = await self.get(tags="a,b", sort="name")
|
|
self.assertEqual(
|
|
self.names(body),
|
|
["day.png", "image_ab.png", "positive_a_image_b.png", "positive_ab.png"],
|
|
)
|
|
self.assertEqual(body["total"], 4)
|
|
self.assertEqual(body["positive_count"], 3)
|
|
_, body = await self.get(tags="a,b", category="Night", sort="name")
|
|
self.assertEqual(
|
|
self.names(body), ["positive_a_image_b.png", "positive_ab.png"]
|
|
)
|
|
# The rating sort ranks an image-admitted linked file by its positive.
|
|
_, body = await self.get(tags="a,b", sort="rating")
|
|
self.assertEqual(
|
|
self.names(body),
|
|
["positive_a_image_b.png", "day.png", "positive_ab.png", "image_ab.png"],
|
|
)
|
|
# With q, a linked positive hit admits as before.
|
|
_, body = await self.get(tags="b", q="only", sort="name")
|
|
self.assertEqual(self.names(body), ["positive_a_image_b.png"])
|
|
|
|
async def test_q_matches_image_tag_names(self):
|
|
tagged = self.media("unlinked_tagged.png", mtime=102)
|
|
self.media("plain.png", mtime=101)
|
|
self.tag(tagged, ["Sand Dunes", "warm_light"])
|
|
for query in ("dune", "DUNES", "sand d", "warm_"):
|
|
with self.subTest(q=query):
|
|
response, body = await self.get(q=query)
|
|
self.assertEqual(response.status, 200, body)
|
|
self.assertEqual(self.names(body), ["unlinked_tagged.png"])
|
|
_, body = await self.get(q="plain")
|
|
self.assertEqual(self.names(body), ["plain.png"])
|
|
_, body = await self.get(q="nothing")
|
|
self.assertEqual(self.names(body), [])
|
|
|
|
async def test_filter_ignores_tags_of_a_changed_fingerprint(self):
|
|
moved = self.media("touched.png", mtime=102)
|
|
grown = self.media("grown.png", mtime=101)
|
|
kept = self.media("kept.png", mtime=100)
|
|
for path in (moved, grown, kept):
|
|
self.tag(path, ["dunes"])
|
|
_, body = await self.get(tags="dunes")
|
|
self.assertEqual(self.names(body), ["touched.png", "grown.png", "kept.png"])
|
|
os.utime(moved, (500, 500))
|
|
with grown.open("ab") as handle:
|
|
handle.write(b"\0")
|
|
os.utime(grown, (101, 101))
|
|
self.api.invalidate_gallery_cache()
|
|
_, body = await self.get(tags="dunes")
|
|
self.assertEqual(self.names(body), ["kept.png"])
|
|
_, body = await self.get(q="dunes")
|
|
self.assertEqual(self.names(body), ["kept.png"])
|
|
|
|
async def test_stat_calls_equal_tagged_candidates_not_listing_size(self):
|
|
from py.api import images as images_module
|
|
|
|
listed = [
|
|
(self.output_dir / f"day-{i % 50:02}" / f"img_{i:05}.png", 1_000.0 + i)
|
|
for i in range(1_997)
|
|
]
|
|
tagged = [self.media(f"tagged_{i}.png", mtime=5_000 + i) for i in range(3)]
|
|
listed += [(path, path.stat().st_mtime) for path in tagged]
|
|
# Tagged but outside the listing, and listed with a non-matching tag.
|
|
outside = self.media("elsewhere.png", root=self.comfy_dir)
|
|
unrelated = self.media("unrelated.png", mtime=4_000)
|
|
listed.append((unrelated, unrelated.stat().st_mtime))
|
|
for path in tagged:
|
|
self.tag(path, ["dunes"])
|
|
self.tag(outside, ["dunes"])
|
|
self.tag(unrelated, ["other"])
|
|
self.assertEqual(len(listed), 2_001)
|
|
stat_calls = 0
|
|
real = getattr(images_module, "_current_fingerprint", None)
|
|
self.assertIsNotNone(real, "the candidate stat helper is not implemented")
|
|
|
|
def counting(path):
|
|
nonlocal stat_calls
|
|
stat_calls += 1
|
|
return real(path)
|
|
|
|
with (
|
|
patch.object(
|
|
self.api, "_get_gallery_files", AsyncMock(return_value=listed)
|
|
),
|
|
patch.object(
|
|
self.api, "_format_workspace_page", TestWorkspaceImages._stub_page
|
|
),
|
|
patch.object(images_module, "_current_fingerprint", counting),
|
|
):
|
|
response, body = await self.get(tags="dunes", limit=10)
|
|
self.assertEqual(response.status, 200, body)
|
|
self.assertEqual(
|
|
sorted(i["filename"] for i in body["images"]),
|
|
["tagged_0.png", "tagged_1.png", "tagged_2.png"],
|
|
)
|
|
self.assertEqual(stat_calls, 3)
|
|
|
|
async def test_plain_listing_reads_no_image_tags(self):
|
|
prompt = self._save_prompt("Linked", category="Night", tags=["a"])
|
|
for index in range(4):
|
|
path = self.media(f"{index}.png", mtime=100 + index)
|
|
self.tag(path, ["dunes"])
|
|
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
|
|
):
|
|
_, first = await self.get(limit=2)
|
|
_, rest = await self.get(limit=2, cursor=first["next_cursor"])
|
|
for params in (
|
|
{"sort": "rating"},
|
|
{"category": "Night"},
|
|
{"media": "image"},
|
|
):
|
|
response, _ = await self.get(**params)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(len(first["images"] + rest["images"]), 4)
|
|
self.assertFalse(
|
|
[q for q in statements if "media_" in q or "annotation" in q], statements
|
|
)
|
|
|
|
async def test_image_tag_read_failure_fails_tags_and_degrades_search(self):
|
|
tagged = self.media("unlinked_tagged.png", mtime=102)
|
|
self.media("dunes_by_name.png", mtime=101)
|
|
self.tag(tagged, ["dunes"])
|
|
leak = "no such table: media_annotation_files"
|
|
with patch.object(
|
|
self.api.db,
|
|
"iter_image_tag_records",
|
|
side_effect=sqlite3.OperationalError(leak),
|
|
):
|
|
# As a failed link read does under a Tags filter: a fixed coded
|
|
# sentence, never the exception or SQL text; the traceback is logged.
|
|
with self.assertLogs("prompt_manager", level="ERROR") as failed:
|
|
response, body = await self.get(tags="dunes")
|
|
self.assertEqual((response.status, body["success"]), (500, False))
|
|
self.assertEqual(body["code"], "listing-failed")
|
|
self.assertNotIn(leak, body["error"])
|
|
self.assertNotIn("media_annotation_files", json.dumps(body))
|
|
self.assertTrue(any(r.exc_info for r in failed.records))
|
|
with self.assertLogs("prompt_manager", level="WARNING") as logged:
|
|
response, body = await self.get(q="dunes")
|
|
self.assertEqual(response.status, 200, body)
|
|
self.assertEqual(self.names(body), ["dunes_by_name.png"])
|
|
self.assertIn("image tags", "\n".join(logged.output))
|
|
|
|
async def test_legacy_output_failure_logs_its_traceback(self):
|
|
self.media("legacy.png", mtime=100)
|
|
with (
|
|
patch.object(
|
|
self.api, "_collect_output_images", side_effect=RuntimeError("x")
|
|
),
|
|
self.assertLogs("prompt_manager", level="ERROR") as logged,
|
|
):
|
|
response = await self.client.get("/prompt_manager/images/output")
|
|
self.assertEqual(response.status, 500)
|
|
self.assertTrue(any(r.exc_info for r in logged.records), logged.output)
|
|
|
|
async def test_legacy_gallery_and_search_shapes_unchanged(self):
|
|
path = self.media("legacy.png", mtime=100)
|
|
self.media("other.png", mtime=99)
|
|
self._link_image(path, "dunes at dusk")
|
|
|
|
async def legacy():
|
|
bodies = []
|
|
for url, params in (
|
|
("/prompt_manager/images/output", {}),
|
|
("/prompt_manager/images/search", {"q": "dunes"}),
|
|
("/prompt_manager/images/all", {}),
|
|
):
|
|
response = await self.client.get(url, params=params)
|
|
self.assertEqual(response.status, 200, url)
|
|
bodies.append(await response.json())
|
|
return bodies
|
|
|
|
before = await legacy()
|
|
self.tag(path, ["dunes", "warm_light"])
|
|
self.tag(self.output_dir / "other.png", ["dunes"])
|
|
self.api.invalidate_gallery_cache()
|
|
self.assertEqual(await legacy(), before)
|
|
|
|
def test_link_facts_resolve_pairs_skip_ambiguous_and_chunk(self):
|
|
from database.models import normalize_image_path
|
|
|
|
db = self.api.db
|
|
plain = self._save_prompt("Plain dunes", category="Day", tags=["a", "b"])
|
|
positive = self._save_prompt("Pair positive", category="Night", tags=["p"])
|
|
pair = db.save_prompt_pair("Pair positive", "Noise", title="Pair")
|
|
owners = (self._save_prompt("Owner one"), self._save_prompt("Owner two"))
|
|
paths = {
|
|
name: str(self.output_dir / f"{name}.png")
|
|
for name in ("plain", "paired", "ambiguous")
|
|
}
|
|
db.link_image_to_prompt(plain, paths["plain"])
|
|
image = db.link_image_to_prompt(positive, paths["paired"])
|
|
db.link_image_to_pair(pair, image, {"source": "fixture"})
|
|
for owner in owners:
|
|
db.link_image_to_prompt(owner, paths["ambiguous"])
|
|
keys = {name: normalize_image_path(path) for name, path in paths.items()}
|
|
padding = [f"/nowhere/{index}.png" for index in range(600)]
|
|
facts = db.get_image_link_facts([*padding, *keys.values(), keys["plain"]])
|
|
self.assertEqual(
|
|
facts,
|
|
{
|
|
keys["plain"]: (True, None, plain, "Day", ("a", "b")),
|
|
keys["paired"]: (True, None, positive, "Night", ("p",)),
|
|
},
|
|
)
|
|
searched = db.get_image_link_facts(list(keys.values()), "dunes")
|
|
self.assertEqual(
|
|
{path: row[0] for path, row in searched.items()},
|
|
{keys["plain"]: True, keys["paired"]: False},
|
|
)
|
|
self.assertEqual(db.get_image_link_facts([]), {})
|
|
|
|
def test_new_link_fact_reads_use_indexes(self):
|
|
from database.operations import (
|
|
_LINK_FACTS_SQL,
|
|
_LINK_FILTER_HIT,
|
|
_POSITIVE_TAG_NAMES_SQL,
|
|
)
|
|
|
|
with self.api.db.model.get_connection() as conn:
|
|
conn.executemany(
|
|
"INSERT INTO prompts(id,text,hash) VALUES(?,?,?)",
|
|
[(i, f"prompt {i}", f"hash-{i}") for i in range(1, 3_001)],
|
|
)
|
|
conn.executemany(
|
|
"INSERT INTO generated_images(prompt_id,image_path,filename,"
|
|
"file_path) VALUES(?,?,?,?)",
|
|
[
|
|
(i, f"/x/{i}.png", f"{i}.png", f"/x/{i}.png")
|
|
for i in range(1, 3_001)
|
|
],
|
|
)
|
|
for hit, params in (
|
|
("1", {}),
|
|
(_LINK_FILTER_HIT, {"pattern": "%prompt 1%"}),
|
|
):
|
|
sql = _LINK_FACTS_SQL.format(hit=hit, marks=":path0,:path1")
|
|
bound = {**params, "path0": "/x/1.png", "path1": "/x/2.png"}
|
|
plan = [
|
|
row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, bound)
|
|
]
|
|
self.assertIn(
|
|
"SEARCH gi USING INDEX idx_image_file_path (file_path=?)", plan
|
|
)
|
|
self.assertFalse(
|
|
[s for s in plan if s.startswith(("SCAN gi", "SCAN other"))], plan
|
|
)
|
|
plan = [
|
|
row[3]
|
|
for row in conn.execute(
|
|
"EXPLAIN QUERY PLAN " + _POSITIVE_TAG_NAMES_SQL.format(marks="?,?"),
|
|
(1, 2),
|
|
)
|
|
]
|
|
self.assertFalse([step for step in plan if step.startswith("SCAN")], plan)
|