Vito Sansevero
9ba6c0ca63
Image paths in API responses are now relative to the ComfyUI tree, so the auto-tag review page can no longer derive a serve URL from image_path; the server always provides url. The lora-dirs serving test assumed an empty LoraManager path, which only held in a checkout without a local config.json.
1557 lines
58 KiB
Python
1557 lines
58 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 json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
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.
|
|
sys.modules.setdefault("folder_paths", MagicMock())
|
|
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 # 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()
|
|
self.api.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):
|
|
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_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, 200)
|
|
|
|
|
|
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=".png"):
|
|
return self.output_dir / "thumbnails" / f"{name}_thumb{suffix}"
|
|
|
|
|
|
class TestThumbnailProgressStream(ThumbnailTestCase):
|
|
"""POST /prompt_manager/images/generate-thumbnails/progress (SSE)."""
|
|
|
|
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, once per file, plus the scan.
|
|
self.assertEqual(calls.count("_generate_one"), 3)
|
|
self.assertIn("_thumbnail_targets", calls)
|
|
|
|
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._find_comfyui_output_dir = lambda: None
|
|
|
|
resp = await self.client.request("POST", self.PROGRESS_URL)
|
|
|
|
events = parse_sse(await resp.text())
|
|
self.assertEqual(events[-1][0], "error")
|
|
|
|
async def test_unknown_quality_falls_back_to_medium(self):
|
|
make_png(self.output_dir / "one.png", size=(400, 400))
|
|
|
|
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), 300)
|
|
|
|
|
|
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._find_comfyui_output_dir = lambda: None
|
|
|
|
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 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 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_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, [])
|
|
|
|
def test_thumbnail_size_ignores_non_string_quality(self):
|
|
self.assertEqual(self.api._thumbnail_size({"w": 99999}), (300, 300))
|
|
self.assertEqual(self.api._thumbnail_size(None), (300, 300))
|
|
self.assertEqual(self.api._thumbnail_size("low"), (150, 150))
|
|
|
|
|
|
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_searches_all_roots(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, 200)
|
|
|
|
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)
|
|
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 = (64, 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", ".jpg"), self.SIZE, True)
|
|
with patch.object(self.api, "_generate_video_thumbnail", return_value=False):
|
|
bad = self.api._generate_one(
|
|
src, self._dst("clip", ".jpg"), 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", ".jpg").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._thumbnail_targets = _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._find_comfyui_output_dir = _raise
|
|
status, _ = await self._json(
|
|
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
|
)
|
|
self.assertEqual(status, 500)
|
|
|
|
def _dst(self, name, suffix=".png"):
|
|
return self.output_dir / "thumbnails" / f"{name}_thumb{suffix}"
|
|
|
|
|
|
class TestVideoThumbnailFallbacks(ImageRouteCoverageCase):
|
|
"""_generate_video_thumbnail degrades: cv2 -> ffmpeg -> placeholder."""
|
|
|
|
SIZE = (48, 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(self):
|
|
with patch.dict(sys.modules, {"cv2": self._fake_cv2(opened=False)}):
|
|
closed = self.api._generate_video_thumbnail(
|
|
self.video, self.thumb, self.SIZE
|
|
)
|
|
with patch.dict(sys.modules, {"cv2": self._fake_cv2(read_ok=False)}):
|
|
unread = self.api._generate_video_thumbnail(
|
|
self.video, self.thumb, self.SIZE
|
|
)
|
|
self.assertEqual((closed, unread), (False, False))
|
|
|
|
def test_cv2_crash_is_contained(self):
|
|
cv2 = MagicMock()
|
|
cv2.VideoCapture.side_effect = RuntimeError("driver crash")
|
|
with patch.dict(sys.modules, {"cv2": cv2}):
|
|
ok = self.api._generate_video_thumbnail(self.video, self.thumb, self.SIZE)
|
|
self.assertFalse(ok)
|
|
|
|
def test_ffmpeg_absent_yields_placeholder(self):
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": None}),
|
|
patch("subprocess.run", 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)
|
|
|
|
def test_ffmpeg_failure_yields_placeholder(self):
|
|
failed = MagicMock(returncode=1, stderr="no stream")
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": None}),
|
|
patch("subprocess.run", return_value=failed),
|
|
):
|
|
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.run", return_value=MagicMock(returncode=0)) 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")
|
|
|
|
def test_placeholder_failure_returns_false(self):
|
|
with (
|
|
patch.dict(sys.modules, {"cv2": None}),
|
|
patch("subprocess.run", 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._find_comfyui_output_dir = lambda: None
|
|
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._find_comfyui_output_dir = _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._find_comfyui_output_dir = lambda: str(out)
|
|
|
|
status, data = await self._json("POST", self.URL)
|
|
|
|
self.assertEqual(status, 400)
|
|
self.assertFalse(data["success"])
|
|
self.assertTrue(sibling.is_file())
|
|
|
|
|
|
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 TestOutputScanHelpers(ImageAPITestCase):
|
|
"""Stable ids for filesystem entries, no symlink reads, stop on disconnect."""
|
|
|
|
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"])
|
|
|
|
async def test_thumbnail_stream_stops_when_the_client_disconnects(self):
|
|
generated = []
|
|
|
|
def fake_generate(src, dst, size, is_video):
|
|
generated.append(src)
|
|
return {
|
|
"action": "generated",
|
|
"file": src.name,
|
|
"dir": src.parent.name,
|
|
"type": "image",
|
|
}
|
|
|
|
self.api._generate_one = fake_generate
|
|
|
|
class GoneResponse:
|
|
async def write(self, data):
|
|
raise ConnectionResetError("client went away")
|
|
|
|
targets = [(Path(f"/x/{i}.png"), Path(f"/x/t{i}.jpg"), False) for i in range(3)]
|
|
await self.api._stream_thumbnails(GoneResponse(), targets, (256, 256))
|
|
self.assertEqual(generated, [])
|