Vito
2964a3b632
* feat(thumbnails): thumbnail format, quality and size as server settings
Thumbnail format (webp default, jxl, jpeg), quality (1-100, default 80)
and size limit (128-2048 px, default 512) are server settings in the
config.json thumbnails section, saved with the settings revision. The
never-read gallery thumbnail_size is gone: older saves wrote its unused
default 256 into config.json.
One module names previews (<stem>_thumb.webp/.jxl/.jpg) and finds the
preview a source owns, preferring the configured format; legacy names
are still found. The encoder module writes every preview atomically and
verifies it decodes. JPEG XL needs no new dependency: Pillow (natively
or with pillow_jxl), ffmpeg with libjxl, or cjxl; without one the
settings refuse jxl (400 jxl-unavailable) and GET says why. Encoder
processes are tracked, bounded and killed at shutdown.
* feat(thumbnails): every preview writer and lookup follows the thumbnail settings
Automatic previews (capture and first view), the thumbnail job, the
legacy gallery routes and the history route write <stem>_thumb.webp
(or .jxl/.jpg) with the configured quality and size limit, for images
and videos alike: video frames, ffmpeg's included (a PNG on stdout,
never upscaled), and the placeholder are encoded like images.
Listings, duplicates, enrichment, the autotag scan and the history
route show any preview the file owns, the current format first, so
older _thumb.png previews keep working until converted. A request for
an older name answers the current preview. Previews are served with
an explicit image/webp or image/jxl type and the same cache headers.
The plan and status count missing, outdated (another format) and
ready previews; the missing job never writes beside an outdated one.
Removing a file removes its previews in every format; the clear
removes .webp and .jxl too. Encoder processes are killed at shutdown.
A preview is only a placeholder: inspection refuses paths under
thumbnails/ and always reads the original, and the inspector cover
accepts the listing's preview in any format the file owns.
* feat(thumbnails): Convert outdated thumbnails as an Activity job
POST /jobs/thumbnails {mode: "convert"} starts "Convert outdated
thumbnails" under the thumbnail job's kind, guarded admission, worker
threads and cancel at a file boundary. For each file whose preview is
in another format, it writes the new preview from the source (atomic,
verified; a video with no decodable frame fails instead of getting the
placeholder), then removes that file's old-format previews in its
thumbnails folder. A failure keeps the old preview; orphan previews
and links are never touched. The final line reports converted, failed
and bytes before and after. The plan counts it with mode=convert; the
legacy gallery stream refuses to attach to a running conversion.
* feat(thumbnails): thumbnail settings in Appearance and Convert outdated in the workspace
System > Appearance gains a Thumbnails surface: format (WebP, JPEG XL,
JPEG), quality and size limit, saved to the server with the settings
revision. JPEG XL is offered only when the server has an encoder, else
the option is disabled with the server's reason (D130); a note says
when this browser cannot show JPEG XL. Saving a new format converts
nothing: the row then says previews keep their old format.
Thumbnail maintenance gains the Convert outdated scope and names the
server's settings instead of its own quality and size fields (D131).
Overview's thumbnail row reads N missing, M outdated, K failed and
offers Convert outdated while previews are outdated.
The workspace decodes a tiny JPEG XL once: where it cannot, cards, the
Library and the inspector cover load the original for a .jxl preview,
and any preview that fails to load falls back to the original. The
viewer and downloads keep using the original (regression test).
README explains enabling JPEG XL; TESTING.md records D130 and D131.
* test(thumbnails): coverage floors for the thumbnail modules and WebP in page workflows
The two new frontend modules enter the coverage baseline (98% lines);
other floors only rise. The visual fixture's preview is WebP and
Appearance has its Thumbnails surface in the System shell check.
* fix(thumbnails): say when the thumbnail settings could not be read
Like Appearance's other server rows, the Thumbnails rows show 'Server
settings could not be read.' while the read failed, instead of only
staying disabled. Firefox exposed the gap: its JPEG XL note arrived
first and changed the row notes the page workflow counted.
* test(thumbnails): accept the legacy preview's spelling a case-insensitive disk finds
On Windows and macOS the '.jpg' lookup finds an older release's
b_thumb.JPG first, so the listed URL spells '.jpg'; it names the same
file there. The test now accepts either spelling and checks the URL
serves that preview.
* fix(thumbnails): address the final review of the thumbnail format work
- ffmpeg writes JPEG XL with -update 1, so a % in a file name is never an
image-sequence pattern, and reads videos as file:<path>. The clear and a
startup sweep remove temp files that interrupted writes left behind.
- Convert outdated converts each group of same-stem files once, from the
first source in the R2 order. An old preview that is already gone counts
as removed, and each file's bytes count once. A current preview is only
kept if it decodes.
- An encoder that fails makes the write detect encoders again once. Another
encoder is used if one is found; with none, the file is written as WebP
and the availability Appearance shows is refreshed.
- Changing the thumbnail format, quality or size while Convert outdated runs
returns 409 busy. If the format still changes underneath, the job stops
without removing anything.
- Legacy preview spellings are matched case-insensitively, so X_thumb.JPG
counts as outdated.
- History's cover uses the same lookup as Images. It serves .jxl only when
Accept lists image/jxl, and otherwise the original or the placeholder.
- A video whose thumbnail the browser cannot show keeps its own element
(Images) or its label (Library).
* fix(thumbnails): keep the new preview, probe encoders only when gone, sweep reused PIDs
- Convert never removes a path that is the preview it just wrote. It
compares the device and inode, or the resolved and normcased paths when
the disk has no inode numbers. This covers x_thumb.JPG on a
case-insensitive disk.
- A JPEG XL encode failure re-runs encoder detection only when the
encoder's binary or module is gone. A "no encoder" result is checked
again on the settings read, on a jxl save, and after 60 s, so one slow
probe no longer turns JPEG XL off until a restart.
- The startup sweep now also removes temp files that carry this process's
PID but predate its start. Containers reuse PIDs.
* fix(history): let the cover ask for JPEG XL when this browser decodes it
Firefox can decode JPEG XL without listing image/jxl in Accept. When the
workspace's probe decodes JPEG XL, it adds jxl=1 to the History cover URL.
The server serves a .jxl preview when that flag is set or when Accept
lists image/jxl. In every other case it falls back as before. The flag
accepts only 0 or 1, and the response stays no-store.
* fix(thumbnails): keep going past the per-batch cap until every preview is made
The missing, rebuild and Convert outdated jobs stopped after
MAX_THUMBNAIL_FILES (10,000) files and said "N more on the next run", but
nothing ever started that next run. The cap now bounds one batch. After
each batch the job plans again, leaving out every file it already worked
on, and continues until nothing is left, the job is cancelled, or a batch
makes no progress. That last case stops on files that fail every time
instead of looping over them.
Activity progress counts against the overall total, and the final line and
the legacy stream's complete frame report every batch. Admission, worker
threads, cancellation and the #316 rules are unchanged. The Thumbnail
maintenance dialog now counts all the work instead of saying "more on the
next run".
2345 lines
92 KiB
Python
2345 lines
92 KiB
Python
"""
|
|
Route tests for the image API (py/api/images.py).
|
|
|
|
Uses aiohttp's test client against a real PromptManagerAPI wired to a
|
|
temporary SQLite database and a temporary output directory.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from contextlib import ExitStack
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# ComfyUI-only modules are stubbed so the API can be imported standalone.
|
|
try:
|
|
from host_stubs import install_folder_paths_stub
|
|
except ImportError: # imported as tests.test_api_images
|
|
from tests.host_stubs import install_folder_paths_stub
|
|
|
|
install_folder_paths_stub()
|
|
sys.modules.setdefault("server", MagicMock())
|
|
|
|
from aiohttp import web # noqa: E402
|
|
from aiohttp.test_utils import AioHTTPTestCase # noqa: E402
|
|
from PIL import Image # noqa: E402
|
|
|
|
from database.operations import PromptDatabase # noqa: E402
|
|
import py.api.images as images_module
|
|
from py.api.thumbnail_settings import ThumbnailSettings # noqa: E402
|
|
from py.api import PromptManagerAPI # noqa: E402
|
|
from py.api.prompts import MAX_PAGE_LIMIT # noqa: E402
|
|
from py.config import GalleryConfig # noqa: E402
|
|
from utils.hashing import generate_prompt_hash # noqa: E402
|
|
|
|
|
|
def make_png(path, size=(4, 4), color=(255, 0, 0)):
|
|
"""Write a tiny but valid PNG to *path* and return it as a Path."""
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with Image.new("RGB", size, color) as img:
|
|
img.save(path, "PNG")
|
|
return path
|
|
|
|
|
|
class ImageAPITestCase(AioHTTPTestCase):
|
|
"""App with PromptManager routes, a temp DB and a temp output directory."""
|
|
|
|
async def get_application(self):
|
|
self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
|
|
self._temp_db.close()
|
|
|
|
self._output_tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True)
|
|
# A fake ComfyUI tree: <tmp>/output is the gallery root and <tmp> is the
|
|
# anchor that public (response) paths are rendered relative to.
|
|
self.comfy_dir = Path(os.path.realpath(self._output_tmp.name))
|
|
self.output_dir = self.comfy_dir / "output"
|
|
self.output_dir.mkdir()
|
|
anchors = [os.path.normcase(str(self.comfy_dir))]
|
|
self._anchor_patch = patch.object(
|
|
GalleryConfig, "path_anchors", classmethod(lambda cls: list(anchors))
|
|
)
|
|
self._anchor_patch.start()
|
|
|
|
app = web.Application()
|
|
routes = web.RouteTableDef()
|
|
|
|
self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name))
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir]
|
|
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
|
|
self.api.add_routes(routes)
|
|
app.router.add_routes(routes)
|
|
return app
|
|
|
|
async def tearDownAsync(self):
|
|
await super().tearDownAsync() # closes the aiohttp test client
|
|
# The automatic preview worker writes into the temp tree: stop it first.
|
|
self.api.close_auto_thumbnails()
|
|
self._anchor_patch.stop()
|
|
# Windows refuses to unlink a database that still has open handles.
|
|
self.api.db.close_all()
|
|
for path in (
|
|
self._temp_db.name,
|
|
self._temp_db.name + "-wal",
|
|
self._temp_db.name + "-shm",
|
|
):
|
|
if os.path.exists(path):
|
|
os.unlink(path)
|
|
self._output_tmp.cleanup()
|
|
|
|
# -- helpers ---------------------------------------------------------
|
|
|
|
def _save_prompt(self, text="Test prompt", **kwargs):
|
|
return self.api.db.save_prompt(
|
|
text=text, prompt_hash=generate_prompt_hash(text), **kwargs
|
|
)
|
|
|
|
def _link_image(self, image_path, text="Linked prompt"):
|
|
"""Create a prompt and link *image_path* to it; return the image id."""
|
|
prompt_id = self._save_prompt(text)
|
|
image_id = self.api.db.link_image_to_prompt(prompt_id, str(image_path))
|
|
self.assertGreater(image_id, 0, "fixture: image link failed")
|
|
return image_id
|
|
|
|
def _make_symlink(self, target, link):
|
|
"""Create a symlink or skip the test when the platform refuses."""
|
|
try:
|
|
os.symlink(str(target), str(link))
|
|
except (OSError, NotImplementedError, AttributeError) as exc:
|
|
self.skipTest(f"symlinks unavailable: {exc}")
|
|
|
|
|
|
class TestServeOutputImage(ImageAPITestCase):
|
|
"""GET /prompt_manager/images/serve/{filepath}"""
|
|
|
|
async def test_serves_png_inside_output_dir(self):
|
|
make_png(self.output_dir / "sub" / "pic.png")
|
|
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/sub/pic.png"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(resp.headers["Content-Type"], "image/png")
|
|
body = await resp.read()
|
|
self.assertTrue(body.startswith(b"\x89PNG"))
|
|
|
|
async def test_non_media_extensions_are_forbidden(self):
|
|
for name in ("secret.py", "prompts.db", "notes.txt"):
|
|
(self.output_dir / name).write_text("not an image")
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/serve/{name}"
|
|
)
|
|
self.assertEqual(resp.status, 403, name)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
async def test_extension_check_is_case_insensitive(self):
|
|
make_png(self.output_dir / "UPPER.PNG")
|
|
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/UPPER.PNG"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
|
|
async def test_empty_allowed_dirs_fails_closed(self):
|
|
make_png(self.output_dir / "pic.png")
|
|
self.api._get_all_output_dirs = lambda: []
|
|
|
|
resp = await self.client.request("GET", "/prompt_manager/images/serve/pic.png")
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
async def test_dot_dot_traversal_is_forbidden(self):
|
|
outside = Path(self._output_tmp.name).parent / "pm_outside_probe.png"
|
|
make_png(outside)
|
|
self.addCleanup(lambda: outside.exists() and outside.unlink())
|
|
|
|
# "..%2f" survives URL normalisation and reaches the handler as "../".
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/..%2fpm_outside_probe.png"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_absolute_path_is_forbidden(self):
|
|
outside = Path(self._output_tmp.name).parent / "pm_abs_probe.png"
|
|
make_png(outside)
|
|
self.addCleanup(lambda: outside.exists() and outside.unlink())
|
|
|
|
encoded = str(outside).replace(os.sep, "%2f").replace("/", "%2f")
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/serve/{encoded}"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_symlink_escaping_output_dir_is_forbidden(self):
|
|
with tempfile.TemporaryDirectory() as other:
|
|
target = make_png(Path(other) / "elsewhere.png")
|
|
self._make_symlink(target, self.output_dir / "escape.png")
|
|
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/escape.png"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_every_gallery_video_format_is_served(self):
|
|
# The gallery lists these as playable cards; opening them must not be a 403.
|
|
for ext in sorted(images_module.VIDEO_EXTENSIONS):
|
|
(self.output_dir / f"clip{ext}").write_bytes(b"\x00\x00\x00\x18ftyp")
|
|
|
|
for ext in sorted(images_module.VIDEO_EXTENSIONS):
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/serve/clip{ext}"
|
|
)
|
|
await resp.read()
|
|
self.assertEqual(resp.status, 200, ext)
|
|
|
|
async def test_missing_file_is_404(self):
|
|
resp = await self.client.request("GET", "/prompt_manager/images/serve/nope.png")
|
|
|
|
self.assertEqual(resp.status, 404)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
async def test_root_index_selects_single_root(self):
|
|
with tempfile.TemporaryDirectory() as second:
|
|
second_dir = Path(os.path.realpath(second))
|
|
make_png(second_dir / "only_here.png")
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second_dir]
|
|
|
|
hit = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/only_here.png?root=1"
|
|
)
|
|
miss = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/only_here.png?root=0"
|
|
)
|
|
bad_index = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/only_here.png?root=abc"
|
|
)
|
|
|
|
self.assertEqual(hit.status, 200)
|
|
self.assertEqual(miss.status, 404)
|
|
self.assertEqual(bad_index.status, 403)
|
|
for resp in (hit, miss, bad_index):
|
|
await resp.read() # release the served file before the directory goes
|
|
|
|
|
|
class TestServeImageById(ImageAPITestCase):
|
|
"""GET /prompt_manager/images/{image_id}/file"""
|
|
|
|
async def test_serves_linked_png(self):
|
|
png = make_png(self.output_dir / "linked.png")
|
|
image_id = self._link_image(png)
|
|
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(resp.headers["Content-Type"], "image/png")
|
|
|
|
async def test_non_media_record_is_forbidden(self):
|
|
txt = self.output_dir / "linked.txt"
|
|
txt.write_text("plain")
|
|
image_id = self._link_image(txt)
|
|
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_empty_allowed_dirs_fails_closed(self):
|
|
png = make_png(self.output_dir / "linked.png")
|
|
image_id = self._link_image(png)
|
|
self.api._get_all_output_dirs = lambda: []
|
|
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_record_outside_allowed_dirs_is_forbidden(self):
|
|
with tempfile.TemporaryDirectory() as other:
|
|
png = make_png(Path(other) / "outside.png")
|
|
image_id = self._link_image(png)
|
|
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_symlink_record_escaping_is_forbidden(self):
|
|
with tempfile.TemporaryDirectory() as other:
|
|
target = make_png(Path(other) / "elsewhere.png")
|
|
link = self.output_dir / "escape.png"
|
|
self._make_symlink(target, link)
|
|
image_id = self._link_image(link)
|
|
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_missing_file_is_404(self):
|
|
image_id = self._link_image(self.output_dir / "gone.png")
|
|
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 404)
|
|
|
|
async def test_unknown_id_is_404(self):
|
|
resp = await self.client.request("GET", "/prompt_manager/images/424242/file")
|
|
|
|
self.assertEqual(resp.status, 404)
|
|
|
|
async def test_non_numeric_id_is_400(self):
|
|
resp = await self.client.request("GET", "/prompt_manager/images/abc/file")
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
|
|
|
|
class TestRecentImagesBounds(ImageAPITestCase):
|
|
"""GET /prompt_manager/images/recent clamps limit/offset and rejects junk."""
|
|
|
|
async def test_limit_is_clamped_to_max_page_limit(self):
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/recent?limit=999999"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
data = await resp.json()
|
|
self.assertTrue(data["success"])
|
|
self.assertLessEqual(len(data["images"]), MAX_PAGE_LIMIT)
|
|
self.assertEqual(data["pagination"]["limit"], MAX_PAGE_LIMIT)
|
|
|
|
async def test_non_integer_limit_is_400(self):
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/recent?limit=abc"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
self.assertIn("error", data)
|
|
|
|
async def test_negative_offset_is_clamped_to_zero(self):
|
|
# Negative offsets are clamped rather than rejected so that a
|
|
# client paging backwards past the first page still gets page one.
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/recent?offset=-5"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
data = await resp.json()
|
|
self.assertEqual(data["pagination"]["offset"], 0)
|
|
|
|
async def test_offset_skips_rows(self):
|
|
for name in ("a.png", "b.png", "c.png"):
|
|
self._link_image(make_png(self.output_dir / name), text=f"p-{name}")
|
|
|
|
page = await self.client.request(
|
|
"GET", "/prompt_manager/images/recent?limit=2&offset=2"
|
|
)
|
|
|
|
data = await page.json()
|
|
self.assertEqual(len(data["images"]), 1)
|
|
self.assertEqual(data["pagination"]["offset"], 2)
|
|
|
|
|
|
def parse_sse(text):
|
|
"""Turn an SSE body into a list of (event, data_dict) tuples."""
|
|
events = []
|
|
for block in text.strip().split("\n\n"):
|
|
event, data = None, None
|
|
for line in block.splitlines():
|
|
if line.startswith("event:"):
|
|
event = line[len("event:") :].strip()
|
|
elif line.startswith("data:"):
|
|
data = json.loads(line[len("data:") :].strip())
|
|
if event is not None:
|
|
events.append((event, data))
|
|
return events
|
|
|
|
|
|
class ThumbnailTestCase(ImageAPITestCase):
|
|
"""Shared helpers for the thumbnail routes."""
|
|
|
|
PROGRESS_URL = "/prompt_manager/images/generate-thumbnails/progress"
|
|
|
|
def _spy_executor(self):
|
|
"""Record the __name__ of every callable sent to _run_in_executor."""
|
|
original = self.api._run_in_executor
|
|
calls = []
|
|
|
|
async def spy(func, *args, **kwargs):
|
|
calls.append(getattr(func, "__name__", repr(func)))
|
|
return await original(func, *args, **kwargs)
|
|
|
|
self.api._run_in_executor = spy
|
|
return calls
|
|
|
|
def _thumb_path(self, name, suffix=".webp"):
|
|
return self.output_dir / "thumbnails" / f"{name}_thumb{suffix}"
|
|
|
|
|
|
class TestThumbnailProgressStream(ThumbnailTestCase):
|
|
"""POST /prompt_manager/images/generate-thumbnails/progress (SSE)."""
|
|
|
|
def _record_threads(self):
|
|
"""Make every thumbnail write record its thread and take a moment."""
|
|
import threading
|
|
import time
|
|
|
|
from py.api import images as images_module
|
|
|
|
seen = set()
|
|
lock = threading.Lock()
|
|
original = images_module._write_image_thumbnail
|
|
|
|
def recording(src, dst, size, *args):
|
|
with lock:
|
|
seen.add(threading.current_thread().name)
|
|
time.sleep(0.05)
|
|
return original(src, dst, size, *args)
|
|
|
|
images_module._write_image_thumbnail = recording
|
|
self.addCleanup(setattr, images_module, "_write_image_thumbnail", original)
|
|
return seen
|
|
|
|
async def test_thumbnails_are_generated_on_several_threads(self):
|
|
from py.config import PromptManagerConfig
|
|
|
|
for i in range(8):
|
|
make_png(self.output_dir / f"p{i}.png", size=(64, 64))
|
|
orig = PromptManagerConfig.WORKER_THREADS
|
|
PromptManagerConfig.WORKER_THREADS = 4
|
|
self.addCleanup(setattr, PromptManagerConfig, "WORKER_THREADS", orig)
|
|
seen = self._record_threads()
|
|
|
|
resp = await self.client.request("POST", self.PROGRESS_URL)
|
|
events = parse_sse(await resp.text())
|
|
|
|
names = [e for e, _ in events]
|
|
self.assertEqual(names.count("progress"), 8)
|
|
self.assertEqual(dict(events)["complete"]["count"], 8)
|
|
self.assertGreaterEqual(len(seen), 2)
|
|
for i in range(8):
|
|
self.assertTrue(self._thumb_path(f"p{i}").is_file())
|
|
|
|
async def test_blocking_endpoint_also_uses_the_worker_threads(self):
|
|
from py.config import PromptManagerConfig
|
|
|
|
for i in range(8):
|
|
make_png(self.output_dir / f"p{i}.png", size=(64, 64))
|
|
orig = PromptManagerConfig.WORKER_THREADS
|
|
PromptManagerConfig.WORKER_THREADS = 4
|
|
self.addCleanup(setattr, PromptManagerConfig, "WORKER_THREADS", orig)
|
|
seen = self._record_threads()
|
|
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
|
|
self.assertEqual((await resp.json())["count"], 8)
|
|
self.assertGreaterEqual(len(seen), 2)
|
|
|
|
async def test_get_is_method_not_allowed(self):
|
|
# The stream writes thumbnails, so a plain GET (prefetch, link
|
|
# preview, <img src>) must not be able to trigger it.
|
|
make_png(self.output_dir / "a.png")
|
|
|
|
resp = await self.client.request("GET", self.PROGRESS_URL)
|
|
|
|
self.assertEqual(resp.status, 405)
|
|
self.assertFalse(self._thumb_path("a").exists())
|
|
|
|
async def test_three_pngs_yield_three_thumbnails_and_events(self):
|
|
for name in ("a", "b", "c"):
|
|
make_png(self.output_dir / f"{name}.png", size=(64, 64))
|
|
calls = self._spy_executor()
|
|
|
|
resp = await self.client.request("POST", f"{self.PROGRESS_URL}?quality=low")
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(resp.headers["Content-Type"], "text/event-stream")
|
|
events = parse_sse(await resp.text())
|
|
names = [e for e, _ in events]
|
|
self.assertEqual(names.count("progress"), 3)
|
|
self.assertEqual(names.count("complete"), 1)
|
|
complete = dict(events)["complete"]
|
|
self.assertEqual(complete["count"], 3)
|
|
for name in ("a", "b", "c"):
|
|
self.assertTrue(self._thumb_path(name).is_file(), name)
|
|
with Image.open(self._thumb_path("a")) as thumb:
|
|
self.assertLessEqual(max(thumb.size), 150)
|
|
# PIL work ran in the executor in batches (never on the loop), plus the scan.
|
|
self.assertGreaterEqual(calls.count("_generate_batch"), 1)
|
|
self.assertNotIn("_generate_one", calls)
|
|
self.assertIn("_plan_gallery_thumbnails", calls)
|
|
|
|
async def test_a_run_past_the_batch_cap_goes_on_to_the_end(self):
|
|
# The cap bounds one batch; the run plans again and finishes (#317).
|
|
for name in ("a", "b", "c"):
|
|
make_png(self.output_dir / f"{name}.png", size=(64, 64))
|
|
|
|
with patch.object(images_module, "MAX_THUMBNAIL_FILES", 2):
|
|
first = parse_sse(
|
|
await (await self.client.request("POST", self.PROGRESS_URL)).text()
|
|
)
|
|
second = parse_sse(
|
|
await (await self.client.request("POST", self.PROGRESS_URL)).text()
|
|
)
|
|
|
|
names = [e for e, _ in first]
|
|
self.assertEqual(names.count("progress"), 3) # only pending files are work
|
|
self.assertEqual(dict(first)["start"]["total_images"], 3)
|
|
self.assertEqual(
|
|
[d["total_images"] for e, d in first if e == "progress"], [3] * 3
|
|
)
|
|
done = dict(first)["complete"]
|
|
self.assertEqual(done["count"], 3)
|
|
self.assertEqual(done["remaining"], 0)
|
|
self.assertEqual(done["total_images"], 3)
|
|
self.assertNotIn("more", done["message"])
|
|
done = dict(second)["complete"]
|
|
self.assertEqual(done["count"], 0)
|
|
self.assertEqual(done["skipped"], 3)
|
|
|
|
async def test_blocking_endpoint_reports_remaining(self):
|
|
for name in ("a", "b", "c"):
|
|
make_png(self.output_dir / f"{name}.png", size=(64, 64))
|
|
|
|
with patch.object(images_module, "MAX_THUMBNAIL_FILES", 2):
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
data = await resp.json()
|
|
|
|
self.assertEqual(data["count"], 2)
|
|
self.assertEqual(data["remaining"], 1)
|
|
self.assertEqual(data["total_images"], 3)
|
|
|
|
async def test_corrupt_png_is_reported_and_loop_continues(self):
|
|
make_png(self.output_dir / "good1.png")
|
|
(self.output_dir / "bad.png").write_bytes(b"\x89PNG definitely not a png")
|
|
make_png(self.output_dir / "good2.png")
|
|
|
|
resp = await self.client.request("POST", self.PROGRESS_URL)
|
|
|
|
events = parse_sse(await resp.text())
|
|
file_errors = [d for e, d in events if e == "file_error"]
|
|
self.assertEqual(len(file_errors), 1)
|
|
self.assertEqual(file_errors[0]["file"], "bad.png")
|
|
complete = dict(events)["complete"]
|
|
self.assertEqual(complete["count"], 2)
|
|
self.assertEqual(complete["error_count"], 1)
|
|
self.assertTrue(self._thumb_path("good1").is_file())
|
|
self.assertTrue(self._thumb_path("good2").is_file())
|
|
self.assertFalse(self._thumb_path("bad").exists())
|
|
|
|
async def test_response_has_no_wildcard_cors_header(self):
|
|
make_png(self.output_dir / "one.png")
|
|
|
|
resp = await self.client.request("POST", self.PROGRESS_URL)
|
|
await resp.text()
|
|
|
|
self.assertNotIn("Access-Control-Allow-Origin", resp.headers)
|
|
|
|
async def test_existing_newer_thumbnail_is_skipped(self):
|
|
make_png(self.output_dir / "one.png")
|
|
first = await self.client.request("POST", self.PROGRESS_URL)
|
|
await first.text()
|
|
|
|
second = await self.client.request("POST", self.PROGRESS_URL)
|
|
|
|
complete = dict(parse_sse(await second.text()))["complete"]
|
|
self.assertEqual(complete["skipped"], 1)
|
|
self.assertEqual(complete["count"], 0)
|
|
|
|
async def test_empty_output_dir_completes_with_zero(self):
|
|
resp = await self.client.request("POST", self.PROGRESS_URL)
|
|
|
|
events = parse_sse(await resp.text())
|
|
self.assertEqual(events[-1][0], "complete")
|
|
self.assertEqual(events[-1][1]["total_images"], 0)
|
|
|
|
async def test_missing_output_dir_sends_error_event(self):
|
|
self.api._get_all_output_dirs = lambda: []
|
|
|
|
resp = await self.client.request("POST", self.PROGRESS_URL)
|
|
|
|
events = parse_sse(await resp.text())
|
|
self.assertEqual(events[-1][0], "error")
|
|
|
|
async def test_any_quality_preset_uses_the_server_size_limit(self):
|
|
# #317 R3: the server's thumbnail settings size every preview.
|
|
make_png(self.output_dir / "one.png", size=(1000, 1000))
|
|
|
|
resp = await self.client.request("POST", f"{self.PROGRESS_URL}?quality=huge")
|
|
await resp.text()
|
|
|
|
with Image.open(self._thumb_path("one")) as thumb:
|
|
self.assertEqual(max(thumb.size), 512)
|
|
|
|
|
|
class TestGenerateThumbnailsPost(ThumbnailTestCase):
|
|
"""POST /prompt_manager/images/generate-thumbnails (blocking variant)."""
|
|
|
|
async def test_generates_then_skips(self):
|
|
make_png(self.output_dir / "sub" / "x.png")
|
|
make_png(self.output_dir / "y.png")
|
|
|
|
first = await self.client.request(
|
|
"POST",
|
|
"/prompt_manager/images/generate-thumbnails",
|
|
json={"quality": "high"},
|
|
)
|
|
second = await self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
|
|
self.assertEqual(first.status, 200)
|
|
data = await first.json()
|
|
self.assertTrue(data["success"])
|
|
self.assertEqual(data["count"], 2)
|
|
self.assertTrue(self._thumb_path("sub/x").is_file())
|
|
data2 = await second.json()
|
|
self.assertEqual(data2["skipped"], 2)
|
|
self.assertEqual(data2["count"], 0)
|
|
|
|
async def test_missing_output_dir_is_404(self):
|
|
self.api._get_all_output_dirs = lambda: []
|
|
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
|
|
self.assertEqual(resp.status, 404)
|
|
|
|
async def test_no_media_reports_zero(self):
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
|
|
data = await resp.json()
|
|
self.assertTrue(data["success"])
|
|
self.assertEqual(data["total_images"], 0)
|
|
|
|
|
|
class TestThumbnailRoutesHoldTheRegistry(ThumbnailTestCase):
|
|
"""Thumbnail work is ``guarded``: never beside a record-changing job."""
|
|
|
|
def _run_scan(self):
|
|
gate = asyncio.Event()
|
|
|
|
async def runner(job):
|
|
await gate.wait()
|
|
|
|
job = self.api._jobs.start(
|
|
"scan", "Scan for metadata", admission="records", scope={}, runner=runner
|
|
)
|
|
|
|
async def release():
|
|
gate.set()
|
|
await asyncio.wait_for(job.task, 5)
|
|
|
|
self.addAsyncCleanup(release)
|
|
return job
|
|
|
|
async def test_thumbnail_generate_refused_while_scan_runs(self):
|
|
make_png(self.output_dir / "a.png")
|
|
self._run_scan()
|
|
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
self.assertEqual(resp.status, 409)
|
|
self.assertEqual(
|
|
await resp.json(),
|
|
{"success": False, "error": "Scan for metadata already running"},
|
|
)
|
|
|
|
resp = await self.client.request("POST", self.PROGRESS_URL)
|
|
text = await resp.text()
|
|
self.assertEqual(resp.status, 200) # the gallery reads an SSE stream
|
|
self.assertEqual(
|
|
text,
|
|
"event: error\ndata: "
|
|
+ json.dumps(
|
|
{
|
|
"error": "Scan for metadata already running",
|
|
"message": "Scan for metadata already running",
|
|
}
|
|
)
|
|
+ "\n\n",
|
|
)
|
|
|
|
thumbs = self.output_dir / "thumbnails"
|
|
make_png(thumbs / "b_thumb.png")
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/images/clear-thumbnails"
|
|
)
|
|
self.assertEqual(resp.status, 409)
|
|
self.assertFalse(self._thumb_path("a").exists())
|
|
self.assertTrue((thumbs / "b_thumb.png").exists())
|
|
|
|
async def test_scan_refused_while_thumbnails_generate(self):
|
|
make_png(self.output_dir / "a.png")
|
|
entered = threading.Event()
|
|
release = threading.Event()
|
|
real = self.api._generate_gallery_thumbnails_sync
|
|
|
|
def gated(*args):
|
|
entered.set()
|
|
release.wait(10)
|
|
return real(*args)
|
|
|
|
self.api._generate_gallery_thumbnails_sync = gated
|
|
loop = asyncio.get_running_loop()
|
|
generate = asyncio.ensure_future(
|
|
self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
)
|
|
self.assertTrue(await loop.run_in_executor(None, entered.wait, 10))
|
|
resp = await self.client.request("POST", "/prompt_manager/jobs/scan")
|
|
body = await resp.json()
|
|
self.assertEqual((resp.status, body["code"]), (409, "busy"))
|
|
self.assertEqual(
|
|
body["running"]["title"], "Thumbnail generation (gallery page)"
|
|
)
|
|
release.set()
|
|
self.assertEqual((await generate).status, 200)
|
|
self.assertTrue(self._thumb_path("a").exists())
|
|
|
|
|
|
class TestRecentImagesCaps(ImageAPITestCase):
|
|
"""Every path through /images/recent hands the DB a bounded window."""
|
|
|
|
def _spy_db(self):
|
|
seen = []
|
|
original = self.api.db.get_recent_images
|
|
|
|
def spy(limit, offset=0):
|
|
seen.append((limit, offset))
|
|
return original(limit, offset)
|
|
|
|
self.api.db.get_recent_images = spy
|
|
return seen
|
|
|
|
async def test_absent_params_use_bounded_default(self):
|
|
seen = self._spy_db()
|
|
|
|
resp = await self.client.request("GET", "/prompt_manager/images/recent")
|
|
|
|
data = await resp.json()
|
|
self.assertEqual(data["pagination"]["limit"], 50)
|
|
self.assertEqual(seen, [(50, 0)])
|
|
|
|
async def test_huge_offset_is_clamped_before_reaching_db(self):
|
|
seen = self._spy_db()
|
|
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/recent?limit=10&offset=999999999999"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
data = await resp.json()
|
|
# The offset is paged in SQL, so the DB sees the clamped value itself.
|
|
self.assertEqual(seen[0], (10, data["pagination"]["offset"]))
|
|
self.assertLess(data["pagination"]["offset"], 999999999999)
|
|
|
|
async def test_float_string_limit_is_400(self):
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/recent?limit=1e9"
|
|
)
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
|
|
async def test_zero_limit_becomes_one(self):
|
|
seen = self._spy_db()
|
|
|
|
await self.client.request("GET", "/prompt_manager/images/recent?limit=0")
|
|
|
|
self.assertEqual(seen, [(1, 0)])
|
|
|
|
|
|
class TestAllImagesCaps(ImageAPITestCase):
|
|
"""/images/all never asks the DB for an unbounded result."""
|
|
|
|
def _spy_db(self):
|
|
seen = []
|
|
original = self.api.db.get_all_images
|
|
|
|
def spy(limit=0, offset=0):
|
|
seen.append((limit, offset))
|
|
return original(limit=limit, offset=offset)
|
|
|
|
self.api.db.get_all_images = spy
|
|
return seen
|
|
|
|
async def test_absent_params_use_bulk_ceiling(self):
|
|
seen = self._spy_db()
|
|
|
|
resp = await self.client.request("GET", "/prompt_manager/images/all")
|
|
|
|
data = await resp.json()
|
|
self.assertTrue(data["success"])
|
|
self.assertEqual(data["pagination"]["limit"], images_module.MAX_BULK_LIMIT)
|
|
self.assertEqual(seen, [(images_module.MAX_BULK_LIMIT, 0)])
|
|
|
|
async def test_limit_above_bulk_ceiling_is_clamped(self):
|
|
seen = self._spy_db()
|
|
|
|
await self.client.request("GET", "/prompt_manager/images/all?limit=99999999")
|
|
|
|
self.assertEqual(seen[0][0], images_module.MAX_BULK_LIMIT)
|
|
|
|
async def test_negative_offset_is_clamped(self):
|
|
seen = self._spy_db()
|
|
|
|
await self.client.request("GET", "/prompt_manager/images/all?offset=-9")
|
|
|
|
self.assertEqual(seen[0][1], 0)
|
|
|
|
async def test_non_integer_limit_is_400(self):
|
|
resp = await self.client.request("GET", "/prompt_manager/images/all?limit=abc")
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
|
|
async def test_returns_linked_images(self):
|
|
self._link_image(make_png(self.output_dir / "one.png"))
|
|
|
|
resp = await self.client.request("GET", "/prompt_manager/images/all?limit=1")
|
|
|
|
data = await resp.json()
|
|
self.assertEqual(data["count"], 1)
|
|
self.assertEqual(data["pagination"]["count"], 1)
|
|
|
|
|
|
class TestOutputImagesCaps(ImageAPITestCase):
|
|
"""GET /prompt_manager/images/output clamps paging on every path."""
|
|
|
|
URL = "/prompt_manager/images/output"
|
|
|
|
async def test_absent_params_use_default_limit(self):
|
|
resp = await self.client.request("GET", self.URL)
|
|
|
|
data = await resp.json()
|
|
self.assertTrue(data["success"])
|
|
self.assertEqual(data["limit"], 100)
|
|
self.assertEqual(data["offset"], 0)
|
|
|
|
async def test_limit_is_clamped_to_max_page_limit(self):
|
|
resp = await self.client.request("GET", f"{self.URL}?limit=999999")
|
|
|
|
data = await resp.json()
|
|
self.assertEqual(data["limit"], MAX_PAGE_LIMIT)
|
|
|
|
async def test_zero_limit_becomes_one(self):
|
|
resp = await self.client.request("GET", f"{self.URL}?limit=0")
|
|
|
|
data = await resp.json()
|
|
self.assertEqual(data["limit"], 1)
|
|
|
|
async def test_float_string_limit_is_400(self):
|
|
resp = await self.client.request("GET", f"{self.URL}?limit=1e9")
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
async def test_negative_offset_is_clamped(self):
|
|
resp = await self.client.request("GET", f"{self.URL}?offset=-5")
|
|
|
|
data = await resp.json()
|
|
self.assertEqual(data["offset"], 0)
|
|
|
|
async def test_pages_and_filters_by_subfolder(self):
|
|
make_png(self.output_dir / "sub" / "a.png")
|
|
make_png(self.output_dir / "sub" / "b.png")
|
|
make_png(self.output_dir / "top.png")
|
|
make_png(self.output_dir / "thumbnails" / "top_thumb.png")
|
|
|
|
page = await self.client.request("GET", f"{self.URL}?subfolder=sub&limit=1")
|
|
rest = await self.client.request("GET", f"{self.URL}?limit=10")
|
|
|
|
data = await page.json()
|
|
self.assertEqual(data["total"], 2)
|
|
self.assertEqual(len(data["images"]), 1)
|
|
self.assertTrue(data["has_more"])
|
|
self.assertEqual(data["images"][0]["root_index"], 0)
|
|
all_data = await rest.json()
|
|
names = sorted(i["filename"] for i in all_data["images"])
|
|
self.assertEqual(names, ["a.png", "b.png", "top.png"])
|
|
top = next(i for i in all_data["images"] if i["filename"] == "top.png")
|
|
self.assertIn("top_thumb.png", top["thumbnail_url"])
|
|
|
|
async def test_no_output_dirs_reports_failure(self):
|
|
self.api._get_all_output_dirs = lambda: []
|
|
|
|
resp = await self.client.request("GET", self.URL)
|
|
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
|
|
class TestThumbnailFilesAreNeverSeenHalfWritten(ThumbnailTestCase):
|
|
"""A browser that lists the gallery mid-run must never fetch a partial file,
|
|
and a regenerated thumbnail must never be served from a stale cache entry."""
|
|
|
|
def test_image_thumbnail_write_is_atomic(self):
|
|
import os
|
|
|
|
make_png(self.output_dir / "a.png", size=(64, 64))
|
|
thumbs = self.output_dir / "thumbnails"
|
|
thumbs.mkdir()
|
|
dst = thumbs / "a_thumb.webp"
|
|
settings = ThumbnailSettings("webp", 80, 32)
|
|
|
|
def crash_mid_write(self_img, fp, *args, **kwargs):
|
|
# Simulates an interrupted encoder: bytes on disk, then failure.
|
|
pathlib.Path(fp).write_bytes(b"\x89PNG partial")
|
|
raise OSError("disk full")
|
|
|
|
with patch.object(Image.Image, "save", crash_mid_write):
|
|
with self.assertRaises(OSError):
|
|
images_module._write_image_thumbnail(
|
|
self.output_dir / "a.png", dst, settings
|
|
)
|
|
self.assertFalse(dst.exists(), "a half-written thumbnail was left in place")
|
|
self.assertEqual(os.listdir(thumbs), [])
|
|
|
|
images_module._write_image_thumbnail(self.output_dir / "a.png", dst, settings)
|
|
self.assertTrue(dst.is_file())
|
|
self.assertEqual(os.listdir(thumbs), ["a_thumb.webp"])
|
|
|
|
def test_thumbnail_url_is_versioned_by_its_mtime(self):
|
|
import os
|
|
import re
|
|
|
|
make_png(self.output_dir / "a.png", size=(64, 64))
|
|
self.api._generate_thumbnails_sync(
|
|
self.output_dir,
|
|
self.output_dir / "thumbnails",
|
|
ThumbnailSettings("webp", 80, 32),
|
|
)
|
|
thumb = self._thumb_path("a")
|
|
|
|
entry = images_module._output_image_entry(
|
|
self.output_dir / "a.png", self.output_dir, 0
|
|
)
|
|
self.assertRegex(
|
|
entry["thumbnail_url"],
|
|
r"^/prompt_manager/images/serve/thumbnails/a_thumb\.webp\?v=\d+&root=0$",
|
|
)
|
|
first = int(re.search(r"v=(\d+)", entry["thumbnail_url"]).group(1))
|
|
self.assertEqual(first, int(thumb.stat().st_mtime))
|
|
|
|
later = thumb.stat().st_mtime + 10
|
|
os.utime(thumb, (later, later))
|
|
entry = images_module._output_image_entry(
|
|
self.output_dir / "a.png", self.output_dir, 0
|
|
)
|
|
second = int(re.search(r"v=(\d+)", entry["thumbnail_url"]).group(1))
|
|
self.assertEqual(second, first + 10)
|
|
|
|
async def test_versioned_thumbnail_url_is_served(self):
|
|
make_png(self.output_dir / "a.png", size=(64, 64))
|
|
self.api._generate_thumbnails_sync(
|
|
self.output_dir,
|
|
self.output_dir / "thumbnails",
|
|
ThumbnailSettings("webp", 80, 32),
|
|
)
|
|
entry = images_module._output_image_entry(
|
|
self.output_dir / "a.png", self.output_dir, 0
|
|
)
|
|
|
|
resp = await self.client.request("GET", entry["thumbnail_url"])
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(resp.headers["Content-Type"], "image/webp")
|
|
|
|
async def test_missing_image_404_is_not_cacheable(self):
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/thumbnails/nope_thumb.png"
|
|
)
|
|
self.assertEqual(resp.status, 404)
|
|
self.assertEqual(resp.headers.get("Cache-Control"), "no-store")
|
|
|
|
|
|
class TestMediaScanCaps(ImageAPITestCase):
|
|
"""Request-triggered filesystem walks are depth- and count-capped."""
|
|
|
|
def test_gallery_scan_honours_file_cap(self):
|
|
for name in ("a", "b", "c"):
|
|
make_png(self.output_dir / f"{name}.png")
|
|
|
|
with patch.object(images_module, "MAX_GALLERY_FILES", 2):
|
|
found = self.api._scan_gallery_files_sync(self.output_dir)
|
|
|
|
self.assertEqual(len(found), 2)
|
|
|
|
def test_gallery_scan_honours_depth_cap(self):
|
|
make_png(self.output_dir / "a" / "shallow.png")
|
|
make_png(self.output_dir / "a" / "b" / "deep.png")
|
|
|
|
with patch.object(images_module, "MAX_SCAN_DEPTH", 1):
|
|
found = [
|
|
p.name for p, _ in self.api._scan_gallery_files_sync(self.output_dir)
|
|
]
|
|
|
|
self.assertEqual(found, ["shallow.png"])
|
|
|
|
def test_gallery_scan_skips_symlinked_directories(self):
|
|
with tempfile.TemporaryDirectory() as other:
|
|
make_png(Path(other) / "outside.png")
|
|
self._make_symlink(other, self.output_dir / "link")
|
|
|
|
found = self.api._scan_gallery_files_sync(self.output_dir)
|
|
|
|
self.assertEqual(found, [])
|
|
|
|
def test_gallery_scan_skips_thumbnails_and_non_media(self):
|
|
make_png(self.output_dir / "keep.png")
|
|
make_png(self.output_dir / "thumbnails" / "keep_thumb.png")
|
|
(self.output_dir / "notes.txt").write_text("x")
|
|
|
|
found = [p.name for p, _ in self.api._scan_gallery_files_sync(self.output_dir)]
|
|
|
|
self.assertEqual(found, ["keep.png"])
|
|
|
|
def test_thumbnail_plan_skips_fresh_thumbnails_before_applying_the_cap(self):
|
|
for name in ("a", "b", "c"):
|
|
make_png(self.output_dir / f"{name}.png")
|
|
self.api._generate_thumbnails_sync(
|
|
self.output_dir,
|
|
self.output_dir / "thumbnails",
|
|
ThumbnailSettings("webp", 80, 64),
|
|
)
|
|
for name in ("d", "e", "f"):
|
|
make_png(self.output_dir / f"{name}.png")
|
|
|
|
with patch.object(images_module, "MAX_THUMBNAIL_FILES", 2):
|
|
plan = self.api._plan_thumbnails(
|
|
self.output_dir, self.output_dir / "thumbnails"
|
|
)
|
|
|
|
self.assertEqual(plan["skipped"], 3)
|
|
self.assertEqual([t[0].name for t in plan["targets"]], ["d.png", "e.png"])
|
|
self.assertEqual(plan["remaining"], 1)
|
|
|
|
def test_thumbnail_walk_is_bounded_like_the_gallery_walk(self):
|
|
for name in ("a", "b", "c"):
|
|
make_png(self.output_dir / f"{name}.png")
|
|
|
|
with patch.object(images_module, "MAX_GALLERY_FILES", 2):
|
|
plan = self.api._plan_thumbnails(
|
|
self.output_dir, self.output_dir / "thumbnails"
|
|
)
|
|
|
|
self.assertEqual(len(plan["targets"]), 2)
|
|
self.assertEqual(plan["remaining"], 0)
|
|
|
|
def test_thumbnail_targets_honour_file_cap(self):
|
|
for name in ("a", "b", "c"):
|
|
make_png(self.output_dir / f"{name}.png")
|
|
|
|
with patch.object(images_module, "MAX_THUMBNAIL_FILES", 2):
|
|
targets = self.api._thumbnail_targets(
|
|
self.output_dir, self.output_dir / "thumbnails"
|
|
)
|
|
|
|
self.assertEqual(len(targets), 2)
|
|
|
|
def test_thumbnail_targets_honour_depth_cap(self):
|
|
make_png(self.output_dir / "a" / "shallow.png")
|
|
make_png(self.output_dir / "a" / "b" / "deep.png")
|
|
|
|
with patch.object(images_module, "MAX_SCAN_DEPTH", 1):
|
|
targets = self.api._thumbnail_targets(
|
|
self.output_dir, self.output_dir / "thumbnails"
|
|
)
|
|
|
|
self.assertEqual([src.name for src, _, _ in targets], ["shallow.png"])
|
|
|
|
def test_thumbnail_targets_skip_symlinked_directories(self):
|
|
with tempfile.TemporaryDirectory() as other:
|
|
make_png(Path(other) / "outside.png")
|
|
self._make_symlink(other, self.output_dir / "link")
|
|
|
|
targets = self.api._thumbnail_targets(
|
|
self.output_dir, self.output_dir / "thumbnails"
|
|
)
|
|
|
|
self.assertEqual(targets, [])
|
|
|
|
|
|
class TestJsonBodyCaps(ImageAPITestCase):
|
|
"""JSON bodies are read in chunks and rejected past MAX_JSON_BODY_BYTES."""
|
|
|
|
LINK_URL = "/prompt_manager/images/link"
|
|
THUMBS_URL = "/prompt_manager/images/generate-thumbnails"
|
|
|
|
def _oversized(self, cap):
|
|
return b'{"image_path": "' + b"a" * (cap + 16) + b'"}'
|
|
|
|
async def test_oversized_content_length_is_413(self):
|
|
with patch.object(images_module, "MAX_JSON_BODY_BYTES", 512):
|
|
resp = await self.client.request(
|
|
"POST",
|
|
self.LINK_URL,
|
|
data=self._oversized(512),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
|
|
self.assertEqual(resp.status, 413)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
async def test_oversized_chunked_body_is_413(self):
|
|
with patch.object(images_module, "MAX_JSON_BODY_BYTES", 512):
|
|
resp = await self.client.request(
|
|
"POST",
|
|
self.THUMBS_URL,
|
|
data=self._oversized(512),
|
|
chunked=True,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
|
|
self.assertEqual(resp.status, 413)
|
|
|
|
async def test_invalid_json_is_400(self):
|
|
resp = await self.client.request(
|
|
"POST",
|
|
self.LINK_URL,
|
|
data=b"{not json",
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
data = await resp.json()
|
|
self.assertFalse(data["success"])
|
|
|
|
async def test_non_object_json_is_400(self):
|
|
resp = await self.client.request("POST", self.THUMBS_URL, json=[1, 2, 3])
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
|
|
|
|
class TestLinkImage(ImageAPITestCase):
|
|
"""POST /prompt_manager/images/link only links media inside allowed dirs."""
|
|
|
|
URL = "/prompt_manager/images/link"
|
|
|
|
async def test_links_png_inside_output_dir(self):
|
|
prompt_id = self._save_prompt("to link")
|
|
png = make_png(self.output_dir / "gen.png")
|
|
|
|
resp = await self.client.request(
|
|
"POST", self.URL, json={"prompt_id": prompt_id, "image_path": str(png)}
|
|
)
|
|
|
|
self.assertEqual(resp.status, 200)
|
|
data = await resp.json()
|
|
self.assertTrue(data["success"])
|
|
self.assertGreater(data["image_id"], 0)
|
|
|
|
async def test_missing_fields_is_400(self):
|
|
resp = await self.client.request("POST", self.URL, json={"prompt_id": 1})
|
|
|
|
self.assertEqual(resp.status, 400)
|
|
|
|
async def test_path_outside_allowed_dirs_is_403(self):
|
|
prompt_id = self._save_prompt("to link")
|
|
with tempfile.TemporaryDirectory() as other:
|
|
png = make_png(Path(other) / "outside.png")
|
|
|
|
resp = await self.client.request(
|
|
"POST", self.URL, json={"prompt_id": prompt_id, "image_path": str(png)}
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_non_media_path_is_403(self):
|
|
prompt_id = self._save_prompt("to link")
|
|
txt = self.output_dir / "secret.txt"
|
|
txt.write_text("x")
|
|
|
|
resp = await self.client.request(
|
|
"POST", self.URL, json={"prompt_id": prompt_id, "image_path": str(txt)}
|
|
)
|
|
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_missing_file_is_404(self):
|
|
prompt_id = self._save_prompt("to link")
|
|
|
|
resp = await self.client.request(
|
|
"POST",
|
|
self.URL,
|
|
json={
|
|
"prompt_id": prompt_id,
|
|
"image_path": str(self.output_dir / "no.png"),
|
|
},
|
|
)
|
|
|
|
self.assertEqual(resp.status, 404)
|
|
|
|
|
|
def _raise(*_args, **_kwargs):
|
|
raise RuntimeError("boom")
|
|
|
|
|
|
class ImageRouteCoverageCase(ImageAPITestCase):
|
|
"""Helpers for driving success and failure branches of each route."""
|
|
|
|
def _break_db(self, db_method):
|
|
setattr(self.api.db, db_method, _raise)
|
|
|
|
async def _json(self, method, path, **kwargs):
|
|
resp = await self.client.request(method, path, **kwargs)
|
|
return resp.status, await resp.json()
|
|
|
|
|
|
class TestPathHelpers(unittest.TestCase):
|
|
|
|
def test_has_traversal(self):
|
|
self.assertTrue(images_module._has_traversal(""))
|
|
self.assertTrue(images_module._has_traversal("a/../b.png"))
|
|
self.assertTrue(images_module._has_traversal("\\\\srv\\share\\x.png"))
|
|
self.assertTrue(images_module._has_traversal(os.path.abspath(os.sep)))
|
|
self.assertFalse(images_module._has_traversal("sub/x.png"))
|
|
|
|
|
|
class TestPromptImagesAndSearch(ImageRouteCoverageCase):
|
|
|
|
async def test_prompt_images_success_and_nan_cleanup(self):
|
|
png = make_png(self.output_dir / "p.png")
|
|
self._link_image(png, text="owner")
|
|
prompt_id = self.api.db.get_prompt_by_hash(generate_prompt_hash("owner"))["id"]
|
|
|
|
status, data = await self._json(
|
|
"GET", f"/prompt_manager/prompts/{prompt_id}/images"
|
|
)
|
|
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(len(data["images"]), 1)
|
|
self.assertTrue(data["images"][0]["url"].endswith("p.png"))
|
|
|
|
self.api.db.get_prompt_images = lambda _pid: [
|
|
{"id": 1, "image_path": str(png), "width": float("nan")}
|
|
]
|
|
status, data = await self._json("GET", "/prompt_manager/prompts/1/images")
|
|
self.assertEqual(status, 200)
|
|
self.assertIsNone(data["images"][0]["width"])
|
|
|
|
async def test_prompt_images_db_failure(self):
|
|
self._break_db("get_prompt_images")
|
|
status, data = await self._json("GET", "/prompt_manager/prompts/1/images")
|
|
self.assertEqual(status, 500)
|
|
self.assertFalse(data["success"])
|
|
|
|
async def test_search_images_paths(self):
|
|
self._link_image(make_png(self.output_dir / "s.png"), text="sunset beach")
|
|
missing, _ = await self._json("GET", "/prompt_manager/images/search")
|
|
ok, data = await self._json("GET", "/prompt_manager/images/search?q=sunset")
|
|
self._break_db("search_images_by_prompt")
|
|
err, _ = await self._json("GET", "/prompt_manager/images/search?q=x")
|
|
self.assertEqual((missing, ok, err), (400, 200, 500))
|
|
self.assertEqual(data["query"], "sunset")
|
|
self.assertEqual(len(data["images"]), 1)
|
|
|
|
async def test_list_routes_db_failures(self):
|
|
self._break_db("get_recent_images")
|
|
self._break_db("get_all_images")
|
|
recent, _ = await self._json("GET", "/prompt_manager/images/recent")
|
|
every, _ = await self._json("GET", "/prompt_manager/images/all")
|
|
self.assertEqual((recent, every), (500, 500))
|
|
|
|
|
|
class TestServeErrorBranches(ImageRouteCoverageCase):
|
|
|
|
async def test_serve_by_id_db_failure(self):
|
|
self._break_db("get_image_by_id")
|
|
status, _ = await self._json("GET", "/prompt_manager/images/1/file")
|
|
self.assertEqual(status, 500)
|
|
|
|
async def test_serve_output_lookup_failure(self):
|
|
self.api._get_all_output_dirs = _raise
|
|
status, _ = await self._json("GET", "/prompt_manager/images/serve/x.png")
|
|
self.assertEqual(status, 500)
|
|
|
|
async def test_out_of_range_root_index_is_rejected(self):
|
|
make_png(self.output_dir / "r.png")
|
|
resp = await self.client.request(
|
|
"GET", "/prompt_manager/images/serve/r.png?root=7"
|
|
)
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_directory_named_like_media_is_404(self):
|
|
(self.output_dir / "dir.png").mkdir()
|
|
status, _ = await self._json("GET", "/prompt_manager/images/serve/dir.png")
|
|
self.assertEqual(status, 404)
|
|
|
|
async def test_lora_dirs_extend_allowed_roots(self):
|
|
from py.config import IntegrationConfig
|
|
|
|
with (
|
|
tempfile.TemporaryDirectory() as lora,
|
|
tempfile.TemporaryDirectory() as cache,
|
|
):
|
|
png = make_png(Path(lora) / "preview.png")
|
|
image_id = self._link_image(png)
|
|
# An empty configured path means auto-detection; a developer's
|
|
# own config.json may hold a real LoraManager path, so pin it.
|
|
with (
|
|
patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True),
|
|
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", ""),
|
|
patch("py.lora_utils.find_lora_directories", return_value=[lora]),
|
|
patch("py.lora_utils.get_lora_image_cache_dir", return_value=cache),
|
|
):
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
self.assertEqual(resp.status, 200)
|
|
|
|
with (
|
|
patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True),
|
|
patch.object(IntegrationConfig, "LORA_MANAGER_PATH", ""),
|
|
patch("py.lora_utils.find_lora_directories", side_effect=RuntimeError),
|
|
):
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
|
|
class TestLoraManagerPathAllowList(ImageRouteCoverageCase):
|
|
"""A hand-edited LoraManager path must not widen what serve_image reads."""
|
|
|
|
def _enter_lora(self, stack, lora, cache, path, resolved):
|
|
from py.config import IntegrationConfig
|
|
|
|
stack.enter_context(
|
|
patch.object(IntegrationConfig, "LORA_MANAGER_ENABLED", True)
|
|
)
|
|
stack.enter_context(patch.object(IntegrationConfig, "LORA_MANAGER_PATH", path))
|
|
stack.enter_context(
|
|
patch("py.lora_utils.find_lora_directories", return_value=[lora])
|
|
)
|
|
stack.enter_context(
|
|
patch("py.lora_utils.get_lora_image_cache_dir", return_value=cache)
|
|
)
|
|
return stack.enter_context(
|
|
patch(
|
|
"py.lora_utils.resolve_lora_manager_path",
|
|
return_value=resolved,
|
|
create=True,
|
|
)
|
|
)
|
|
|
|
async def test_configured_path_outside_custom_nodes_is_ignored(self):
|
|
with ExitStack() as stack:
|
|
lora = stack.enter_context(tempfile.TemporaryDirectory())
|
|
cache = stack.enter_context(tempfile.TemporaryDirectory())
|
|
image_id = self._link_image(make_png(Path(lora) / "preview.png"))
|
|
self._enter_lora(stack, lora, cache, lora, None)
|
|
|
|
self.assertEqual(self.api._lora_image_dirs(), [])
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
self.assertEqual(resp.status, 403)
|
|
|
|
async def test_configured_custom_nodes_child_is_served(self):
|
|
configured = "custom_nodes/ComfyUI-Lora-Manager"
|
|
with ExitStack() as stack:
|
|
lora = stack.enter_context(tempfile.TemporaryDirectory())
|
|
cache = stack.enter_context(tempfile.TemporaryDirectory())
|
|
image_id = self._link_image(make_png(Path(lora) / "preview.png"))
|
|
resolve = self._enter_lora(stack, lora, cache, configured, lora)
|
|
|
|
resp = await self.client.request(
|
|
"GET", f"/prompt_manager/images/{image_id}/file"
|
|
)
|
|
self.assertEqual(resp.status, 200)
|
|
await resp.read() # release the served file before the directory goes
|
|
resolve.assert_called_once_with(configured)
|
|
|
|
|
|
class TestLinkImageRelativePaths(ImageAPITestCase):
|
|
"""POST /images/link accepts the path forms the API hands out."""
|
|
|
|
URL = "/prompt_manager/images/link"
|
|
|
|
async def setUpAsync(self):
|
|
await super().setUpAsync()
|
|
self.png = make_png(self.output_dir / "sub" / "gen.png")
|
|
|
|
async def _link(self, image_path):
|
|
prompt_id = self._save_prompt(f"link {image_path}")
|
|
resp = await self.client.request(
|
|
"POST", self.URL, json={"prompt_id": prompt_id, "image_path": image_path}
|
|
)
|
|
return resp.status, await resp.json()
|
|
|
|
async def test_root_relative_and_base_relative_paths_link(self):
|
|
for form in ("sub/gen.png", "output/sub/gen.png"):
|
|
status, data = await self._link(form)
|
|
|
|
self.assertEqual(status, 200, form)
|
|
image = self.api.db.get_image_by_id(data["image_id"])
|
|
self.assertEqual(Path(image["image_path"]), self.png, form)
|
|
|
|
async def test_relative_traversal_is_forbidden(self):
|
|
make_png(self.comfy_dir / "secret.png") # removed with the temp tree
|
|
|
|
status, _ = await self._link("../secret.png")
|
|
|
|
self.assertEqual(status, 403)
|
|
|
|
async def test_unknown_relative_path_is_404(self):
|
|
status, _ = await self._link("sub/missing.png")
|
|
|
|
self.assertEqual(status, 404)
|
|
|
|
|
|
class TestGalleryListing(ImageRouteCoverageCase):
|
|
|
|
async def test_subfolders_with_and_without_ancestors(self):
|
|
make_png(self.output_dir / "a" / "b" / "x.png")
|
|
make_png(self.output_dir / "c" / "y.png")
|
|
plain, data = await self._json("GET", "/prompt_manager/gallery/subfolders")
|
|
_, with_anc = await self._json(
|
|
"GET", "/prompt_manager/gallery/subfolders?include_ancestors=true"
|
|
)
|
|
self.assertEqual(plain, 200)
|
|
self.assertEqual(data["subfolders"], [os.path.join("a", "b"), "c"])
|
|
self.assertIn("a", with_anc["subfolders"])
|
|
|
|
async def test_subfolders_and_output_lookup_failures(self):
|
|
self.api._get_all_output_dirs = _raise
|
|
sub, _ = await self._json("GET", "/prompt_manager/gallery/subfolders")
|
|
out, _ = await self._json("GET", "/prompt_manager/images/output")
|
|
self.assertEqual((sub, out), (500, 500))
|
|
|
|
async def test_vanished_file_is_dropped_from_page(self):
|
|
png = make_png(self.output_dir / "gone.png")
|
|
await self._json("GET", "/prompt_manager/images/output")
|
|
png.unlink()
|
|
status, data = await self._json("GET", "/prompt_manager/images/output")
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["total"], 1)
|
|
self.assertEqual(data["images"], [])
|
|
|
|
def test_gallery_scan_skips_unreadable_entries(self):
|
|
ghost = self.output_dir / "ghost.png"
|
|
with patch.object(
|
|
images_module, "_iter_media_files", return_value=iter([ghost])
|
|
):
|
|
found = self.api._scan_gallery_files_sync(self.output_dir)
|
|
self.assertEqual(found, [])
|
|
|
|
|
|
class TestThumbnailUnits(ImageRouteCoverageCase):
|
|
|
|
SIZE = ThumbnailSettings("webp", 80, 64)
|
|
|
|
def test_generate_one_reports_vanished_source(self):
|
|
result = self.api._generate_one(
|
|
self.output_dir / "nope.png", self._dst("nope"), self.SIZE, False
|
|
)
|
|
self.assertEqual(result["action"], "error")
|
|
self.assertIn("no longer exists", result["error"])
|
|
|
|
def test_generate_one_video_success_and_failure(self):
|
|
src = self.output_dir / "clip.mp4"
|
|
src.write_bytes(b"\x00")
|
|
with patch.object(self.api, "_generate_video_thumbnail", return_value=True):
|
|
ok = self.api._generate_one(src, self._dst("clip"), self.SIZE, True)
|
|
with patch.object(self.api, "_generate_video_thumbnail", return_value=False):
|
|
bad = self.api._generate_one(src, self._dst("clip"), self.SIZE, True)
|
|
self.assertEqual((ok["action"], ok["type"]), ("generated", "video"))
|
|
self.assertEqual(bad["action"], "error")
|
|
|
|
def test_generate_one_unexpected_error_is_safe(self):
|
|
src = make_png(self.output_dir / "boom.png")
|
|
with patch.object(
|
|
images_module, "_write_image_thumbnail", side_effect=RuntimeError("bad")
|
|
):
|
|
result = self.api._generate_one(src, self._dst("boom"), self.SIZE, False)
|
|
self.assertEqual(result["action"], "error")
|
|
self.assertEqual(result["error"], "bad")
|
|
|
|
def test_rgba_and_jpeg_sources(self):
|
|
with Image.new("RGBA", (32, 32), (0, 255, 0, 128)) as img:
|
|
img.save(self.output_dir / "alpha.png")
|
|
with Image.new("RGB", (32, 32), (0, 0, 255)) as img:
|
|
img.save(self.output_dir / "photo.jpg", "JPEG")
|
|
result = self.api._generate_thumbnails_sync(
|
|
self.output_dir, self.output_dir / "thumbnails", self.SIZE
|
|
)
|
|
self.assertEqual(result["count"], 2)
|
|
self.assertTrue(self._dst("alpha").is_file())
|
|
self.assertTrue(self._dst("photo").is_file())
|
|
|
|
async def test_emit_sse_swallows_write_errors(self):
|
|
response = MagicMock()
|
|
response.write = AsyncMock(side_effect=ConnectionResetError("gone"))
|
|
await self.api._emit_sse(response, "progress", {"n": 1})
|
|
response.write.assert_awaited_once()
|
|
|
|
async def test_progress_stream_reports_internal_failure(self):
|
|
self.api._plan_thumbnails = _raise
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails/progress"
|
|
)
|
|
events = parse_sse(await resp.text())
|
|
self.assertEqual(events[-1][0], "error")
|
|
|
|
async def test_post_generate_reports_internal_failure(self):
|
|
self.api._get_all_output_dirs = _raise
|
|
status, _ = await self._json(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
self.assertEqual(status, 500)
|
|
|
|
def _dst(self, name, suffix=".webp"):
|
|
return self.output_dir / "thumbnails" / f"{name}_thumb{suffix}"
|
|
|
|
|
|
class TestVideoThumbnailFallbacks(ImageRouteCoverageCase):
|
|
"""_generate_video_thumbnail degrades: cv2 -> ffmpeg -> placeholder."""
|
|
|
|
# JPEG keeps these checks format-specific; #317 settings carry the edge.
|
|
SIZE = ThumbnailSettings("jpeg", 85, 48)
|
|
|
|
@property
|
|
def video(self):
|
|
# output_dir only exists once the app is built, so create lazily.
|
|
path = self.output_dir / "clip.mp4"
|
|
if not path.exists():
|
|
path.write_bytes(b"\x00\x00")
|
|
return path
|
|
|
|
@property
|
|
def thumb(self):
|
|
return self.output_dir / "clip_thumb.jpg"
|
|
|
|
def _fake_cv2(self, opened=True, read_ok=True):
|
|
numpy = self._numpy()
|
|
cv2 = MagicMock()
|
|
cap = cv2.VideoCapture.return_value
|
|
cap.isOpened.return_value = opened
|
|
cap.get.return_value = 30
|
|
frame = numpy.zeros((8, 8, 3), dtype=numpy.uint8)
|
|
cap.read.return_value = (read_ok, frame if read_ok else None)
|
|
cv2.cvtColor.side_effect = lambda f, _code: f
|
|
return cv2
|
|
|
|
def _numpy(self):
|
|
try:
|
|
import numpy
|
|
except ImportError:
|
|
self.skipTest("numpy not installed")
|
|
return numpy
|
|
|
|
def test_cv2_success(self):
|
|
with patch.dict(sys.modules, {"cv2": self._fake_cv2()}):
|
|
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
|
self.assertTrue(ok)
|
|
with Image.open(self.thumb) as img:
|
|
self.assertEqual(img.format, "JPEG")
|
|
|
|
def test_cv2_cannot_open_or_read_falls_through_to_ffmpeg_then_placeholder(self):
|
|
# #249 I1: an OpenCV that cannot decode the clip (AV1) no longer ends
|
|
# the run; ffmpeg and then the placeholder are tried.
|
|
for fake in (self._fake_cv2(opened=False), self._fake_cv2(read_ok=False)):
|
|
self.thumb.unlink(missing_ok=True)
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": fake}),
|
|
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
|
):
|
|
ok = self.api._generate_video_thumbnail(
|
|
self.video, self.thumb, self.SIZE
|
|
)
|
|
self.assertTrue(ok)
|
|
with Image.open(self.thumb) as img:
|
|
self.assertEqual((img.format, img.size), ("JPEG", self.SIZE.size))
|
|
|
|
def test_cv2_crash_is_contained(self):
|
|
cv2 = MagicMock()
|
|
cv2.VideoCapture.side_effect = RuntimeError("driver crash")
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": cv2}),
|
|
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
|
):
|
|
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
|
self.assertTrue(ok, "the placeholder is written")
|
|
|
|
def test_cv2_failure_keeps_an_existing_preview_on_rebuild(self):
|
|
make_png(self.thumb)
|
|
before = self.thumb.read_bytes()
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": self._fake_cv2(opened=False)}),
|
|
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
|
):
|
|
ok = self.api._generate_video_thumbnail(
|
|
self.video, self.thumb, self.SIZE, keep_existing=True
|
|
)
|
|
self.assertFalse(ok)
|
|
self.assertEqual(self.thumb.read_bytes(), before)
|
|
|
|
def test_ffmpeg_absent_yields_placeholder(self):
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": None}),
|
|
patch("subprocess.Popen", side_effect=FileNotFoundError("ffmpeg")),
|
|
):
|
|
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
|
self.assertTrue(ok)
|
|
with Image.open(self.thumb) as img:
|
|
self.assertEqual(img.size, self.SIZE.size)
|
|
|
|
@staticmethod
|
|
def _process(returncode, stderr=b"", stdout=b""):
|
|
process = MagicMock(returncode=returncode)
|
|
process.communicate.return_value = (stdout, stderr)
|
|
return process
|
|
|
|
@staticmethod
|
|
def _png_frame(size=(96, 64)):
|
|
import io
|
|
|
|
buffer = io.BytesIO()
|
|
Image.new("RGB", size, (200, 0, 0)).save(buffer, "PNG")
|
|
return buffer.getvalue()
|
|
|
|
def test_ffmpeg_failure_yields_placeholder(self):
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": None}),
|
|
patch("subprocess.Popen", return_value=self._process(1, b"no stream")),
|
|
):
|
|
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
|
self.assertTrue(ok)
|
|
self.assertTrue(self.thumb.is_file())
|
|
|
|
def test_ffmpeg_success(self):
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": None}),
|
|
patch(
|
|
"subprocess.Popen",
|
|
return_value=self._process(0, stdout=self._png_frame()),
|
|
) as run,
|
|
):
|
|
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
|
self.assertTrue(ok)
|
|
self.assertEqual(run.call_args.args[0][0], "ffmpeg")
|
|
with Image.open(self.thumb) as img:
|
|
# ffmpeg's frame (a PNG on stdout) fitted inside the edge (#317).
|
|
self.assertEqual((img.format, img.size), ("JPEG", (48, 32)))
|
|
|
|
def test_ffmpeg_is_tracked_and_a_hold_before_the_rename_writes_nothing(self):
|
|
tracked = []
|
|
|
|
def writing(cmd, **kwargs):
|
|
return self._process(0, stdout=self._png_frame())
|
|
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": None}),
|
|
patch("subprocess.Popen", side_effect=writing),
|
|
):
|
|
with self.assertRaises(images_module.PreviewWriteHeld):
|
|
self.api._generate_video_thumbnail(
|
|
self.video,
|
|
self.thumb,
|
|
self.SIZE,
|
|
guard=lambda: False,
|
|
track=tracked.append,
|
|
)
|
|
self.assertEqual(len(tracked), 1)
|
|
self.assertFalse(self.thumb.exists())
|
|
self.assertEqual(
|
|
[p.name for p in self.thumb.parent.iterdir() if "part" in p.name], []
|
|
)
|
|
|
|
def test_placeholder_failure_returns_false(self):
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": None}),
|
|
patch("subprocess.Popen", side_effect=FileNotFoundError),
|
|
patch.object(
|
|
images_module.Image, "new", side_effect=RuntimeError("no PIL")
|
|
),
|
|
):
|
|
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
|
self.assertFalse(ok)
|
|
|
|
|
|
class TestClearThumbnails(ImageRouteCoverageCase):
|
|
|
|
URL = "/prompt_manager/images/clear-thumbnails"
|
|
|
|
async def test_missing_output_dir_is_404(self):
|
|
self.api._get_all_output_dirs = lambda: []
|
|
status, _ = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 404)
|
|
|
|
async def test_no_thumbnails_dir_is_noop(self):
|
|
status, data = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["cleared_files"], 0)
|
|
|
|
async def test_clears_only_generated_thumbnails(self):
|
|
thumbs = self.output_dir / "thumbnails"
|
|
make_png(thumbs / "sub" / "a_thumb.png")
|
|
make_png(thumbs / "b_thumb.png")
|
|
(thumbs / "readme.txt").write_text("keep")
|
|
make_png(self.output_dir / "original.png")
|
|
|
|
status, data = await self._json("POST", self.URL)
|
|
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["cleared_files"], 2)
|
|
self.assertIn("cleared_size_formatted", data)
|
|
self.assertFalse((thumbs / "sub").exists())
|
|
self.assertTrue((thumbs / "readme.txt").is_file())
|
|
self.assertTrue((self.output_dir / "original.png").is_file())
|
|
|
|
async def test_unlink_failure_is_logged_not_raised(self):
|
|
make_png(self.output_dir / "thumbnails" / "a_thumb.png")
|
|
with patch.object(Path, "unlink", side_effect=PermissionError("locked")):
|
|
status, data = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["cleared_files"], 0)
|
|
|
|
async def test_internal_failure_is_500(self):
|
|
self.api._get_all_output_dirs = _raise
|
|
status, _ = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 500)
|
|
|
|
async def test_sibling_prefix_directory_is_refused(self):
|
|
# "<base>/out2/thumbnails" starts with "<base>/out" as a string but is
|
|
# not inside it; a symlinked thumbnails dir pointing there must be
|
|
# refused rather than emptied.
|
|
out = self.comfy_dir / "out"
|
|
out.mkdir()
|
|
sibling = make_png(self.comfy_dir / "out2" / "thumbnails" / "a_thumb.png")
|
|
self._make_symlink(sibling.parent, out / "thumbnails")
|
|
self.api._get_all_output_dirs = lambda: [out]
|
|
|
|
status, data = await self._json("POST", self.URL)
|
|
|
|
self.assertEqual(status, 400)
|
|
self.assertFalse(data["success"])
|
|
self.assertTrue(sibling.is_file())
|
|
|
|
async def test_clears_all_roots_and_combines_counts_without_touching_originals(
|
|
self,
|
|
):
|
|
second = self.comfy_dir / "extra-output"
|
|
roots = [self.output_dir, second]
|
|
self.api._get_all_output_dirs = lambda: roots
|
|
expected_size = 0
|
|
for root in roots:
|
|
make_png(root / "original_thumb.png")
|
|
thumb = make_png(root / "thumbnails" / "sub" / "generated_thumb.png")
|
|
expected_size += thumb.stat().st_size
|
|
(root / "thumbnails" / "notes.txt").write_text("keep")
|
|
status, data = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["cleared_files"], 2)
|
|
self.assertEqual(data["cleared_size"], expected_size)
|
|
self.assertNotIn("{format_size", data["message"])
|
|
for root in roots:
|
|
self.assertTrue((root / "original_thumb.png").is_file())
|
|
self.assertTrue((root / "thumbnails" / "notes.txt").is_file())
|
|
self.assertFalse((root / "thumbnails" / "sub").exists())
|
|
|
|
async def test_missing_thumbnails_in_first_root_does_not_skip_later_roots(self):
|
|
second = self.comfy_dir / "extra-output"
|
|
thumb = make_png(second / "thumbnails" / "a_thumb.png")
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
|
|
status, data = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["cleared_files"], 1)
|
|
self.assertFalse(thumb.exists())
|
|
|
|
async def test_invalid_later_root_is_rejected_before_deleting_any_files(self):
|
|
first_thumb = make_png(self.output_dir / "thumbnails" / "a_thumb.png")
|
|
second = self.comfy_dir / "extra-output"
|
|
second.mkdir()
|
|
outside = make_png(self.comfy_dir / "elsewhere" / "a_thumb.png")
|
|
self._make_symlink(outside.parent, second / "thumbnails")
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
|
|
status, data = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 400)
|
|
self.assertFalse(data["success"])
|
|
self.assertTrue(first_thumb.is_file())
|
|
self.assertTrue(outside.is_file())
|
|
|
|
async def test_thumbnail_symlink_back_to_root_cannot_delete_originals(self):
|
|
original = make_png(self.output_dir / "original_thumb.png")
|
|
self._make_symlink(self.output_dir, self.output_dir / "thumbnails")
|
|
status, data = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 400)
|
|
self.assertFalse(data["success"])
|
|
self.assertTrue(original.is_file())
|
|
|
|
async def test_cleanup_runs_off_loop_and_invalidates_gallery_cache(self):
|
|
import threading
|
|
|
|
make_png(self.output_dir / "thumbnails" / "a_thumb.png")
|
|
loop_thread = threading.get_ident()
|
|
deletion_threads = []
|
|
original_unlink = Path.unlink
|
|
|
|
def unlink(path, *args, **kwargs):
|
|
deletion_threads.append(threading.get_ident())
|
|
return original_unlink(path, *args, **kwargs)
|
|
|
|
self.api._gallery_cache = {"previous": object()}
|
|
with patch.object(Path, "unlink", unlink):
|
|
status, data = await self._json("POST", self.URL)
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["cleared_files"], 1)
|
|
self.assertTrue(deletion_threads)
|
|
self.assertNotIn(loop_thread, deletion_threads)
|
|
self.assertEqual(self.api._gallery_cache, {})
|
|
|
|
|
|
class TestImagePromptLookup(ImageRouteCoverageCase):
|
|
|
|
URL = "/prompt_manager/images/prompt/"
|
|
|
|
async def test_relative_path_resolves_against_output_dir(self):
|
|
self._link_image(make_png(self.output_dir / "gen.png"), text="the prompt")
|
|
status, data = await self._json("GET", f"{self.URL}gen.png")
|
|
self.assertEqual(status, 200)
|
|
self.assertTrue(data["success"])
|
|
self.assertEqual(data["prompt"]["text"], "the prompt")
|
|
self.assertEqual(data["prompt"]["image_path"], "output/gen.png")
|
|
|
|
async def test_unknown_image_reports_no_prompt(self):
|
|
status, data = await self._json("GET", f"{self.URL}unknown.png")
|
|
self.assertEqual(status, 200)
|
|
self.assertFalse(data["success"])
|
|
self.assertIn("image_path", data)
|
|
|
|
async def test_empty_path_is_400(self):
|
|
status, _ = await self._json("GET", self.URL)
|
|
self.assertEqual(status, 400)
|
|
|
|
async def test_db_failure_is_500_without_details(self):
|
|
self._break_db("get_image_prompt_info")
|
|
status, data = await self._json("GET", f"{self.URL}x.png")
|
|
self.assertEqual(status, 500)
|
|
self.assertEqual(data["error"], "Database error occurred")
|
|
|
|
async def test_output_dir_lookup_failure_is_500(self):
|
|
self.api._find_comfyui_output_dir = _raise
|
|
status, _ = await self._json("GET", f"{self.URL}x.png")
|
|
self.assertEqual(status, 500)
|
|
|
|
|
|
class TestPublicPathsInResponses(ImageAPITestCase):
|
|
"""Image routes describe files relative to the ComfyUI tree, never absolutely."""
|
|
|
|
async def setUpAsync(self):
|
|
await super().setUpAsync()
|
|
self.png = make_png(self.output_dir / "sub" / "gen.png")
|
|
self.prompt_id = self._save_prompt("public paths")
|
|
self.assertGreater(
|
|
self.api.db.link_image_to_prompt(self.prompt_id, str(self.png)), 0
|
|
)
|
|
|
|
async def _body(self, method, path, **kwargs):
|
|
resp = await self.client.request(method, path, **kwargs)
|
|
text = await resp.text()
|
|
self.assertNotIn(str(self.comfy_dir), text, path)
|
|
return resp.status, json.loads(text)
|
|
|
|
async def test_output_listing_paths_are_root_relative(self):
|
|
status, data = await self._body("GET", "/prompt_manager/images/output")
|
|
|
|
self.assertEqual(status, 200)
|
|
entry = data["images"][0]
|
|
self.assertEqual(entry["path"], "sub/gen.png")
|
|
self.assertEqual(entry["root_dir"], "output")
|
|
self.assertEqual(entry["root_index"], 0)
|
|
|
|
async def test_image_lists_publish_image_path_and_drop_file_path(self):
|
|
routes = (
|
|
f"/prompt_manager/prompts/{self.prompt_id}/images",
|
|
"/prompt_manager/images/recent",
|
|
"/prompt_manager/images/all",
|
|
"/prompt_manager/images/search?q=public",
|
|
)
|
|
for route in routes:
|
|
status, data = await self._body("GET", route)
|
|
|
|
self.assertEqual(status, 200, route)
|
|
image = data["images"][0]
|
|
self.assertEqual(image["image_path"], "output/sub/gen.png", route)
|
|
self.assertNotIn("file_path", image, route)
|
|
self.assertTrue(image["url"].endswith("sub/gen.png"), route)
|
|
|
|
async def test_image_prompt_lookup_publishes_path(self):
|
|
_, found = await self._body("GET", "/prompt_manager/images/prompt/sub/gen.png")
|
|
_, missing = await self._body("GET", "/prompt_manager/images/prompt/sub/no.png")
|
|
|
|
self.assertEqual(found["prompt"]["image_path"], "output/sub/gen.png")
|
|
self.assertEqual(missing["image_path"], "output/sub/no.png")
|
|
|
|
async def test_thumbnails_path_is_public(self):
|
|
status, data = await self._body(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
|
|
self.assertEqual(status, 200)
|
|
self.assertEqual(data["thumbnails_path"], "output/thumbnails")
|
|
|
|
async def test_oserror_bodies_keep_only_the_basename(self):
|
|
secret = str(self.comfy_dir / "hidden.png")
|
|
|
|
def boom(*_args, **_kwargs):
|
|
raise FileNotFoundError(2, "No such file", secret)
|
|
|
|
self.api._get_all_output_dirs = boom
|
|
self.api._find_comfyui_output_dir = boom
|
|
self.api.db.get_prompt_images = boom
|
|
self.api.db.delete_image = boom
|
|
routes = (
|
|
("GET", "/prompt_manager/gallery/subfolders"),
|
|
("GET", "/prompt_manager/prompts/1/images"),
|
|
("DELETE", "/prompt_manager/images/1"),
|
|
("POST", "/prompt_manager/images/clear-thumbnails"),
|
|
("GET", "/prompt_manager/images/prompt/x.png"),
|
|
)
|
|
for method, route in routes:
|
|
status, data = await self._body(method, route)
|
|
|
|
self.assertEqual(status, 500, route)
|
|
self.assertIn("hidden.png", data["error"], route)
|
|
|
|
|
|
class TestDeleteAndLinkErrors(ImageRouteCoverageCase):
|
|
|
|
async def test_delete_image_paths(self):
|
|
image_id = self._link_image(make_png(self.output_dir / "d.png"))
|
|
ok, _ = await self._json("DELETE", f"/prompt_manager/images/{image_id}")
|
|
gone, _ = await self._json("DELETE", f"/prompt_manager/images/{image_id}")
|
|
bad, _ = await self._json("DELETE", "/prompt_manager/images/abc")
|
|
self._break_db("delete_image")
|
|
err, _ = await self._json("DELETE", "/prompt_manager/images/1")
|
|
self.assertEqual((ok, gone, bad, err), (200, 404, 400, 500))
|
|
|
|
async def test_link_db_failure_is_500(self):
|
|
prompt_id = self._save_prompt("to link")
|
|
png = make_png(self.output_dir / "l.png")
|
|
self._break_db("link_image_to_prompt")
|
|
status, _ = await self._json(
|
|
"POST",
|
|
"/prompt_manager/images/link",
|
|
json={"prompt_id": prompt_id, "image_path": str(png)},
|
|
)
|
|
self.assertEqual(status, 500)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|
|
|
|
class TestGalleryRescan(ImageAPITestCase):
|
|
"""POST /prompt_manager/gallery/rescan drops the listing cache and rewalks."""
|
|
|
|
URL = "/prompt_manager/images/output"
|
|
|
|
async def _json(self, method, path, **kwargs):
|
|
resp = await self.client.request(method, path, **kwargs)
|
|
return resp.status, await resp.json()
|
|
|
|
async def test_rescan_makes_new_files_visible_before_the_cache_expires(self):
|
|
self.api._gallery_cache_ttl = 3600 # the cache alone would hide the new file
|
|
make_png(self.output_dir / "a.png")
|
|
_, first = await self._json("GET", self.URL)
|
|
self.assertEqual(first["total"], 1)
|
|
|
|
make_png(self.output_dir / "b.png")
|
|
_, stale = await self._json("GET", self.URL)
|
|
self.assertEqual(stale["total"], 1)
|
|
|
|
status, data = await self._json("POST", "/prompt_manager/gallery/rescan")
|
|
_, fresh = await self._json("GET", self.URL)
|
|
|
|
self.assertEqual(status, 200)
|
|
self.assertTrue(data["success"])
|
|
self.assertEqual(data["total"], 2)
|
|
self.assertEqual(data["roots"], 1)
|
|
self.assertEqual(fresh["total"], 2)
|
|
|
|
async def test_rescan_is_post_only(self):
|
|
resp = await self.client.request("GET", "/prompt_manager/gallery/rescan")
|
|
self.assertEqual(resp.status, 405)
|
|
|
|
async def test_rescan_reports_internal_failure(self):
|
|
self.api._scan_gallery_files_sync = _raise
|
|
status, data = await self._json("POST", "/prompt_manager/gallery/rescan")
|
|
self.assertEqual(status, 500)
|
|
self.assertFalse(data["success"])
|
|
self.assertNotIn("boom", json.dumps(data))
|
|
|
|
|
|
class TestOutputScanHelpers(ImageAPITestCase):
|
|
"""Stable ids for filesystem entries and no symlink reads.
|
|
|
|
The thumbnail stream no longer stops when its client goes away: see
|
|
``tests/test_thumbnail_jobs.py`` (the work is a job that outlives it).
|
|
"""
|
|
|
|
def test_output_entry_id_is_a_stable_digest_of_the_path(self):
|
|
import hashlib
|
|
|
|
output = Path(self.output_dir)
|
|
media = output / "stable.png"
|
|
make_png(media)
|
|
entry = images_module._output_image_entry(media, output, 0)
|
|
expected = hashlib.sha1(str(media).encode("utf-8")).hexdigest()[:16]
|
|
self.assertEqual(entry["id"], expected)
|
|
|
|
def test_iter_media_files_skips_symlinked_files(self):
|
|
output = Path(self.output_dir)
|
|
make_png(output / "real.png")
|
|
with tempfile.TemporaryDirectory() as other:
|
|
outside = Path(other) / "outside.png"
|
|
make_png(outside)
|
|
try:
|
|
os.symlink(outside, output / "link.png")
|
|
except (OSError, NotImplementedError):
|
|
self.skipTest("symlinks not available")
|
|
names = sorted(p.name for p in images_module._iter_media_files(output, 100))
|
|
self.assertEqual(names, ["real.png"])
|
|
|
|
|
|
class TestThumbnailStartupStatus(ThumbnailTestCase):
|
|
"""Startup counts and generation cover the same configured collection."""
|
|
|
|
STATUS_URL = "/prompt_manager/images/thumbnails/status"
|
|
|
|
async def test_status_reports_exact_pending_count_without_generating(self):
|
|
for name in ("ready", "missing", "outdated"):
|
|
make_png(self.output_dir / f"{name}.png")
|
|
thumbs = self.output_dir / "thumbnails"
|
|
ready = make_png(thumbs / "ready_thumb.png")
|
|
outdated = make_png(thumbs / "outdated_thumb.png")
|
|
source_time = (self.output_dir / "ready.png").stat().st_mtime
|
|
os.utime(ready, (source_time + 10, source_time + 10))
|
|
os.utime(outdated, (source_time - 10, source_time - 10))
|
|
calls = self._spy_executor()
|
|
|
|
response = await self.client.get(self.STATUS_URL)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
data = await response.json()
|
|
self.assertTrue(data["success"])
|
|
self.assertEqual(data["total_images"], 3)
|
|
self.assertEqual(data["pending"], 2)
|
|
self.assertEqual(data["ready"], 1)
|
|
self.assertTrue(data["collection_id"])
|
|
self.assertIn("_plan_gallery_thumbnails", calls)
|
|
self.assertFalse((thumbs / "missing_thumb.png").exists())
|
|
|
|
async def test_both_generation_routes_cover_all_gallery_roots(self):
|
|
second = self.comfy_dir / "another-output"
|
|
make_png(self.output_dir / "a.png")
|
|
make_png(second / "b.png")
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
|
|
for route in (
|
|
"/prompt_manager/images/generate-thumbnails",
|
|
self.PROGRESS_URL,
|
|
):
|
|
with self.subTest(route=route):
|
|
for root in (self.output_dir, second):
|
|
for thumb in (root / "thumbnails").glob("*_thumb.webp"):
|
|
thumb.unlink()
|
|
response = await self.client.post(route, json={})
|
|
if route == self.PROGRESS_URL:
|
|
data = dict(parse_sse(await response.text()))["complete"]
|
|
else:
|
|
data = await response.json()
|
|
self.assertEqual(data["count"], 2)
|
|
self.assertTrue((second / "thumbnails" / "b_thumb.webp").exists())
|
|
status = await (await self.client.get(self.STATUS_URL)).json()
|
|
self.assertEqual(status["pending"], 0)
|
|
self.assertEqual(status["total_images"], 2)
|
|
|
|
async def test_pending_count_includes_files_beyond_the_generation_cap(self):
|
|
for name in ("a", "b", "c"):
|
|
make_png(self.output_dir / f"{name}.png")
|
|
with patch.object(images_module, "MAX_THUMBNAIL_FILES", 1):
|
|
response = await self.client.get(self.STATUS_URL)
|
|
self.assertEqual(response.status, 200)
|
|
data = await response.json()
|
|
self.assertEqual(data["pending"], 3)
|
|
|
|
async def test_generation_cap_is_shared_across_roots(self):
|
|
second = self.comfy_dir / "another-output"
|
|
make_png(self.output_dir / "a.png")
|
|
make_png(second / "b.png")
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
|
|
with patch.object(images_module, "MAX_THUMBNAIL_FILES", 1):
|
|
plan = self.api._plan_gallery_thumbnails([self.output_dir, second])
|
|
response = await self.client.post(self.PROGRESS_URL)
|
|
data = dict(parse_sse(await response.text()))["complete"]
|
|
# One batch holds one file across both roots; the run goes on (#317).
|
|
self.assertEqual((len(plan["targets"]), plan["remaining"]), (1, 1))
|
|
self.assertEqual(data["count"], 2)
|
|
self.assertEqual(data["remaining"], 0)
|
|
self.assertEqual(data["total_images"], 2)
|
|
|
|
async def test_equal_timestamps_do_not_regenerate_existing_thumbnails(self):
|
|
source = make_png(self.output_dir / "ready.png")
|
|
thumb = make_png(self.output_dir / "thumbnails" / "ready_thumb.png")
|
|
timestamp = int(source.stat().st_mtime)
|
|
os.utime(source, (timestamp, timestamp))
|
|
os.utime(thumb, (timestamp, timestamp))
|
|
response = await self.client.post(self.PROGRESS_URL)
|
|
data = dict(parse_sse(await response.text()))["complete"]
|
|
self.assertEqual(data["count"], 0)
|
|
self.assertEqual(data["skipped"], 1)
|
|
|
|
|
|
class TestGalleryRootIdentity(ThumbnailTestCase):
|
|
async def test_colliding_original_and_thumbnail_urls_serve_their_own_roots(self):
|
|
roots = [self.output_dir, self.comfy_dir / "second-output"]
|
|
for index, root in enumerate(roots):
|
|
make_png(root / "same.png", color=(index * 255, 0, 255 - index * 255))
|
|
self.api._get_all_output_dirs = lambda: roots
|
|
await self.client.post("/prompt_manager/images/generate-thumbnails", json={})
|
|
data = await (await self.client.get("/prompt_manager/images/output")).json()
|
|
self.assertEqual(len({image["url"] for image in data["images"]}), 2)
|
|
self.assertEqual(len({image["thumbnail_url"] for image in data["images"]}), 2)
|
|
for image in data["images"]:
|
|
root = roots[image["root_index"]]
|
|
for key, path in (
|
|
("url", root / "same.png"),
|
|
("thumbnail_url", root / "thumbnails" / "same_thumb.webp"),
|
|
):
|
|
response = await self.client.get(image[key])
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(await response.read(), path.read_bytes())
|
|
|
|
async def test_root_scoped_prompt_lookup_uses_the_matching_database_image(self):
|
|
roots = [self.output_dir, self.comfy_dir / "second-output"]
|
|
for index, root in enumerate(roots):
|
|
self._link_image(make_png(root / "same.png"), text=f"root {index} prompt")
|
|
self.api._get_all_output_dirs = lambda: roots
|
|
for index in range(2):
|
|
response = await self.client.get(
|
|
f"/prompt_manager/images/prompt/same.png?root={index}"
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
data = await response.json()
|
|
self.assertEqual(data["prompt"]["text"], f"root {index} prompt")
|
|
|
|
async def test_special_characters_in_paths_survive_serving_and_prompt_lookup(self):
|
|
from urllib.parse import quote
|
|
|
|
filenames = ["space #%20 café.png"]
|
|
if os.name != "nt":
|
|
filenames.append("space #?%20 café.png")
|
|
for filename in filenames:
|
|
self._link_image(make_png(self.output_dir / filename), text="special path")
|
|
for filename in filenames:
|
|
with self.subTest(filename=filename):
|
|
original = self.output_dir / filename
|
|
data = await (
|
|
await self.client.get("/prompt_manager/images/output")
|
|
).json()
|
|
entry = next(
|
|
image for image in data["images"] if image["filename"] == filename
|
|
)
|
|
response = await self.client.get(entry["url"])
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(await response.read(), original.read_bytes())
|
|
response = await self.client.get(
|
|
f"/prompt_manager/images/prompt/{quote(filename, safe='')}?root=0"
|
|
)
|
|
data = await response.json()
|
|
self.assertEqual(data["prompt"]["text"], "special path")
|
|
|
|
async def test_invalid_explicit_roots_never_fall_back_to_another_image(self):
|
|
self._link_image(make_png(self.output_dir / "same.png"), text="first root")
|
|
for root in ("bad", "-1", "7", ""):
|
|
for route in ("serve", "prompt"):
|
|
with self.subTest(root=root, route=route):
|
|
response = await self.client.get(
|
|
f"/prompt_manager/images/{route}/same.png?root={root}"
|
|
)
|
|
self.assertEqual(response.status, 403)
|
|
|
|
async def test_explicit_root_prompt_lookup_rejects_paths_outside_that_root(self):
|
|
from urllib.parse import quote
|
|
|
|
outside = make_png(self.comfy_dir / "outside.png")
|
|
self._link_image(outside, text="outside")
|
|
response = await self.client.get(
|
|
f"/prompt_manager/images/prompt/{quote(str(outside), safe='')}?root=0"
|
|
)
|
|
self.assertEqual(response.status, 403)
|
|
|
|
async def test_unlinked_image_does_not_inherit_a_prompt_from_another_root(self):
|
|
second = self.comfy_dir / "second-output"
|
|
self._link_image(make_png(self.output_dir / "same.png"), text="first root")
|
|
make_png(second / "same.png")
|
|
self.api._get_all_output_dirs = lambda: [self.output_dir, second]
|
|
response = await self.client.get(
|
|
"/prompt_manager/images/prompt/same.png?root=1"
|
|
)
|
|
self.assertEqual(response.status, 200)
|
|
self.assertFalse((await response.json())["success"])
|
|
|
|
|
|
class TestRootScopeServing(ImageAPITestCase):
|
|
async def scoped_request(self, path="image.png", root="0", scope="0" * 64):
|
|
params = {"root_scope": scope}
|
|
if root is not None:
|
|
params["root"] = root
|
|
return await self.client.get(
|
|
f"/prompt_manager/images/serve/{path}", params=params
|
|
)
|
|
|
|
async def assert_no_access(self, path, root="0", scope="0" * 64, status=409):
|
|
original_stat = Path.stat
|
|
|
|
def stat(candidate, *args, **kwargs):
|
|
if candidate.name in ("image.png", "image_thumb.png"):
|
|
raise AssertionError("candidate stat before scope guard")
|
|
return original_stat(candidate, *args, **kwargs)
|
|
|
|
with (
|
|
patch.object(Path, "stat", stat),
|
|
patch.object(
|
|
self.api, "_file_response", side_effect=AssertionError("file response")
|
|
),
|
|
patch.object(
|
|
self.api.db,
|
|
"get_image_prompt_info",
|
|
side_effect=AssertionError("saved lookup"),
|
|
),
|
|
patch(
|
|
"py.api.inspection_service.read_inspection_file",
|
|
side_effect=AssertionError("metadata read"),
|
|
),
|
|
patch(
|
|
"utils.prompt_inspector.subprocess.Popen",
|
|
side_effect=AssertionError("probe"),
|
|
),
|
|
):
|
|
response = await self.scoped_request(path, root, scope)
|
|
self.assertEqual(response.status, status)
|
|
result = await response.json()
|
|
self.assertFalse(result["success"])
|
|
if status == 409:
|
|
self.assertEqual(result["code"], "source-changed")
|
|
self.assertEqual(response.headers.get("Cache-Control"), "no-store")
|
|
self.assertNotIn(str(self.comfy_dir), json.dumps(result))
|
|
|
|
async def test_root_scope_stale_original_and_thumbnail_fail_before_access(self):
|
|
make_png(self.output_dir / "image.png")
|
|
make_png(self.output_dir / "thumbnails" / "image_thumb.png")
|
|
for path in ("image.png", "thumbnails/image_thumb.png"):
|
|
await self.assert_no_access(path)
|
|
|
|
async def test_root_scope_malformed_tokens_are_bad_requests(self):
|
|
make_png(self.output_dir / "image.png")
|
|
for scope in ("", "g" * 64, "a" * 63, "a" * 65, "a" * 4097, "A" * 64):
|
|
await self.assert_no_access("image.png", scope=scope, status=400)
|
|
for root in (None, "", "x", "-1", "0.0", "+0"):
|
|
await self.assert_no_access("image.png", root=root, status=400)
|
|
|
|
async def test_root_scope_rotated_key_and_removed_index_require_refresh(self):
|
|
make_png(self.output_dir / "image.png")
|
|
listing = await self.client.get("/prompt_manager/images/output?mode=workspace")
|
|
descriptor = (await listing.json())["images"][0]
|
|
self.assertRegex(descriptor.get("root_scope", ""), r"^[0-9a-f]{64}$")
|
|
with patch.object(
|
|
images_module, "_WORKSPACE_CURSOR_SECRET", b"different-process-key"
|
|
):
|
|
await self.assert_no_access("image.png", scope=descriptor["root_scope"])
|
|
self.api._get_all_output_dirs = lambda: []
|
|
await self.assert_no_access("image.png", scope=descriptor["root_scope"])
|
|
|
|
async def test_root_scope_keeps_containment_and_unscoped_access(self):
|
|
make_png(self.output_dir / "image.png")
|
|
listing = await self.client.get("/prompt_manager/images/output?mode=workspace")
|
|
descriptor = (await listing.json())["images"][0]
|
|
self.assertRegex(descriptor.get("root_scope", ""), r"^[0-9a-f]{64}$")
|
|
scope = descriptor["root_scope"]
|
|
self.assertEqual((await self.scoped_request(scope=scope)).status, 200)
|
|
outside = make_png(self.comfy_dir / "outside.png")
|
|
self._make_symlink(outside, self.output_dir / "escape.png")
|
|
for path in ("..%2Foutside.png", "escape.png"):
|
|
self.assertEqual((await self.scoped_request(path, scope=scope)).status, 403)
|
|
self.assertEqual(
|
|
(await self.client.get("/prompt_manager/images/serve/image.png")).status,
|
|
200,
|
|
)
|
|
|
|
async def test_root_scope_equal_label_reorder_replacement_and_symlink_are_stale(
|
|
self,
|
|
):
|
|
roots = [self.comfy_dir / name / "output" for name in ("a", "b", "c")]
|
|
for index, root in enumerate(roots):
|
|
make_png(root / "image.png", color=(index, 0, 0))
|
|
make_png(root / "thumbnails" / "image_thumb.png", color=(index, 0, 0))
|
|
# Both roots have the same display label under disposable public anchors.
|
|
with patch.object(
|
|
GalleryConfig,
|
|
"path_anchors",
|
|
classmethod(lambda cls: [str(root.parent) for root in roots]),
|
|
):
|
|
self.api._get_all_output_dirs = lambda: roots[:2]
|
|
listing = await self.client.get(
|
|
"/prompt_manager/images/output?mode=workspace"
|
|
)
|
|
descriptors = (await listing.json())["images"]
|
|
self.assertEqual({entry["root_dir"] for entry in descriptors}, {"output"})
|
|
token = next(
|
|
entry for entry in descriptors if entry["root_index"] == 0
|
|
).get("root_scope", "0" * 64)
|
|
for changed in ([roots[1], roots[0]], [roots[2], roots[1]]):
|
|
self.api._get_all_output_dirs = lambda: changed
|
|
for path in ("image.png", "thumbnails/image_thumb.png"):
|
|
await self.assert_no_access(path, scope=token)
|
|
link = self.comfy_dir / "linked-output"
|
|
self._make_symlink(roots[0], link)
|
|
self.api._get_all_output_dirs = lambda: [link]
|
|
listing = await self.client.get(
|
|
"/prompt_manager/images/output?mode=workspace"
|
|
)
|
|
token = (await listing.json())["images"][0]["root_scope"]
|
|
link.unlink()
|
|
self._make_symlink(roots[1], link)
|
|
for path in ("image.png", "thumbnails/image_thumb.png"):
|
|
await self.assert_no_access(path, scope=token)
|