Files
ComfyAssets-ComfyUI_PromptM…/tests/test_api_admin.py
T
Vito Sansevero 38f4c0af43 perf: run duplicate and output scans off the event loop
The duplicate scan walked the output tree and hashed every file on the
aiohttp event loop, and the duplicate delete removed files inline; the
output scan ran its per-file database work one executor hop at a time.

- _find_duplicate_images_sync / _describe_media_file do the walk and
  hashing; find_duplicate_images() runs them via _run_in_executor.
- _delete_duplicate_images_sync / _delete_one_duplicate do the removal
  (absolute or output-relative paths) off the loop.
- The output scan is split into _collect_output_media_sync,
  _extract_batch_metadata_sync and _ingest_scan_batch_sync (one
  executor call per batch for all DB work) with a small SSE generator.
- Media collection is shared and de-duplicates by normcase'd path.
2026-10-08 18:01:34 -07:00

567 lines
21 KiB
Python

"""Route tests for the admin API (py/api/admin.py).
Uses aiohttp's test client against a PromptManagerAPI wired to a temporary
database and a stub ``folder_paths`` module, so the tests run without ComfyUI.
"""
import json
import os
import shutil
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import MagicMock
_mock_server = MagicMock()
_mock_server.PromptServer.instance.routes = MagicMock()
sys.modules["server"] = _mock_server
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from aiohttp import web # noqa: E402
from aiohttp.test_utils import AioHTTPTestCase # noqa: E402
from database.operations import PromptDatabase # noqa: E402
from py.api import PromptManagerAPI # noqa: E402
from py.config import GalleryConfig, PromptManagerConfig # noqa: E402
CONFIG_PATH_ENV = "PROMPT_MANAGER_CONFIG_PATH"
EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS"
class AdminAPITestCase(AioHTTPTestCase):
"""Test app with PromptManager routes, temp DB, temp config and fake ComfyUI tree."""
async def get_application(self):
self.tmpdir = tempfile.mkdtemp()
self.comfy_dir = Path(self.tmpdir) / "ComfyUI"
self.output_dir = self.comfy_dir / "output"
self.output_dir.mkdir(parents=True)
self.outside_dir = Path(self.tmpdir) / "outside"
self.outside_dir.mkdir()
self.config_path = os.path.join(self.tmpdir, "config.json")
self._orig_env = {
key: os.environ.pop(key, None) for key in (CONFIG_PATH_ENV, EXTRA_ROOTS_ENV)
}
os.environ[CONFIG_PATH_ENV] = self.config_path
self._orig_folder_paths = sys.modules.get("folder_paths")
sys.modules["folder_paths"] = types.SimpleNamespace(
base_path=str(self.comfy_dir),
get_output_directory=lambda: str(self.output_dir),
)
self._orig_gallery_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
self._orig_pm_config = PromptManagerConfig.get_config()
GalleryConfig.MONITORING_DIRECTORIES = []
db_path = os.path.join(self.tmpdir, "prompts.db")
app = web.Application()
routes = web.RouteTableDef()
self.api = PromptManagerAPI()
self.api.db = PromptDatabase(db_path)
self.api.add_routes(routes)
app.router.add_routes(routes)
return app
async def tearDownAsync(self):
GalleryConfig.MONITORING_DIRECTORIES = self._orig_gallery_dirs
PromptManagerConfig.update_config(self._orig_pm_config)
for key, value in self._orig_env.items():
if value is None:
os.environ.pop(key, None)
else:
os.environ[key] = value
if self._orig_folder_paths is None:
sys.modules.pop("folder_paths", None)
else:
sys.modules["folder_paths"] = self._orig_folder_paths
shutil.rmtree(self.tmpdir, ignore_errors=True)
# ── helpers ────────────────────────────────────────────────────────
def _write_config(self, content):
with open(self.config_path, "w") as f:
json.dump(content, f)
def _read_config(self):
with open(self.config_path) as f:
return json.load(f)
class TestSaveSettingsGalleryRoots(AdminAPITestCase):
async def test_rejects_filesystem_root_and_keeps_config(self):
self._write_config({"marker": "untouched"})
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": [os.path.abspath(os.sep)]},
)
self.assertEqual(resp.status, 400)
data = await resp.json()
self.assertFalse(data["success"])
self.assertEqual(self._read_config(), {"marker": "untouched"})
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [])
async def test_rejects_directory_outside_comfyui(self):
self._write_config({"marker": "untouched"})
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": [str(self.outside_dir)]},
)
self.assertEqual(resp.status, 400)
self.assertEqual(self._read_config(), {"marker": "untouched"})
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [])
async def test_one_bad_root_rejects_whole_request(self):
good = self.output_dir / "good"
good.mkdir()
self._write_config({"marker": "untouched"})
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": [str(good), str(self.outside_dir)]},
)
self.assertEqual(resp.status, 400)
self.assertEqual(self._read_config(), {"marker": "untouched"})
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [])
async def test_legacy_single_path_rejected_too(self):
self._write_config({"marker": "untouched"})
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_path": os.path.expanduser("~")},
)
self.assertEqual(resp.status, 400)
self.assertEqual(self._read_config(), {"marker": "untouched"})
async def test_accepts_subdir_of_output_and_persists(self):
good = self.output_dir / "renders"
good.mkdir()
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": [str(good)]},
)
self.assertEqual(resp.status, 200)
data = await resp.json()
self.assertTrue(data["success"])
self.assertTrue(data["restart_required"])
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(good)])
saved = self._read_config()
self.assertEqual(saved["gallery"]["monitoring"]["directories"], [str(good)])
async def test_accepts_env_listed_directory(self):
os.environ[EXTRA_ROOTS_ENV] = str(self.outside_dir)
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": [str(self.outside_dir)]},
)
self.assertEqual(resp.status, 200)
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(self.outside_dir)])
async def test_legacy_single_path_accepted(self):
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_path": str(self.output_dir)},
)
self.assertEqual(resp.status, 200)
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [str(self.output_dir)])
async def test_empty_paths_clear_roots(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": ["", " "]},
)
self.assertEqual(resp.status, 200)
self.assertEqual(GalleryConfig.MONITORING_DIRECTORIES, [])
self.assertEqual(
self._read_config()["gallery"]["monitoring"]["directories"], []
)
async def test_paths_must_be_a_list(self):
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": str(self.output_dir)},
)
self.assertEqual(resp.status, 400)
async def test_path_entries_must_be_strings(self):
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": [42]},
)
self.assertEqual(resp.status, 400)
class FakeMonitor:
def __init__(self, dirs):
self.monitored_directories = list(dirs)
def get_status(self):
return {
"running": True,
"monitored_directories": self.monitored_directories,
"handler_active": True,
"observer_alive": True,
}
class TestNoAbsolutePathsInResponses(AdminAPITestCase):
"""No admin response may reveal the server's absolute filesystem layout."""
def _install_monitor(self, dirs):
import utils.image_monitor as im_mod
orig = im_mod._monitor_instance
im_mod._monitor_instance = FakeMonitor(dirs)
self.addCleanup(setattr, im_mod, "_monitor_instance", orig)
async def _body(self, method, path, **kwargs):
resp = await self.client.request(method, path, **kwargs)
return resp, await resp.text()
async def test_settings_show_roots_relative_to_comfyui(self):
renders = self.output_dir / "renders"
renders.mkdir()
GalleryConfig.MONITORING_DIRECTORIES = [str(renders)]
self._install_monitor([str(renders)])
resp, body = await self._body("GET", "/prompt_manager/settings")
self.assertEqual(resp.status, 200)
self.assertNotIn(self.tmpdir, body)
data = json.loads(body)
self.assertEqual(data["settings"]["gallery_root_paths"], ["output/renders"])
self.assertEqual(data["settings"]["gallery_root_path"], "output/renders")
self.assertEqual(data["settings"]["monitored_directories"], ["output/renders"])
async def test_settings_relative_root_round_trips(self):
"""What GET returns can be POSTed back unchanged from any CWD."""
renders = self.output_dir / "renders"
renders.mkdir()
resp = await self.client.request(
"POST",
"/prompt_manager/settings",
json={"gallery_root_paths": ["output/renders"]},
)
self.assertEqual(resp.status, 200)
self.assertEqual(
[os.path.normcase(p) for p in GalleryConfig.MONITORING_DIRECTORIES],
[os.path.normcase(os.path.realpath(renders))],
)
async def test_diagnostics_paths_are_relative(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
self._install_monitor([str(self.output_dir)])
resp, body = await self._body("GET", "/prompt_manager/diagnostics")
self.assertEqual(resp.status, 200)
self.assertNotIn(self.tmpdir, body)
data = json.loads(body)
self.assertEqual(
data["diagnostics"]["comfyui_output"]["output_dirs"], ["output"]
)
self.assertEqual(
data["diagnostics"]["image_monitor"]["monitored_directories"], ["output"]
)
async def test_diagnostics_missing_db_message_has_no_path(self):
self.api.db.model.db_path = os.path.join(self.tmpdir, "missing", "prompts.db")
resp, body = await self._body("GET", "/prompt_manager/diagnostics")
self.assertNotIn(self.tmpdir, body)
data = json.loads(body)
self.assertEqual(data["diagnostics"]["database"]["status"], "error")
self.assertIn("prompts.db", data["diagnostics"]["database"]["message"])
async def test_scan_duplicates_paths_are_relative_to_output(self):
(self.output_dir / "a.png").write_bytes(b"same-bytes")
(self.output_dir / "sub").mkdir()
(self.output_dir / "sub" / "b.png").write_bytes(b"same-bytes")
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
resp, body = await self._body("GET", "/prompt_manager/scan_duplicates")
self.assertEqual(resp.status, 200)
self.assertNotIn(self.tmpdir, body)
data = json.loads(body)
self.assertEqual(len(data["duplicates"]), 1)
paths = sorted(img["path"] for img in data["duplicates"][0]["images"])
self.assertEqual(paths, ["a.png", os.path.join("sub", "b.png")])
async def test_delete_duplicates_accepts_relative_paths(self):
target = self.output_dir / "sub" / "b.png"
target.parent.mkdir()
target.write_bytes(b"bytes")
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
resp, body = await self._body(
"POST",
"/prompt_manager/delete_duplicate_images",
json={"image_paths": [os.path.join("sub", "b.png")]},
)
self.assertEqual(resp.status, 200)
data = json.loads(body)
self.assertEqual(data["deleted_count"], 1)
self.assertFalse(target.exists())
async def test_scan_images_stream_has_no_absolute_paths(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
resp, body = await self._body("POST", "/prompt_manager/scan")
self.assertEqual(resp.status, 200)
self.assertNotIn(self.tmpdir, body)
self.assertIn('"directories": ["output"]', body)
def _write_png(path, prompt_text=None):
"""Write a tiny PNG, optionally carrying a ComfyUI-style 'prompt' chunk."""
from PIL import Image, PngImagePlugin
img = Image.new("RGB", (2, 2), (255, 0, 0))
info = PngImagePlugin.PngInfo()
if prompt_text is not None:
graph = {
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt_text}},
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "bad"}},
"3": {
"class_type": "KSampler",
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
},
}
info.add_text("prompt", json.dumps(graph))
img.save(str(path), pnginfo=info)
class TestScansRunOffTheEventLoop(AdminAPITestCase):
"""Blocking filesystem work must go through _run_in_executor."""
def _spy_executor(self):
calls = []
original = self.api._run_in_executor
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
async def test_scan_duplicates_runs_in_executor(self):
(self.output_dir / "a.png").write_bytes(b"dup")
(self.output_dir / "b.png").write_bytes(b"dup")
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
calls = self._spy_executor()
resp = await self.client.request("GET", "/prompt_manager/scan_duplicates")
self.assertEqual(resp.status, 200)
self.assertEqual((await resp.json())["duplicates"][0]["count"], 2)
self.assertIn("_find_duplicate_images_sync", calls)
async def test_delete_duplicates_runs_in_executor(self):
(self.output_dir / "a.png").write_bytes(b"dup")
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
calls = self._spy_executor()
resp = await self.client.request(
"POST",
"/prompt_manager/delete_duplicate_images",
json={"image_paths": ["a.png"]},
)
self.assertEqual(resp.status, 200)
self.assertEqual((await resp.json())["deleted_count"], 1)
self.assertIn("_delete_duplicate_images_sync", calls)
async def test_output_scan_runs_in_executor(self):
_write_png(self.output_dir / "gen.png", prompt_text="a red square")
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
calls = self._spy_executor()
resp = await self.client.request("POST", "/prompt_manager/scan")
body = await resp.text()
self.assertEqual(resp.status, 200)
self.assertIn('"type": "complete"', body)
self.assertIn('"added": 1', body)
for name in (
"_collect_output_media_sync",
"_extract_batch_metadata_sync",
"_ingest_scan_batch_sync",
):
self.assertIn(name, calls)
async def test_output_scan_links_existing_prompt_on_second_pass(self):
_write_png(self.output_dir / "gen.png", prompt_text="a red square")
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
first = await (await self.client.request("POST", "/prompt_manager/scan")).text()
second = await (
await self.client.request("POST", "/prompt_manager/scan")
).text()
self.assertIn('"added": 1', first)
self.assertIn('"linked": 1', second)
self.assertIn('"added": 0', second)
prompts = self.api.db.get_recent_prompts(limit=10)["prompts"]
self.assertEqual([p["text"] for p in prompts], ["a red square"])
async def test_output_scan_without_directories_reports_error(self):
GalleryConfig.MONITORING_DIRECTORIES = [str(self.outside_dir)]
sys.modules.pop("folder_paths", None)
self.api._find_comfyui_output_dir = lambda: None
body = await (await self.client.request("POST", "/prompt_manager/scan")).text()
self.assertIn('"type": "error"', body)
class TestDuplicateScanSyncBody(AdminAPITestCase):
"""Behaviour of the synchronous duplicate scan against a real tree."""
def _scan(self):
return self.api._find_duplicate_images_sync(str(self.output_dir))
def test_groups_identical_files_and_sorts_by_mtime(self):
older = self.output_dir / "older.png"
newer = self.output_dir / "sub" / "newer.PNG"
newer.parent.mkdir()
older.write_bytes(b"same")
newer.write_bytes(b"same")
os.utime(older, (1_000_000, 1_000_000))
os.utime(newer, (2_000_000, 2_000_000))
(self.output_dir / "unique.png").write_bytes(b"different")
groups = self._scan()
self.assertEqual(len(groups), 1)
self.assertEqual(groups[0]["count"], 2)
self.assertEqual(
[img["filename"] for img in groups[0]["images"]], ["older.png", "newer.PNG"]
)
self.assertEqual(
groups[0]["images"][1]["path"], os.path.join("sub", "newer.PNG")
)
self.assertEqual(
groups[0]["images"][1]["url"], "/prompt_manager/images/serve/sub/newer.PNG"
)
self.assertNotIn(self.tmpdir, json.dumps(groups))
def test_thumbnails_dir_is_skipped_and_thumbnail_url_reported(self):
(self.output_dir / "a.png").write_bytes(b"same")
(self.output_dir / "b.png").write_bytes(b"same")
thumbs = self.output_dir / "thumbnails"
thumbs.mkdir()
(thumbs / "a_thumb.png").write_bytes(b"same")
(thumbs / "b_thumb.png").write_bytes(b"other")
groups = self._scan()
self.assertEqual(len(groups), 1)
by_name = {img["filename"]: img for img in groups[0]["images"]}
self.assertEqual(set(by_name), {"a.png", "b.png"})
self.assertEqual(
by_name["a.png"]["thumbnail_url"],
"/prompt_manager/images/serve/thumbnails/a_thumb.png",
)
self.assertEqual(
by_name["b.png"]["thumbnail_url"],
"/prompt_manager/images/serve/thumbnails/b_thumb.png",
)
def test_videos_are_flagged_and_unreadable_entries_skipped(self):
(self.output_dir / "clip.mp4").write_bytes(b"vid")
(self.output_dir / "copy.mp4").write_bytes(b"vid")
(self.output_dir / "not_a_file.png").mkdir() # rglob matches, hashing fails
groups = self._scan()
self.assertEqual(len(groups), 1)
self.assertTrue(all(img["is_video"] for img in groups[0]["images"]))
self.assertTrue(
all(img["media_type"] == "video" for img in groups[0]["images"])
)
def test_missing_output_dir_gives_empty_list(self):
self.assertEqual(
self.api._find_duplicate_images_sync(os.path.join(self.tmpdir, "nope")), []
)
class TestDeleteDuplicatesSyncBody(AdminAPITestCase):
def test_deletes_file_and_thumbnail_inside_output(self):
target = self.output_dir / "sub" / "x.png"
target.parent.mkdir()
target.write_bytes(b"x")
thumb = self.output_dir / "thumbnails" / "sub" / "x_thumb.png"
thumb.parent.mkdir(parents=True)
thumb.write_bytes(b"t")
result = self.api._delete_duplicate_images_sync(
[str(target)], str(self.output_dir)
)
self.assertEqual(result["deleted_count"], 1)
self.assertEqual(result["failed_count"], 0)
self.assertFalse(target.exists())
self.assertFalse(thumb.exists())
def test_refuses_paths_outside_output_and_reports_missing(self):
secret = self.outside_dir / "secret.txt"
secret.write_text("s")
result = self.api._delete_duplicate_images_sync(
[str(secret), "ghost.png"], str(self.output_dir)
)
self.assertEqual(result["deleted_count"], 0)
self.assertEqual(result["failed_count"], 2)
self.assertTrue(secret.exists())
self.assertTrue(any("outside" in f for f in result["failed_files"]))
self.assertTrue(any("not found" in f for f in result["failed_files"]))
def test_without_output_dir_everything_fails(self):
result = self.api._delete_duplicate_images_sync(["a.png"], None)
self.assertEqual(result["failed_count"], 1)
self.assertEqual(result["deleted_count"], 0)
if __name__ == "__main__":
unittest.main()