"""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()