Files
ComfyAssets-ComfyUI_PromptM…/tests/test_api_admin.py
T
Vito Sansevero 879247f1ef refactor: settings reads the image monitor through its module
get_settings located the ImageMonitor singleton by scanning sys.modules
for any module with 'image_monitor' in its name, a workaround for the
package being importable under two identities. It now uses the same
try-package / except-top-level import as the rest of the API
(_current_image_monitor) and falls back to the configured roots when no
monitor is running.
2026-10-08 18:01:34 -07:00

609 lines
22 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)
class TestSettingsReadMonitorThroughModule(AdminAPITestCase):
"""get_settings reaches the image monitor through utils.image_monitor."""
def test_no_sys_modules_scan(self):
import inspect
from py.api.admin import AdminRoutesMixin
source = inspect.getsource(AdminRoutesMixin.get_settings)
self.assertNotIn("sys.modules", source)
async def test_reports_directories_of_running_monitor(self):
import utils.image_monitor as im_mod
from utils.image_monitor import get_image_monitor
orig = im_mod._monitor_instance
im_mod._monitor_instance = None
self.addCleanup(setattr, im_mod, "_monitor_instance", orig)
monitor = get_image_monitor(MagicMock(), MagicMock())
monitor.monitored_directories = [str(self.output_dir)]
resp = await self.client.request("GET", "/prompt_manager/settings")
data = await resp.json()
self.assertEqual(resp.status, 200)
self.assertEqual(data["settings"]["monitored_directories"], ["output"])
async def test_falls_back_to_configured_roots_without_monitor(self):
import utils.image_monitor as im_mod
orig = im_mod._monitor_instance
im_mod._monitor_instance = None
self.addCleanup(setattr, im_mod, "_monitor_instance", orig)
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
resp = await self.client.request("GET", "/prompt_manager/settings")
data = await resp.json()
self.assertEqual(data["settings"]["monitored_directories"], ["output"])
self.assertEqual(data["settings"]["gallery_root_paths"], ["output"])
if __name__ == "__main__":
unittest.main()