"""Tests for the core API module (py/api/__init__.py): helpers and base routes.""" 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.setdefault("server", _mock_server) sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) try: from tests.open_handles import assert_closed # noqa: E402 except ImportError: # discovered with tests/ as the top-level directory from open_handles import assert_closed # noqa: E402 from database.models import PromptModel # noqa: E402 from aiohttp import web # noqa: E402 from aiohttp.test_utils import AioHTTPTestCase # noqa: E402 import py.api as api_module # noqa: E402 from database.operations import PromptDatabase # noqa: E402 from py.api import PromptManagerAPI, _public_error, _public_path # noqa: E402 from py.config import GalleryConfig # noqa: E402 EXTRA_ROOTS_ENV = "PROMPT_MANAGER_EXTRA_GALLERY_ROOTS" class FolderPathsFixture(unittest.TestCase): """Temp ComfyUI tree exposed through a stub folder_paths module.""" def setUp(self): self.tmpdir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, self.tmpdir, True) self.addCleanup(assert_closed, self, self.tmpdir) # just before the rmtree self.addCleanup(PromptModel.close_all_instances) self.comfy_dir = Path(self.tmpdir) / "ComfyUI" self.output_dir = self.comfy_dir / "output" self.output_dir.mkdir(parents=True) self._orig_folder_paths = sys.modules.get("folder_paths") self.addCleanup(self._restore_folder_paths) self._orig_extra = os.environ.pop(EXTRA_ROOTS_ENV, None) self.addCleanup(self._restore_extra_env) self.install_folder_paths(base_path=str(self.comfy_dir)) def install_folder_paths(self, **attrs): attrs.setdefault("get_output_directory", lambda: str(self.output_dir)) sys.modules["folder_paths"] = types.SimpleNamespace(**attrs) def _restore_folder_paths(self): if self._orig_folder_paths is None: sys.modules.pop("folder_paths", None) else: sys.modules["folder_paths"] = self._orig_folder_paths def _restore_extra_env(self): if self._orig_extra is None: os.environ.pop(EXTRA_ROOTS_ENV, None) else: os.environ[EXTRA_ROOTS_ENV] = self._orig_extra class TestPublicPath(FolderPathsFixture): def test_path_under_base_is_relative_posix(self): target = self.output_dir / "renders" / "a.png" self.assertEqual(_public_path(str(target)), "output/renders/a.png") def test_base_itself_is_dot(self): self.assertEqual(_public_path(str(self.comfy_dir)), ".") def test_path_outside_base_is_basename_only(self): secret = Path(self.tmpdir) / "secret" / "passwords.txt" secret.parent.mkdir() secret.write_text("x") self.assertEqual(_public_path(str(secret)), "passwords.txt") self.assertNotIn(self.tmpdir, _public_path(str(secret))) def test_path_object_accepted(self): self.assertEqual(_public_path(self.output_dir), "output") def test_falls_back_to_output_parent_without_base_path(self): self.install_folder_paths() # no base_path attribute self.assertEqual(_public_path(str(self.output_dir / "x.png")), "output/x.png") def test_env_extra_root_keeps_its_own_name(self): extra = Path(self.tmpdir) / "nas" / "gallery" (extra / "sub").mkdir(parents=True) os.environ[EXTRA_ROOTS_ENV] = str(extra) self.assertEqual(_public_path(str(extra / "sub")), "gallery/sub") def test_without_folder_paths_returns_basename(self): sys.modules.pop("folder_paths", None) self.assertEqual(_public_path(str(self.output_dir / "img.png")), "img.png") def test_empty_or_none_gives_empty_string(self): self.assertEqual(_public_path(""), "") self.assertEqual(_public_path(None), "") class TestPublicError(FolderPathsFixture): def test_oserror_uses_strerror_and_basename(self): err = FileNotFoundError( 2, "No such file or directory", str(self.tmpdir) + "/a/b.png" ) text = _public_error(err) self.assertIn("No such file or directory", text) self.assertIn("b.png", text) self.assertNotIn(self.tmpdir, text) def test_oserror_without_filename(self): err = PermissionError(13, "Permission denied") self.assertEqual(_public_error(err), "Permission denied") def test_oserror_without_strerror_uses_type_name(self): err = OSError() self.assertEqual(_public_error(err), "OSError") def test_non_oserror_is_str(self): self.assertEqual(_public_error(ValueError("bad value")), "bad value") def test_non_oserror_with_empty_message_uses_type_name(self): self.assertEqual(_public_error(RuntimeError()), "RuntimeError") # ── gzip middleware ─────────────────────────────────────────────────── class TestGzipMiddleware(AioHTTPTestCase): async def get_application(self): big = json.dumps({"data": "x" * 4000}) async def big_json(request): return web.json_response(text=big) async def small_json(request): return web.json_response({"ok": True}) async def already_encoded(request): return web.Response( text=big, content_type="text/plain", headers={"Content-Encoding": "br"} ) async def binary(request): return web.Response(body=b"\x89PNG" * 1000, content_type="image/png") async def incompressible(request): return web.Response(body=os.urandom(4096), content_type="text/plain") async def streamed(request): resp = web.StreamResponse() await resp.prepare(request) await resp.write(b"x" * 4000) await resp.write_eof() return resp app = web.Application(middlewares=[api_module._gzip_middleware]) app.router.add_get("/prompt_manager/big", big_json) app.router.add_get("/prompt_manager/small", small_json) app.router.add_get("/prompt_manager/encoded", already_encoded) app.router.add_get("/prompt_manager/binary", binary) app.router.add_get("/prompt_manager/random", incompressible) app.router.add_get("/prompt_manager/stream", streamed) app.router.add_get("/other/big", big_json) return app async def _get(self, path, accept="gzip"): headers = {"Accept-Encoding": accept} if accept else {} return await self.client.request( "GET", path, headers=headers, auto_decompress=False, skip_auto_headers=["Accept-Encoding"], ) async def test_large_json_is_gzipped(self): resp = await self._get("/prompt_manager/big") self.assertEqual(resp.headers.get("Content-Encoding"), "gzip") self.assertEqual(resp.headers.get("Vary"), "Accept-Encoding") import gzip self.assertIn(b'"data"', gzip.decompress(await resp.read())) async def test_small_body_not_compressed(self): resp = await self._get("/prompt_manager/small") self.assertIsNone(resp.headers.get("Content-Encoding")) async def test_other_paths_untouched(self): resp = await self._get("/other/big") self.assertIsNone(resp.headers.get("Content-Encoding")) async def test_without_accept_encoding(self): resp = await self._get("/prompt_manager/big", accept=None) self.assertIsNone(resp.headers.get("Content-Encoding")) async def test_already_encoded_untouched(self): resp = await self._get("/prompt_manager/encoded") self.assertEqual(resp.headers.get("Content-Encoding"), "br") async def test_binary_type_untouched(self): resp = await self._get("/prompt_manager/binary") self.assertIsNone(resp.headers.get("Content-Encoding")) async def test_incompressible_body_untouched(self): resp = await self._get("/prompt_manager/random") self.assertIsNone(resp.headers.get("Content-Encoding")) async def test_stream_response_untouched(self): resp = await self._get("/prompt_manager/stream") self.assertIsNone(resp.headers.get("Content-Encoding")) self.assertEqual(len(await resp.read()), 4000) # ── JSON body reader ────────────────────────────────────────────────── class TestReadJsonBody(AioHTTPTestCase): """_read_json_body: chunked read with a byte cap, object-only payloads.""" async def get_application(self): async def echo(request): data, error = await api_module._read_json_body(request) if error is not None: return error return web.json_response({"success": True, "keys": sorted(data)}) async def small(request): data, error = await api_module._read_json_body(request, max_bytes=64) if error is not None: return error return web.json_response({"success": True, "data": data}) app = web.Application(client_max_size=4 * 1024 * 1024) app.router.add_post("/echo", echo) app.router.add_post("/small", small) return app async def test_valid_object_is_parsed(self): resp = await self.client.request("POST", "/echo", json={"b": 1, "a": 2}) self.assertEqual(resp.status, 200) self.assertEqual((await resp.json())["keys"], ["a", "b"]) async def test_default_cap_is_one_megabyte(self): self.assertEqual(api_module.JSON_BODY_MAX_BYTES, 1_000_000) body = json.dumps({"pad": "x" * 1_000_100}).encode() resp = await self.client.request( "POST", "/echo", data=body, headers={"Content-Type": "application/json"} ) self.assertEqual(resp.status, 413) data = await resp.json() self.assertFalse(data["success"]) self.assertIn("too large", data["error"].lower()) async def test_body_just_under_cap_is_accepted(self): payload = {"pad": "x" * (1_000_000 - 20)} body = json.dumps(payload).encode() self.assertLessEqual(len(body), 1_000_000) resp = await self.client.request( "POST", "/echo", data=body, headers={"Content-Type": "application/json"} ) self.assertEqual(resp.status, 200) async def test_chunked_upload_without_content_length_is_capped(self): async def gen(): for _ in range(10): yield b'{"pad": "' + b"x" * 20 + b'"}' resp = await self.client.request( "POST", "/small", data=gen(), headers={"Content-Type": "application/json"}, ) self.assertEqual(resp.status, 413) async def test_declared_oversize_content_length_is_413(self): body = b'{"pad": "' + b"x" * 200 + b'"}' resp = await self.client.request( "POST", "/small", data=body, headers={"Content-Type": "application/json"} ) self.assertEqual(resp.status, 413) async def test_invalid_json_is_400(self): for body in (b"{oops", b"\xff\xfe", b""): resp = await self.client.request( "POST", "/echo", data=body, headers={"Content-Type": "application/json"}, ) self.assertEqual(resp.status, 400, body) self.assertFalse((await resp.json())["success"]) async def test_non_object_json_is_400(self): for body in (b"[1, 2]", b'"text"', b"3", b"null", b"true"): resp = await self.client.request( "POST", "/echo", data=body, headers={"Content-Type": "application/json"}, ) self.assertEqual(resp.status, 400, body) self.assertIn("object", (await resp.json())["error"]) def test_helper_is_exposed_on_the_api_object(self): self.assertIs(PromptManagerAPI._read_json_body, api_module._read_json_body) class TestLongJobLock(unittest.TestCase): """_acquire_long_job: one job per family at a time, released on exit.""" def setUp(self): database = MagicMock() database.cleanup_duplicates.return_value = 0 self.api = PromptManagerAPI(db=database) def _run(self, coro): import asyncio return asyncio.run(coro) def test_second_acquire_is_refused_until_first_releases(self): async def scenario(): first = self.api._acquire_long_job("duplicates") self.assertIsNotNone(first) async with first: self.assertIsNone(self.api._acquire_long_job("duplicates")) # Other families are independent other = self.api._acquire_long_job("thumbnails") self.assertIsNotNone(other) async with other: pass again = self.api._acquire_long_job("duplicates") self.assertIsNotNone(again) async with again: pass self._run(scenario()) def test_lock_is_released_when_the_job_raises(self): async def scenario(): with self.assertRaises(RuntimeError): async with self.api._acquire_long_job("duplicates"): raise RuntimeError("boom") self.assertIsNotNone(self.api._acquire_long_job("duplicates")) self._run(scenario()) def test_busy_response_is_409_with_envelope(self): resp = self.api._long_job_busy_response("Duplicate scan") self.assertEqual(resp.status, 409) data = json.loads(resp.body) self.assertFalse(data["success"]) self.assertEqual(data["error"], "Duplicate scan already running") # ── UI and static routes ────────────────────────────────────────────── class CoreRoutesTestCase(AioHTTPTestCase): async def get_application(self): self.tmpdir = tempfile.mkdtemp() app = web.Application() routes = web.RouteTableDef() self.api = PromptManagerAPI( db=PromptDatabase(os.path.join(self.tmpdir, "prompts.db")) ) self.api.add_routes(routes) app.router.add_routes(routes) return app async def tearDownAsync(self): await super().tearDownAsync() # closes the aiohttp test client self.api.db.close_all() assert_closed(self, self.tmpdir) shutil.rmtree(self.tmpdir, ignore_errors=True) def _fake_root(self): """Point _get_project_root at the temp dir and return its web/ dir.""" web_dir = Path(self.tmpdir) / "web" (web_dir / "js").mkdir(parents=True) (web_dir / "lib").mkdir() orig_root = api_module._get_project_root api_module._get_project_root = lambda: self.tmpdir self.addCleanup(setattr, api_module, "_get_project_root", orig_root) return web_dir UI_PATHS = ( "/prompt_manager/workspace", "/prompt_manager/web", "/prompt_manager/gallery.html", "/prompt_manager/admin", "/prompt_manager/gallery", ) class TestUiRoutes(CoreRoutesTestCase): async def test_pages_are_served(self): for path in UI_PATHS: resp = await self.client.request("GET", path) self.assertEqual(resp.status, 200, path) self.assertIn("text/html", resp.headers["Content-Type"]) self.assertIn("<", await resp.text()) async def test_pages_are_cached_after_first_read(self): await self.client.request("GET", "/prompt_manager/admin") self.assertTrue(any(p.endswith("admin.html") for p in self.api._html_cache)) resp = await self.client.request("GET", "/prompt_manager/admin") self.assertEqual(resp.status, 200) async def test_missing_pages_are_404(self): self._fake_root() for path in UI_PATHS: resp = await self.client.request("GET", path) self.assertEqual(resp.status, 404, path) async def test_unreadable_pages_are_500(self): web_dir = self._fake_root() for name in ( "workspace.html", "index.html", "metadata.html", "admin.html", "gallery.html", ): (web_dir / name).mkdir() # a directory cannot be read as a file for path in UI_PATHS: resp = await self.client.request("GET", path) self.assertEqual(resp.status, 500, path) self.assertNotIn(self.tmpdir, await resp.text()) async def test_test_route(self): resp = await self.client.request("GET", "/prompt_manager/test") self.assertTrue((await resp.json())["success"]) class TestStaticRoutes(CoreRoutesTestCase): async def test_js_served_with_mime(self): resp = await self.client.request("GET", "/prompt_manager/js/admin.js") self.assertEqual(resp.status, 200) self.assertIn("application/javascript", resp.headers["Content-Type"]) async def test_workspace_module_served_with_javascript_mime(self): resp = await self.client.request( "GET", "/prompt_manager/js/workspace/index.mjs" ) self.assertEqual(resp.status, 200) self.assertIn("application/javascript", resp.headers["Content-Type"]) async def test_lib_served(self): web_dir = self._fake_root() (web_dir / "lib" / "x.css").write_text("body{}") resp = await self.client.request("GET", "/prompt_manager/lib/x.css") self.assertEqual(resp.status, 200) self.assertIn("text/css", resp.headers["Content-Type"]) self.assertEqual(await resp.text(), "body{}") async def test_real_lib_file_served_with_content_type(self): resp = await self.client.request( "GET", "/prompt_manager/lib/viewerjs/viewer.min.css" ) self.assertEqual(resp.status, 200) self.assertIn("text/css", resp.headers["Content-Type"]) async def test_allowed_asset_types_get_their_mime(self): web_dir = self._fake_root() expected = { "a.map": "application/json", "f.woff": "font/woff", "f.woff2": "font/woff2", "f.ttf": "font/ttf", "i.png": "image/png", "i.svg": "image/svg+xml", "i.ico": "image/x-icon", "p.html": "text/html", } for name in expected: (web_dir / "lib" / name).write_bytes(b"\x00\x01") for name, mime in expected.items(): resp = await self.client.request("GET", f"/prompt_manager/lib/{name}") self.assertEqual(resp.status, 200, name) self.assertIn(mime, resp.headers["Content-Type"], name) async def test_disallowed_extensions_are_403_even_when_present(self): web_dir = self._fake_root() for name in ("x.bin", "x.py", "x.json", "noext", "x.JS.bak"): (web_dir / "lib" / name).write_bytes(b"\x00") (web_dir / "js" / name).write_bytes(b"\x00") for prefix in ("lib", "js"): for name in ("x.bin", "x.py", "x.json", "noext", "x.JS.bak"): resp = await self.client.request( "GET", f"/prompt_manager/{prefix}/{name}" ) self.assertEqual(resp.status, 403, f"{prefix}/{name}") async def test_traversal_is_403(self): for prefix in ("lib", "js"): for encoded in ( "..%2F..%2Fpyproject.toml", "..%5C..%5Cpyproject.toml", "sub%2F..%2F..%2F..%2Fpyproject.toml", "%2E%2E%2Fpyproject.toml", ): resp = await self.client.request( "GET", f"/prompt_manager/{prefix}/{encoded}" ) self.assertEqual(resp.status, 403, f"{prefix}/{encoded}") async def test_absolute_drive_and_unc_paths_are_rejected_on_every_os(self): # Built literally so the Windows-only escapes are exercised everywhere: # os.path.join(root, "C:/x") drops root on Windows, "\\x" and UNC # shares re-anchor the path as well. candidates = ( "C:/Windows/win.ini", "C:%5CWindows%5Cwin.ini", "c:win.ini", "%5C%5Cserver%5Cshare%5Cx.js", "%2F%2Fserver%2Fshare%2Fx.js", "%5CWindows%5Cwin.ini", "%2Fetc%2Fpasswd", ) for prefix in ("lib", "js"): for encoded in candidates: resp = await self.client.request( "GET", f"/prompt_manager/{prefix}/{encoded}" ) self.assertIn(resp.status, (403, 404), f"{prefix}/{encoded}") self.assertNotIn("[extensions]", await resp.text()) async def test_symlink_escaping_web_dir_is_403(self): web_dir = self._fake_root() secret = Path(self.tmpdir) / "secret.js" secret.write_text("secret") link = web_dir / "lib" / "link.js" try: os.symlink(secret, link) except (OSError, NotImplementedError): self.skipTest("symlinks not available") resp = await self.client.request("GET", "/prompt_manager/lib/link.js") self.assertEqual(resp.status, 403) async def test_empty_path_is_not_served(self): for prefix in ("lib", "js"): resp = await self.client.request("GET", f"/prompt_manager/{prefix}/") self.assertIn(resp.status, (403, 404), prefix) async def test_missing_is_404(self): for prefix in ("lib", "js"): resp = await self.client.request("GET", f"/prompt_manager/{prefix}/nope.js") self.assertEqual(resp.status, 404, prefix) async def test_directory_is_404(self): web_dir = self._fake_root() (web_dir / "lib" / "dir.js").mkdir() resp = await self.client.request("GET", "/prompt_manager/lib/dir.js") self.assertEqual(resp.status, 404) async def test_root_of_static_dir_is_not_listed(self): resp = await self.client.request("GET", "/prompt_manager/lib/tailwind") self.assertEqual(resp.status, 403) # ── instance helpers ────────────────────────────────────────────────── class TestInstanceHelpers(FolderPathsFixture): def setUp(self): super().setUp() self.api = PromptManagerAPI( db=PromptDatabase(os.path.join(self.tmpdir, "prompts.db")) ) self.addCleanup(self.api.db.close_all) # runs before the rmtree cleanup orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES) self.addCleanup(setattr, GalleryConfig, "MONITORING_DIRECTORIES", orig_dirs) def _run(self, coro): import asyncio return asyncio.run(coro) def test_run_in_executor_with_kwargs(self): def add(a, b=0): return a + b self.assertEqual(self._run(self.api._run_in_executor(add, 1, b=2)), 3) self.assertEqual(self._run(self.api._run_in_executor(add, 1, 2)), 3) def test_invalidate_gallery_cache(self): self.api._gallery_cache = {"x": 1} self.api.invalidate_gallery_cache() self.assertEqual(self.api._gallery_cache, {}) def test_enrich_images_adds_urls_and_thumbnails(self): GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] img = self.output_dir / "sub" / "a.png" img.parent.mkdir() img.write_bytes(b"x") thumb = self.output_dir / "thumbnails" / "sub" / "a_thumb.png" thumb.parent.mkdir(parents=True) thumb.write_bytes(b"t") outside = Path(self.tmpdir) / "elsewhere.png" outside.write_bytes(b"o") images = [ {"id": 7, "image_path": str(img)}, {"id": 8, "image_path": str(outside)}, {"id": 9, "image_path": ""}, {"image_path": str(self.output_dir / "missing.png")}, ] result = self.api._enrich_images(images) self.assertEqual(result[0]["url"], "/prompt_manager/images/serve/sub/a.png") self.assertRegex( result[0]["thumbnail_url"], r"^/prompt_manager/images/serve/thumbnails/sub/a_thumb\.png\?v=\d+$", ) self.assertEqual(result[0]["relative_path"], os.path.join("sub", "a.png")) self.assertEqual(result[1]["url"], "/prompt_manager/images/8/file") self.assertNotIn("relative_path", result[1]) self.assertNotIn("url", result[2]) self.assertEqual(result[3]["url"], "/prompt_manager/images/serve/missing.png") self.assertNotIn("thumbnail_url", result[3]) def test_enrich_prompt_images_walks_prompts(self): GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] prompts = [ {"images": [{"id": 1, "image_path": str(self.output_dir / "p.png")}]}, {}, ] result = self.api._enrich_prompt_images(prompts) self.assertEqual( result[0]["images"][0]["url"], "/prompt_manager/images/serve/p.png" ) def test_clean_nan_recursive(self): nan = float("nan") cleaned = self.api._clean_nan_recursive( {"a": nan, "b": [nan, 1.5, {"c": nan}], "d": "x"} ) self.assertEqual(cleaned, {"a": None, "b": [None, 1.5, {"c": None}], "d": "x"}) def test_find_output_dir_uses_configured_and_caches(self): GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] self.api._cached_output_dir = None first = self.api._find_comfyui_output_dir() self.assertEqual(first, str(self.output_dir.resolve())) GalleryConfig.MONITORING_DIRECTORIES = [] self.assertEqual(self.api._find_comfyui_output_dir(), first) def test_find_output_dir_skips_configured_root_outside_comfyui(self): outside = Path(self.tmpdir) / "outside" outside.mkdir() GalleryConfig.MONITORING_DIRECTORIES = [str(outside)] self.api._cached_output_dir = None result = self.api._find_comfyui_output_dir() self.assertNotEqual( os.path.normcase(str(result or "")), os.path.normcase(str(outside)) ) def test_find_output_dir_skips_missing_configured(self): GalleryConfig.MONITORING_DIRECTORIES = [os.path.join(self.tmpdir, "nope")] self.api._cached_output_dir = None result = self.api._find_comfyui_output_dir() self.assertNotEqual(result, os.path.join(self.tmpdir, "nope")) # ── metadata parsing ────────────────────────────────────────────────── class TestMetadataParsing(unittest.TestCase): def setUp(self): database = MagicMock() database.cleanup_duplicates.return_value = 0 self.api = PromptManagerAPI(db=database) self.tmpdir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, self.tmpdir, True) def test_extract_metadata_from_png_and_garbage(self): from PIL import Image, PngImagePlugin path = os.path.join(self.tmpdir, "m.png") info = PngImagePlugin.PngInfo() info.add_text("prompt", "{}") Image.new("RGB", (1, 1)).save(path, pnginfo=info) self.assertEqual(self.api._extract_comfyui_metadata(path), {"prompt": "{}"}) garbage = os.path.join(self.tmpdir, "g.png") with open(garbage, "wb") as f: f.write(b"nope") self.assertEqual(self.api._extract_comfyui_metadata(garbage), {}) def test_parse_a1111_parameters(self): meta = {"parameters": "a cat\nNegative prompt: dog\nSteps: 20"} parsed = self.api._parse_comfyui_prompt(meta) self.assertEqual(parsed["positive_prompt"], "a cat") self.assertEqual(parsed["negative_prompt"], "dog") self.assertEqual(parsed["parameters"]["parameters"], meta["parameters"]) self.assertEqual(self.api._extract_readable_prompt(parsed), "a cat") def test_parse_comfyui_fields(self): meta = { "prompt": json.dumps({"1": {"class_type": "X", "inputs": {}}}), "workflow": "not json", "steps": "20", "cfg": "not-json-either", } parsed = self.api._parse_comfyui_prompt(meta) self.assertIsInstance(parsed["prompt"], dict) self.assertEqual(parsed["workflow"], "not json") self.assertEqual(parsed["parameters"]["steps"], 20) self.assertEqual(parsed["parameters"]["cfg"], "not-json-either") parsed = self.api._parse_comfyui_prompt( {"prompt": "plain text", "workflow": json.dumps({"nodes": []})} ) self.assertEqual(parsed["prompt"], "plain text") self.assertEqual(parsed["workflow"], {"nodes": []}) def test_extract_readable_prompt_variants(self): extract = self.api._extract_readable_prompt self.assertEqual(extract({"prompt": "plain"}), "plain") self.assertEqual(extract({"prompt": ["a", "", "b"]}), "a b") self.assertEqual(extract({"prompt": 42}), "42") self.assertEqual( extract({"prompt": None, "parameters": {"positive": ["x", "y"]}}), "x y" ) self.assertIsNone(extract({"prompt": {}, "parameters": {}})) graph = { "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}}, "2": {"class_type": "CLIPTextEncode", "inputs": {"text": "neg"}}, "3": { "class_type": "KSampler", "inputs": {"positive": ["1", 0], "negative": ["2", 0]}, }, } self.assertEqual(extract({"prompt": graph}), "pos") self.assertEqual(extract({"prompt": {"x": 1}, "workflow": graph}), "pos") def test_node_inputs_and_text_lookup(self): self.assertEqual(self.api._get_node_inputs("nope"), {}) self.assertEqual(self.api._get_node_inputs({"inputs": 5}), {}) self.assertEqual(self.api._get_node_inputs({"inputs": {"a": 1}}), {"a": 1}) listed = self.api._get_node_inputs( {"inputs": [{"name": "text", "link": None}, {"nope": 1}, "junk"]} ) self.assertEqual(list(listed), ["text"]) self.assertIsNone(self.api._find_text_in_node(None)) self.assertEqual(self.api._find_text_in_node({"inputs": {"text": "hi"}}), "hi") self.assertEqual( self.api._find_text_in_node( {"type": "CLIPTextEncode", "widgets_values": ["w"]} ), "w", ) self.assertIsNone( self.api._find_text_in_node( {"type": "CLIPTextEncode", "widgets_values": [" "]} ) ) self.assertIsNone(self.api._find_text_in_node({"type": "KSampler"})) def test_positive_prompt_from_nodes_array_and_fallbacks(self): extract = self.api._extract_positive_prompt_from_comfyui_data self.assertIsNone(extract("x")) self.assertIsNone(extract({})) self.assertIsNone(extract({"nodes": []})) workflow = { "nodes": [ {"id": 1, "type": "CLIPTextEncode", "widgets_values": ["pos"]}, { "id": 2, "type": "CLIPTextEncode", "widgets_values": ["neg"], "title": "Negative", }, {"id": 3, "type": "KSampler", "inputs": [{"name": "seed", "link": 4}]}, "junk", ] } self.assertEqual(extract(workflow), "pos") only_negative = { "nodes": [ { "id": 2, "type": "CLIPTextEncode", "widgets_values": ["neg"], "title": "neg", }, ] } self.assertEqual(extract(only_negative), "neg") bad_link = { "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "pos"}}, "3": { "class_type": "KSampler", "inputs": {"positive": "oops", "negative": []}, }, "4": { "class_type": "KSampler", "inputs": {"positive": ["zz", 0], "negative": []}, }, } self.assertEqual(extract(bad_link), "pos") if __name__ == "__main__": unittest.main()