Vito Sansevero
b1bfce702d
The first push of the branch ran the Tests matrix on all three operating systems. Every Python job failed; the causes were: - numpy missing in CI: the WD14 engine tests import it (requirements-dev) - macOS: temp directories resolve through /private; the admin API tests now use the real path like the gallery roots they compare against - Windows: the startup banner printed an emoji, which raises UnicodeEncodeError on a legacy-code-page console and aborted loading the whole node; the banner is ASCII now and a test loads the package under a cp1252 stdout - Windows: gallery subfolders came back with backslashes, which neither the folder tree nor the search filter understand; they are "/" joined on every platform. A relative DEFAULT_DB_PATH kept its forward slashes after being joined to the extension root; it is normalised now - Windows: sqlite connections opened with `with sqlite3.connect()` are not closed by the context manager and linger until garbage collection, which holds a file lock there. Production sites (vacuum, diagnostics, restore inspection) close explicitly; a restore that swaps the API's database object closes the previous one; PromptModel keeps a registry so a shutdown or a test can close every instance - Windows: served files were still open when their temp directory was removed (tests now read the body), a log download test compared text mode newlines, two tests compared paths with a hard-coded separator tests/open_handles.py makes the Linux run fail the way Windows would: every database test base checks, right before deleting its files, that nothing in the process still holds them open.
810 lines
32 KiB
Python
810 lines
32 KiB
Python
"""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):
|
|
self.api = PromptManagerAPI()
|
|
|
|
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()
|
|
self.api.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()
|
|
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/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 ("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_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()
|
|
self.api.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):
|
|
self.api = PromptManagerAPI()
|
|
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()
|