nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
b0a025a326b523ab88ebb3b41bc71e204e4a91b9
ComfyAssets-ComfyUI_PromptM…/tests/test_api_core.py
T
Vito b2cc489fd4 feat(inspect): copy settings summary and local workflow roles (#286) 
* feat(inspect): copy settings summary in the full inspector

The legacy metadata page's COPY ALL (LI-4) becomes "Copy settings summary"
after the six facts of the shared full inspector, so the Images viewer,
the recorded history's current file and Inspect all offer it; the Images
side panel keeps the concept's footer. The text is the legacy copyAllData
template with its placeholders, filled from what the inspector shows. A
unit test runs the template read from web/metadata.html on the legacy
parse of every parity PNG: 22 of 25 are byte-equal, and the other three
differ only by the ZeroOut guard (#249) or facts the legacy parse missed.

Copying follows the concept's prompt copy: the toast "Settings summary
copied.", or the select-to-copy dialog (now titled per copy) without a
clipboard. The concept comparisons map the button as D-NEW-1.

* feat(inspect): local execution evidence and prompt tracing

A browser port of the server's inspection resolver for one sampler
(PR L3 Task 8, LI-1). local-resolver.mjs answers inspect_prompt_evidence's
{prompts, resolution} for a local read's records: execution evidence
(read_pair_execution and its per-node match), every traced strategy
(encoders, strings, concatenations, PromptManager and its historical
classes, passthrough, zeroed, selector branches, composites, guiders),
the latent stage walk, the R1 overwritten-text guard and the inspection
profile's visit, depth, text, path, built and scanned bounds, with the
same reason codes and outcome lines. py-json.mjs reads the records as
CPython's json.loads does (dict order and last-wins keys, int versus
float, BigInt past 2^53, NaN) and measures and strips text as str does.
Several sampler candidates, detailer stages and A1111 parameters raise
LocalResolverGap until Task 9. Nothing is wired into Inspect yet.

Python is the oracle: test_local_inspect_parity compares the port with
the server on the parity corpus, 107 named record sets, every work bound
at its edge and 600 seeded fuzzed graphs, and requires gaps to be exactly
the Task 9 paths. Sixteen porting mutations each fail the suite. The
Node unit test replays committed server answers, which the Python suite
keeps current.

* fix(inspect): a list or object workflow node type answers malformed

The R1 overwritten-text check looked every UI workflow node's `type` up
among the subgraph IDs. A `type` that is a JSON list or object raised
TypeError there, so the inspect route answered 500 "Prompt inspection
failed". Such a workflow is now malformed evidence: the resolution is
unresolved with reason_code `malformed` and the inspector's own
"Prompt evidence is malformed" line, never exception text.

The browser port answers the same, and the parity corpus gains the
list, object, nested-subgraph and skipped-ID shapes.

* fix(inspect): a list or object graph class_type is never a preview output

With several sampler candidates the inspector first drops stages whose
outputs reach only previews. That pass looked each API graph node's
`class_type` up in the preview class set, so a list or object
`class_type` raised TypeError and the inspect route answered 500. Only
a string names a class now; such a node is an ordinary terminal, as the
rest of the resolver already treats it.

* feat(inspect): local candidate selection matching the inspector

The browser port of the inspection resolver now answers every record set
the server does, so no path is left to the server any more:

- several sampler candidates, with the identical-branch rule (#270): a
  compare grid whose branches all resolve to one pair (text and recipe)
  is resolved, any other mix is ambiguous;
- stages whose outputs reach only previews are dropped first;
- detailer stages are judged with the samplers in their image ancestry
  (multistage-differs, cycles, depth and the per-stage visit budgets),
  and a sampler that only feeds detailers is not a candidate of its own;
- A1111 `parameters`, with Python's multiline anchors, `\s` and Unicode
  `\d`, the bare footer and Template trailers, counted in code points.

The parity corpus gains the selection cases and a seeded fuzzer, and
every inspection vector now gives the server's prompts, resolution and
facts locally. 16 of 17 mutations are caught; the uncaught one (not
flattening an ambiguous upstream stage) cannot change any answer.

* feat(inspect): local prompts resolved like the server's

A file chosen on the Inspect page now shows the prompts and resolution
the server's inspector gives the same records: buildLocalPayload takes
its roles from the local resolver instead of the interim open roles, so
the legacy fill applies exactly where it does for Images.

The resolver reads the server's own node table. The page fetches
web/js/inspection-semantics.json once through the API client as it
mounts, before any selection; the static route serves that one JSON
file by exact path (every other JSON stays forbidden). A table that
cannot be fetched answers "The inspector's node table is unavailable."
as the server does, and a later selection asks again.

The Copy settings summary deviation is numbered D117.

* fix(inspect): skip a workflow node whose type is a list or object

The overwritten-PromptManager walk now skips such a node, on the server
and in the browser port, instead of marking the whole file malformed.
A prompt that resolves still resolves; the other nodes still feed the
check.

* fix(inspect): one JSON depth cap for both browser parsers

py-json.mjs and local-video.mjs now refuse past the same 1,000 nested
arrays or objects, counting every container as CPython's scanner does.
Boundary vectors at 999, 1,000 and 1,001 for a deep prompt and a deep
workflow run through the Python oracle, and TESTING.md gives the
measured per-version server limits and both divergence directions.

* fix(inspect): realpath containment for the exact-path node table route

inspection-semantics.json now goes through the same realpath check as
every other static file, so a symlink out of web/ is refused. The
forbidden-path test asserts exactly 403.

* fix(inspect): abortable, stepped local reads with a busy state

The wait for the node table now ends with the selection's signal. A
local read resolves and projects in bounded steps with a turn of the
event loop between them, and the inspector carries aria-busy while it
loads. A budget test covers maximal 2 MiB graphs and checks the stepped
read is byte-identical to one synchronous build. The D117 wrap rule now
targets only the summary button, and TESTING.md scopes the restart line
to PRs L1 and L2.

* fix(inspect): record the two CPython-JSON emulations as a follow-up
2026-10-23 01:38:36 -07:00

900 lines
35 KiB
Python
Raw Blame History

"""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"])
async def test_image_annotation_routes_are_registered(self):
table = {
(route.method, route.resource.canonical)
for route in self.app.router.routes()
if route.resource is not None
}
for entry in (
("POST", "/prompt_manager/image-annotations/observe"),
("GET", "/prompt_manager/image-annotations/capture/{token}"),
("DELETE", "/prompt_manager/image-annotations/capture/{token}"),
("PUT", "/prompt_manager/image-annotations/tags"),
):
self.assertIn(entry, table)
async def test_image_tag_job_route_is_registered_beside_the_scan_job(self):
table = {
(route.method, route.resource.canonical)
for route in self.app.router.routes()
if route.resource is not None
}
self.assertIn(("POST", "/prompt_manager/jobs/scan"), table)
self.assertIn(("POST", "/prompt_manager/jobs/image-tags"), table)
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_inspection_semantics_table_is_the_one_json_served(self):
"""Local Inspect reads the node table the server resolves with (#249 PR L3).
Only that exact file is served; every other JSON stays forbidden.
"""
resp = await self.client.request(
"GET", "/prompt_manager/js/inspection-semantics.json"
)
self.assertEqual(resp.status, 200)
self.assertIn("application/json", resp.headers["Content-Type"])
table = Path(__file__).resolve().parents[1] / "web" / "js"
self.assertEqual(
await resp.read(), (table / "inspection-semantics.json").read_bytes()
)
for path in (
"js/shared/toolkit/provenance.json",
"lib/viewerjs/provenance.json",
"lib/inspection-semantics.json",
"js/workspace/..%2Finspection-semantics.json",
"js/INSPECTION-SEMANTICS.json",
):
resp = await self.client.request("GET", f"/prompt_manager/{path}")
self.assertEqual(resp.status, 403, path)
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_inspection_semantics_symlink_escaping_web_dir_is_403(self):
"""The exact-path JSON gets the same realpath containment as any file."""
web_dir = self._fake_root()
secret = Path(self.tmpdir) / "secret.json"
secret.write_text('{"secret": true}')
try:
os.symlink(secret, web_dir / "js" / "inspection-semantics.json")
except (OSError, NotImplementedError):
self.skipTest("symlinks not available")
resp = await self.client.request(
"GET", "/prompt_manager/js/inspection-semantics.json"
)
self.assertEqual(resp.status, 403)
self.assertNotIn("secret", await resp.text())
async def test_inspection_semantics_missing_is_404(self):
self._fake_root()
resp = await self.client.request(
"GET", "/prompt_manager/js/inspection-semantics.json"
)
self.assertEqual(resp.status, 404)
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()
Reference in New Issue View Git Blame Copy Permalink