nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
d2ea705a44a2164fffb503cff0192a734d9afa90
ComfyAssets-ComfyUI_PromptM…/tests/test_api_lora.py
T
Vito d868745194 fix(lora): cache LoraManager's layout and close the #355 follow-ups (#362) (#364) 
* feat(lora): cache LoraManager's layout behind file signatures (#362, phase 1)

A new py/lora_layout_cache.py keeps one (install, layout) entry and reuses it
while the configured path, the plugin settings file's and the settings file's
signatures (st_mtime_ns, st_size), the install and every media root's real
path are unchanged and the entry is younger than 30 s. A load whose watched
file changed within 3 s (git's racy rule), could not be checked, or whose
settings were unreadable serves its request only. clear() bumps a generation
so a load that started earlier stores nothing. Hits take no lock.

Not wired yet; phase 2 puts it behind the image hot paths.

* fix(lora): treat an offline Windows share as unvalidated and make cache tests portable (#362)

Review findings on phase 1: Windows raises FileNotFoundError with winerror
53/67 for an offline share, which is a temporary failure, not a missing file.
Test times are whole milliseconds (NTFS keeps 100 ns steps) and a file system
that rounds them skips; cleanup no longer ignores errors; tests added for
ValueError, NotADirectoryError, a failed reload, resolve failures and the
watched-file set.

* perf(lora): serve the image hot paths from the layout cache (#362, phase 2)

_lora_image_dirs takes (install, layout) from lora_utils.cached_lora_manager_layout
instead of resolving the install and reading LoraManager's settings files on
every Library, inspect, reveal and serve request. serve_image builds the
allow-list in the executor instead of on the event loop. A finished import
(System job or legacy scan, both through replace_lora_imports) clears the
cache, so files linked under a folder LoraManager just created are served at
once. The integration-off check and the configured-path rule still run before
the cache.

* fix(lora): fingerprint every scanned sidecar, kept or skipped (#362, phase 3)

Discovery keeps every metadata file it saw (scanned_files, resolved and
sorted; scanned is now its length). The fingerprint hashes each one's path,
size, modification time and whether it was kept, so a skipped sidecar added,
removed or edited between the review and Start refuses Start with
plan-changed, as the review's counts promised.

* fix(lora): warn about refused default LoRA folders and log what failed (#362, phase 4)

ComfyUI's default LoRA folders already pass the folder guard; a refused one
that exists is now listed in the layout's refused_defaults, and the layout
cache warns once per newly refused folder. The cache writes one INFO line
when a stored layout differs from the last one logged. BAD_SIDECAR_ROOT and
the module docstring name a folder that contains the home folder. A relative
preview_url writes a DEBUG line before the suffix search; a non-OSError
failure while loading trigger words logs its traceback. Review follow-ups
harden two serve tests.

* fix(lora): skip a sidecar that is not UTF-8 instead of failing the load (#362)

Review follow-up on phase 4: read_lora_metadata caught JSONDecodeError but not
UnicodeDecodeError, so one non-UTF-8 sidecar failed the whole trigger-word
load (now logged with a traceback every retry) and the import review. It is
skipped like any unreadable sidecar. The relative preview_url DEBUG line uses
lazy formatting.

* docs(testing): the LoRA linked-only gate, layout cache and fingerprint scope (#362, phase 5)

TESTING.md's LoRA reference import contract now states what #358 and #362
shipped: the linked-only gate and _media_roots, linked_only media, the cover
link order, the guard on ComfyUI's default LoRA folders and refused_defaults,
the folder-unreadable code, the trigger cache's catch and 60 s retry, the
fingerprint over every scanned sidecar, and the layout cache's key,
signatures, racy window, maximum age, root re-check and clears. The targeted
command includes tests.test_lora_layout_cache and tests.test_api_images.

* fix(lora): a cached layout needs the same home folder; give the rebuild poll its time (#362)

Final review follow-ups: the folder guard compares against the home folder,
so a cache hit now also needs Path.home() unchanged. TESTING.md says the
configured-path rule runs on the cache's result and that a non-UTF-8
sidecar is skipped.

CI: 'Rebuild all is explicit...' (webkit) polled the job for the default
5 s while it regenerated 24 previews; the failure screenshot shows 12 of 24
done at the deadline on the busy runner (0.17 s locally). The poll now
allows 30 s for that real work.
2026-10-27 05:00:55 -07:00

1077 lines
42 KiB
Python
Raw Blame History

"""Route tests for the LoraManager integration endpoints.
Filesystem detection and network downloads are patched out so the tests run
anywhere without a LoraManager install or a CivitAI connection.
"""
import asyncio
import json
import os
import shutil
import sys
import tempfile
import threading
import types
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
if "folder_paths" not in sys.modules:
_stub = types.ModuleType("folder_paths")
_stub.models_dir = tempfile.gettempdir()
_stub.base_path = tempfile.gettempdir()
_stub.get_output_directory = tempfile.gettempdir
_stub.get_folder_paths = lambda name: []
sys.modules["folder_paths"] = _stub
if "server" not in sys.modules:
_server = MagicMock()
_server.PromptServer.instance.routes = MagicMock()
sys.modules["server"] = _server
from aiohttp import web
from aiohttp.test_utils import AioHTTPTestCase
from database.operations import PromptDatabase
from py.api import PromptManagerAPI
from py import config_store, lora_import, lora_utils
from py.config import IntegrationConfig, PromptManagerConfig
from py.api import settings_routes as settings_module
from py.api import lora_integration as lora_module
from py.lora_utils import get_trigger_cache
from utils.hashing import generate_prompt_hash
try:
from tests.lock_probe import ProbeLock
from tests.open_handles import assert_closed
except ImportError: # discovered with tests/ as the top-level directory
from lock_probe import ProbeLock
from open_handles import assert_closed
SECRET_KEY = "sk-civitai-super-secret-0123456789"
class LoraAPITestCase(AioHTTPTestCase):
"""Stands up the PromptManager routes with a temp DB and neutral LoRA config."""
async def get_application(self):
self._temp_db = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
self._temp_db.close()
self.tmp_root = Path(tempfile.mkdtemp())
self.custom_nodes = self.tmp_root / "custom_nodes"
self.custom_nodes.mkdir()
self._custom_nodes_patch = patch(
"py.lora_utils.custom_nodes_directories",
return_value=[self.custom_nodes],
)
self._custom_nodes_patch.start()
# Never the repository's config.json (live user data): every save in
# these tests goes to a file in this test's temp directory.
self.config_file = self.tmp_root / "config.json"
# Never the developer's LoraManager settings either (#355): an empty
# settings folder of this test's own.
self.lm_settings = self.tmp_root / "lora-manager-settings"
self.lm_settings.mkdir()
self._env_patch = patch.dict(
os.environ,
{
"PROMPT_MANAGER_CONFIG_PATH": str(self.config_file),
"LORA_MANAGER_SETTINGS_DIR": str(self.lm_settings),
},
)
self._env_patch.start()
self._config_patch = patch.multiple(
IntegrationConfig,
LORA_MANAGER_ENABLED=False,
LORA_MANAGER_PATH="",
LORA_TRIGGER_WORDS_ENABLED=False,
CIVITAI_API_KEY="",
)
self._config_patch.start()
self._detect_patch = patch(
"py.lora_utils.detect_lora_manager", return_value=None
)
self.detect_mock = self._detect_patch.start()
get_trigger_cache().clear()
app = web.Application()
routes = web.RouteTableDef()
self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name))
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._env_patch.stop()
self._detect_patch.stop()
self._config_patch.stop()
self._custom_nodes_patch.stop()
get_trigger_cache().clear()
self.api.db.close_all()
for path in (
self._temp_db.name,
self._temp_db.name + "-wal",
self._temp_db.name + "-shm",
):
if os.path.exists(path):
os.unlink(path)
assert_closed(self, self.tmp_root)
shutil.rmtree(self.tmp_root, ignore_errors=True)
def _saved(self):
return json.loads(self.config_file.read_text())
def _make_unwritable(self):
"""Point the config path under a regular file, so no write can succeed."""
blocker = self.tmp_root / "not-a-dir"
blocker.write_text("x")
self.config_file = blocker / "config.json"
os.environ["PROMPT_MANAGER_CONFIG_PATH"] = str(self.config_file)
async def _post_json(self, url, body):
return await self.client.request(
"POST",
url,
data=json.dumps(body),
headers={"Content-Type": "application/json"},
)
class TestLoraStatusNeverReturnsKey(LoraAPITestCase):
async def test_status_reports_flag_without_the_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
resp = await self.client.request("GET", "/prompt_manager/lora/status")
self.assertEqual(resp.status, 200)
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["success"])
self.assertIs(data["has_civitai_api_key"], True)
self.assertNotIn("civitai_api_key", data)
self.assertNotIn(SECRET_KEY, body)
async def test_status_flag_is_false_without_a_key(self):
resp = await self.client.request("GET", "/prompt_manager/lora/status")
data = await resp.json()
self.assertIs(data["has_civitai_api_key"], False)
self.assertNotIn("civitai_api_key", data)
def _png_bytes():
from io import BytesIO
from PIL import Image
buf = BytesIO()
Image.new("RGB", (2, 2), (1, 2, 3)).save(buf, "PNG")
return buf.getvalue()
class _FakeResponse:
def __init__(self, body):
self._body = body
self._pos = 0
self.headers = {}
def read(self, amt=-1):
if amt is None or amt < 0:
amt = len(self._body)
chunk = self._body[self._pos : self._pos + amt]
self._pos += len(chunk)
return chunk
def __enter__(self):
return self
def __exit__(self, *exc):
return False
def _make_lora_manager(where):
lm = where / "ComfyUI-Lora-Manager"
(lm / "py").mkdir(parents=True)
(lm / "__init__.py").write_text("")
return lm
def _write_metadata(directory, stem, words=None, images=None, prompt=None):
"""A LoraManager sidecar beside its model, as LoraManager names them:
``<stem>.metadata.json`` next to ``<stem>.safetensors`` (#355)."""
meta = {"file_name": f"{stem}.safetensors", "model_name": stem, "civitai": {}}
if words is not None:
meta["civitai"]["trainedWords"] = words
if images is not None:
meta["civitai"]["images"] = [
{"url": url, "meta": {"prompt": prompt} if prompt else {}} for url in images
]
(directory / f"{stem}.safetensors").write_bytes(b"model")
path = directory / f"{stem}.metadata.json"
path.write_text(json.dumps(meta))
return path
class TestLoraDetectRoute(LoraAPITestCase):
async def test_detected_path_is_public(self):
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["success"])
self.assertTrue(data["detected"])
self.assertEqual(data["path"], self.api._public_path(str(lm)))
self.assertNotIn(str(self.tmp_root), body)
async def test_not_detected(self):
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
data = await resp.json()
self.assertFalse(data["detected"])
self.assertEqual(data["path"], "")
async def test_detection_error_is_500(self):
self.detect_mock.side_effect = RuntimeError("disk gone")
resp = await self.client.request("GET", "/prompt_manager/lora/detect")
self.assertEqual(resp.status, 500)
self.assertFalse((await resp.json())["success"])
class TestLoraStatusRoute(LoraAPITestCase):
async def test_status_reports_detection_and_cache_with_public_paths(self):
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(lm)
resp = await self.client.request("GET", "/prompt_manager/lora/status")
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["enabled"])
self.assertTrue(data["detected"])
public = self.api._public_path(str(lm))
self.assertEqual(data["path"], public)
self.assertEqual(data["detected_path"], public)
self.assertFalse(os.path.isabs(public))
self.assertNotIn(str(self.tmp_root), body)
self.assertFalse(data["trigger_cache_loaded"])
self.assertTrue(data["detected_usable"])
async def test_status_says_a_detected_install_outside_custom_nodes_is_unusable(
self,
):
# S10 review N5: found (a linked install), but enabling it is refused.
outside = _make_lora_manager(self.tmp_root)
self.detect_mock.return_value = str(outside)
resp = await self.client.request("GET", "/prompt_manager/lora/status")
data = await resp.json()
self.assertEqual((data["detected"], data["detected_usable"]), (True, False))
self.detect_mock.return_value = None
resp = await self.client.request("GET", "/prompt_manager/lora/status")
self.assertFalse((await resp.json())["detected_usable"])
async def test_status_error_is_500(self):
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
resp = await self.client.request("GET", "/prompt_manager/lora/status")
self.assertEqual(resp.status, 500)
async def test_status_error_hides_the_server_path(self):
err = PermissionError(13, "Permission denied", str(self.tmp_root / "c.json"))
with patch("py.lora_utils.get_trigger_cache", side_effect=err):
resp = await self.client.request("GET", "/prompt_manager/lora/status")
self.assertEqual(resp.status, 500)
body = await resp.text()
self.assertNotIn(str(self.tmp_root), body)
self.assertIn("c.json", body)
class TestLoraEnableRoute(LoraAPITestCase):
async def test_enable_with_unknown_path_is_400(self):
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": True, "path": "/nowhere"}
)
self.assertEqual(resp.status, 400)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertFalse(self.config_file.exists())
async def test_enable_persists_and_loads_trigger_cache(self):
lm = _make_lora_manager(self.custom_nodes)
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
self.detect_mock.return_value = str(lm)
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{
"enabled": True,
"path": str(lm),
"trigger_words_enabled": True,
"civitai_api_key": SECRET_KEY,
},
)
self.assertEqual(resp.status, 200)
body = await resp.text()
data = json.loads(body)
self.assertTrue(data["success"])
self.assertEqual(data["path"], self.api._public_path(str(lm)))
self.assertNotIn(str(self.tmp_root), body)
self.assertNotIn(SECRET_KEY, body)
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertEqual(
IntegrationConfig.LORA_MANAGER_PATH,
os.path.normcase(os.path.realpath(str(lm))),
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
saved = self._saved()["integrations"]["lora_manager"]
self.assertEqual(saved["civitai_api_key"], SECRET_KEY)
self.assertEqual(saved["path"], IntegrationConfig.LORA_MANAGER_PATH)
self.assertIs(saved["enabled"], True)
self.assertIs(saved["trigger_words_enabled"], True)
self.assertEqual(
(data["persisted"], data["runtime_applied"], data["cache_loaded"]),
(True, True, True),
)
self.assertTrue(get_trigger_cache().is_loaded)
self.assertEqual(get_trigger_cache().get_trigger_words("neon"), ["glow"])
async def test_enable_accepts_the_public_relative_path(self):
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
public = self.api._public_path(str(lm))
self.assertFalse(os.path.isabs(public))
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": True, "path": public}
)
self.assertEqual(resp.status, 200)
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertEqual(
IntegrationConfig.LORA_MANAGER_PATH,
os.path.normcase(os.path.realpath(str(lm))),
)
async def test_enable_with_path_outside_custom_nodes_is_400(self):
lm = _make_lora_manager(self.tmp_root)
self.detect_mock.return_value = str(lm)
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": True, "path": str(lm)}
)
self.assertEqual(resp.status, 400)
body = await resp.text()
self.assertFalse(json.loads(body)["success"])
self.assertNotIn(str(self.tmp_root), body)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertFalse(self.config_file.exists())
async def test_enable_with_custom_node_lacking_lora_in_its_name_is_400(self):
lm = _make_lora_manager(self.custom_nodes).parent / "ComfyUI-Other"
(lm / "py").mkdir(parents=True)
(lm / "__init__.py").write_text("")
self.detect_mock.return_value = str(lm)
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": True, "path": str(lm)}
)
self.assertEqual(resp.status, 400)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
async def test_enable_rejects_an_auto_detected_path_outside_custom_nodes(self):
self.detect_mock.return_value = str(_make_lora_manager(self.tmp_root))
resp = await self._post_json("/prompt_manager/lora/enable", {"enabled": True})
self.assertEqual(resp.status, 400)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
async def test_enable_writes_the_configured_config_file(self):
config_file = self.tmp_root / "cfg" / "config.json"
with patch.dict(os.environ, {"PROMPT_MANAGER_CONFIG_PATH": str(config_file)}):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
self.assertEqual(resp.status, 200)
self.assertTrue(config_file.is_file())
saved = json.loads(config_file.read_text())["integrations"]["lora_manager"]
self.assertFalse(saved["enabled"])
self.assertTrue(saved["trigger_words_enabled"])
async def test_disable_clears_cache(self):
get_trigger_cache()._cache = {"x": ["y"]}
get_trigger_cache()._loaded = True
resp = await self._post_json("/prompt_manager/lora/enable", {"enabled": False})
self.assertEqual(resp.status, 200)
self.assertFalse(get_trigger_cache().is_loaded)
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
async def test_empty_key_keeps_the_stored_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": False, "civitai_api_key": ""}
)
self.assertEqual(resp.status, 200)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
async def test_absent_key_keeps_the_stored_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
await self._post_json("/prompt_manager/lora/enable", {"enabled": False})
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
async def test_new_key_replaces_the_stored_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "civitai_api_key": " new-key "},
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "new-key")
async def test_clear_flag_removes_the_stored_key(self):
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "clear_civitai_api_key": True},
)
self.assertEqual(resp.status, 200)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "")
async def test_non_string_key_is_400(self):
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": False, "civitai_api_key": 123}
)
self.assertEqual(resp.status, 400)
async def test_persist_failure_is_500(self):
self._make_unwritable()
with self.assertLogs("prompt_manager.api", level="ERROR"):
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": False}
)
self.assertEqual(resp.status, 500)
self.assertEqual((await resp.json())["code"], "config-write-failed")
async def test_enable_write_failure_reports_and_keeps_runtime(self):
# #220: a save that cannot be written publishes nothing and says so.
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
self._make_unwritable()
with self.assertLogs("prompt_manager.api", level="ERROR"):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{
"enabled": True,
"path": str(lm),
"trigger_words_enabled": True,
"clear_civitai_api_key": True,
},
)
text = await resp.text()
body = json.loads(text)
self.assertEqual(resp.status, 500)
self.assertEqual(body["code"], "config-write-failed")
self.assertEqual((body["persisted"], body["runtime_applied"]), (False, False))
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertEqual(IntegrationConfig.LORA_MANAGER_PATH, "")
self.assertFalse(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, SECRET_KEY)
self.assertFalse(get_trigger_cache().is_loaded)
self.assertNotIn(SECRET_KEY, text)
self.assertNotIn(str(self.tmp_root), text)
async def test_enable_preserves_unknown_sections(self):
stored = "stored-value" # a plain marker: what is on disk must stay
self.config_file.write_text(
json.dumps(
{
"custom_section": {"keep": True},
"web_ui": {"result_timeout": 11},
"integrations": {
"lora_manager": {"civitai_api_key": stored, "extra": 1},
"other_tool": {"on": True},
},
}
)
)
IntegrationConfig.CIVITAI_API_KEY = stored
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
self.assertEqual(resp.status, 200)
saved = self._saved()
self.assertEqual(saved["custom_section"], {"keep": True})
self.assertEqual(saved["web_ui"], {"result_timeout": 11})
self.assertEqual(saved["integrations"]["other_tool"], {"on": True})
lora = saved["integrations"]["lora_manager"]
self.assertEqual(lora["extra"], 1)
self.assertEqual(lora["civitai_api_key"], stored)
self.assertIs(lora["trigger_words_enabled"], True)
self.assertNotIn("database", saved) # no runtime defaults written
async def test_enable_refuses_a_corrupt_config_file(self):
self.config_file.write_bytes(b"{not json")
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
body = await resp.json()
self.assertEqual((resp.status, body["code"]), (409, "config-unreadable"))
self.assertEqual(self.config_file.read_bytes(), b"{not json")
self.assertFalse(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED)
async def test_enable_with_a_stale_revision_changes_nothing(self):
# System S10: the Save row sends the revision it read; a file changed
# since is a settings-changed conflict and nothing is written.
self.config_file.write_text(json.dumps({"web_ui": {"result_timeout": 9}}))
resp = await self.client.request("GET", "/prompt_manager/settings")
revision = (await resp.json())["revision"]
self.config_file.write_text(json.dumps({"web_ui": {"result_timeout": 10}}))
before = self.config_file.read_bytes()
resp = await self._post_json(
"/prompt_manager/lora/enable",
{
"enabled": False,
"trigger_words_enabled": True,
"expected_revision": revision,
},
)
body = await resp.json()
self.assertEqual((resp.status, body["code"]), (409, "settings-changed"))
self.assertEqual(self.config_file.read_bytes(), before)
self.assertFalse(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED)
resp = await self._post_json(
"/prompt_manager/lora/enable",
{
"enabled": False,
"trigger_words_enabled": True,
"expected_revision": body["revision"],
},
)
saved = await resp.json()
self.assertEqual(resp.status, 200)
self.assertTrue(saved["revision"].startswith("sha256:"))
self.assertTrue(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED)
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "expected_revision": 7},
)
self.assertEqual(resp.status, 400)
async def test_enable_requires_boolean_flags_and_body_cap(self):
for body in (
{"enabled": "yes"},
{"enabled": False, "trigger_words_enabled": 1},
{"enabled": False, "clear_civitai_api_key": "true"},
):
with self.subTest(body=body):
resp = await self._post_json("/prompt_manager/lora/enable", body)
self.assertEqual(resp.status, 400)
resp = await self._post_json("/prompt_manager/lora/enable", [True])
self.assertEqual(resp.status, 400)
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "civitai_api_key": "k" * (2 * 1024 * 1024)},
)
self.assertEqual(resp.status, 413)
self.assertFalse(self.config_file.exists())
self.assertFalse(IntegrationConfig.LORA_MANAGER_ENABLED)
async def test_blank_key_unchanged_and_explicit_clear(self):
self.config_file.write_text(
json.dumps({"integrations": {"lora_manager": {"civitai_api_key": "old"}}})
)
IntegrationConfig.CIVITAI_API_KEY = "old"
resp = await self._post_json(
"/prompt_manager/lora/enable", {"enabled": False, "civitai_api_key": " "}
)
self.assertEqual(resp.status, 200)
self.assertEqual(
self._saved()["integrations"]["lora_manager"]["civitai_api_key"], "old"
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "old")
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "clear_civitai_api_key": True},
)
self.assertEqual(resp.status, 200)
self.assertEqual(
self._saved()["integrations"]["lora_manager"]["civitai_api_key"], ""
)
self.assertEqual(IntegrationConfig.CIVITAI_API_KEY, "")
async def test_cache_load_failure_reports_persisted_but_unloaded(self):
lm = _make_lora_manager(self.custom_nodes)
self.detect_mock.return_value = str(lm)
with (
patch.object(
type(get_trigger_cache()), "load", side_effect=OSError("unreadable")
),
self.assertLogs("prompt_manager.api", level="WARNING"),
):
resp = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": True, "path": str(lm), "trigger_words_enabled": True},
)
text = await resp.text()
body = json.loads(text)
self.assertEqual(resp.status, 200)
self.assertIs(body["success"], True)
self.assertEqual((body["persisted"], body["cache_loaded"]), (True, False))
self.assertTrue(IntegrationConfig.LORA_MANAGER_ENABLED)
self.assertIs(self._saved()["integrations"]["lora_manager"]["enabled"], True)
self.assertNotIn("unreadable", text)
class TestConfigWritersShareOneStore(LoraAPITestCase):
async def asyncSetUp(self):
await super().asyncSetUp()
self._orig_timeout = PromptManagerConfig.RESULT_TIMEOUT
async def asyncTearDown(self):
PromptManagerConfig.RESULT_TIMEOUT = self._orig_timeout
await super().asyncTearDown()
async def test_settings_and_lora_saves_do_not_overwrite_each_other(self):
# #220: the LoRA save is sent while the settings save is inside its
# commit (lock held, file already read). The LoRA commit re-reads the
# file under the same lock, so both edits survive and nothing else
# is lost (before, the LoRA save rewrote the whole runtime config).
self.config_file.write_text(json.dumps({"custom_section": {"keep": True}}))
inside = threading.Event()
release = threading.Event()
real_mutation = settings_module._settings_mutation
def gated_mutation(updates, new_roots):
mutate = real_mutation(updates, new_roots)
def gated(config):
inside.set()
release.wait(10)
return mutate(config)
return gated
lora_calls = []
real_lora_mutation = lora_module._lora_settings_mutation
def recorded_lora_mutation(values):
mutate = real_lora_mutation(values)
def recorded(config):
lora_calls.append(dict(config))
return mutate(config)
return recorded
probe = ProbeLock()
loop = asyncio.get_running_loop()
with (
patch.object(settings_module, "_settings_mutation", gated_mutation),
patch.object(
lora_module, "_lora_settings_mutation", recorded_lora_mutation
),
patch.object(config_store, "_COMMIT_LOCK", probe),
):
settings = asyncio.ensure_future(
self._post_json("/prompt_manager/settings", {"result_timeout": 13})
)
self.assertTrue(await loop.run_in_executor(None, inside.wait, 10))
lora = asyncio.ensure_future(
self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
)
# Wait until the LoRA commit is blocked on the lock the settings
# commit holds; its mutation must not have run yet.
self.assertTrue(await loop.run_in_executor(None, probe.contended.wait, 10))
self.assertEqual(lora_calls, [])
release.set()
settings_resp, lora_resp = await asyncio.gather(settings, lora)
self.assertEqual(len(lora_calls), 1)
self.assertEqual(lora_calls[0]["web_ui"]["result_timeout"], 13)
self.assertEqual((settings_resp.status, lora_resp.status), (200, 200))
saved = self._saved()
self.assertEqual(saved["custom_section"], {"keep": True})
self.assertEqual(saved["web_ui"]["result_timeout"], 13)
self.assertIs(
saved["integrations"]["lora_manager"]["trigger_words_enabled"], True
)
settings_body = await settings_resp.json()
lora_body = await lora_resp.json()
self.assertNotEqual(settings_body["revision"], lora_body["revision"])
def _delay_first_commit_result(self):
"""Hold back the first config commit's result from its handler.
The commit itself has run (the file is written); only the handler's
continuation waits until the returned event is set. Before #220's fix
round the handler published after this point, so a later save could
publish first and the earlier one overwrite it at runtime.
"""
original = self.api._run_in_executor
written = asyncio.Event()
resume = asyncio.Event()
delayed = []
async def run(func, *args, **kwargs):
result = await original(func, *args, **kwargs)
if getattr(func, "__func__", None) is config_store.ConfigStore.commit:
if not delayed:
delayed.append(func)
written.set()
await resume.wait()
return result
patcher = patch.object(self.api, "_run_in_executor", run)
patcher.start()
self.addCleanup(patcher.stop)
return written, resume
async def test_concurrent_settings_saves_end_on_the_last_written(self):
written, resume = self._delay_first_commit_result()
first = asyncio.ensure_future(
self._post_json("/prompt_manager/settings", {"result_timeout": 11})
)
await asyncio.wait_for(written.wait(), 10)
second = await self._post_json(
"/prompt_manager/settings", {"result_timeout": 22}
)
resume.set()
first = await first
self.assertEqual((first.status, second.status), (200, 200))
self.assertEqual(self._saved()["web_ui"]["result_timeout"], 22)
self.assertEqual(PromptManagerConfig.RESULT_TIMEOUT, 22)
async def test_concurrent_lora_saves_end_on_the_last_written(self):
written, resume = self._delay_first_commit_result()
first = asyncio.ensure_future(
self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": True},
)
)
await asyncio.wait_for(written.wait(), 10)
second = await self._post_json(
"/prompt_manager/lora/enable",
{"enabled": False, "trigger_words_enabled": False},
)
resume.set()
first = await first
self.assertEqual((first.status, second.status), (200, 200))
lora = self._saved()["integrations"]["lora_manager"]
self.assertIs(lora["trigger_words_enabled"], False)
self.assertIs(IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED, False)
class TestLoraScanRoute(LoraAPITestCase):
def _enable(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(_make_lora_manager(self.custom_nodes))
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
async def test_configured_path_outside_custom_nodes_is_400(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(_make_lora_manager(self.tmp_root))
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
self.assertEqual(resp.status, 400)
body = await resp.text()
self.assertFalse(json.loads(body)["success"])
self.assertNotIn(str(self.tmp_root), body)
async def test_scan_error_hides_the_server_path(self):
self._enable()
err = PermissionError(13, "Permission denied", str(self.tmp_root / "x.db"))
with patch.object(self.api.db, "replace_category_prompts", side_effect=err):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
body = await resp.text() # the stream ends inside the patch
self.assertNotIn(str(self.tmp_root), body)
self.assertIn("x.db", body)
# The admin page's modal shows the failure as its status and closes.
error = self._sse_events(body)[-1]
self.assertEqual(
(error["type"], error["status"], error["progress"]),
("error", error["message"], 0),
)
@staticmethod
def _sse_events(text):
return [
json.loads(line[len("data: ") :])
for line in text.splitlines()
if line.startswith("data: ")
]
async def _scan(self, loras):
opener_calls = []
def fake_open(req, timeout=None):
opener_calls.append(req)
return _FakeResponse(_png_bytes())
cache_dir = self.tmp_root / "img_cache"
with (
patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]),
patch.object(
lora_utils, "get_lora_image_cache_dir", return_value=cache_dir
),
patch.object(lora_utils, "_open_url", fake_open),
):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
text = await resp.text()
return resp, self._sse_events(text), opener_calls
async def test_not_enabled_is_400(self):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
self.assertEqual(resp.status, 400)
async def test_no_path_is_400(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
self.assertEqual(resp.status, 400)
async def test_a_finished_scan_clears_the_layout_cache(self):
# #362: the legacy scan replaces through the same method.
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "plain")
(loras / "plain.png").write_bytes(_png_bytes())
with patch("py.lora_utils.clear_layout_cache") as clear:
resp, events, _calls = await self._scan(loras)
self.assertEqual(events[-1]["type"], "complete", events)
clear.assert_called_once_with()
async def test_scan_imports_loras_with_images_and_links_existing(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(
loras,
"neon",
words=["glow"],
images=["https://civitai.com/a.png", "https://evil.example/b.png"],
prompt="neon street",
)
# No local image: its CivitAI images are downloaded (#355).
_write_metadata(loras, "plain")
(loras / "broken.metadata.json").write_text("{{{")
(loras / "broken.safetensors").write_bytes(b"model")
existing_id = self.api.db.save_prompt(
text="plain",
category=None,
tags=[],
rating=None,
prompt_hash=generate_prompt_hash("plain"),
)
resp, events, opener_calls = await self._scan(loras)
self.assertEqual(resp.status, 200)
done = events[-1]
self.assertEqual(done["type"], "complete")
self.assertEqual(done["total"], 3)
self.assertEqual(done["imported"], 1)
self.assertEqual(done["skipped"], 2)
self.assertEqual(
[r.full_url for r in opener_calls], ["https://civitai.com/a.png"]
)
self.assertEqual(
opener_calls[0].get_header("Authorization"), f"Bearer {SECRET_KEY}"
)
imported = self.api.db.search_prompts(category="lora-manager")
self.assertEqual(len(imported), 1)
self.assertEqual(imported[0]["text"], "neon street")
self.assertIn("lora:neon", imported[0]["tags"])
self.assertIn("glow", imported[0]["tags"])
self.assertEqual(len(self.api.db.get_prompt_images(imported[0]["id"])), 1)
self.assertEqual(len(self.api.db.get_prompt_images(existing_id)), 0)
# The unreadable sidecar is reported by its file name.
statuses = [e.get("status") for e in events if e.get("processed")]
self.assertIn("broken.metadata.json", statuses)
async def test_scan_links_local_media_without_downloading(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(
loras, "neon", images=["https://civitai.com/a.png"], prompt="neon street"
)
(loras / "neon.preview.webp").write_bytes(b"x")
resp, events, opener_calls = await self._scan(loras)
self.assertEqual(events[-1]["imported"], 1)
self.assertEqual(opener_calls, [])
imported = self.api.db.search_prompts(category="lora-manager")
images = self.api.db.get_prompt_images(imported[0]["id"])
self.assertEqual(
[os.path.basename(i["image_path"]) for i in images], ["neon.preview.webp"]
)
async def test_scan_reports_unusable_settings_without_replacing(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon")
(self.lm_settings / "settings.json").write_text("{not json")
old_id = self.api.db.save_prompt(
text="kept",
category="lora-manager",
tags=[],
rating=None,
prompt_hash=generate_prompt_hash("kept"),
)
resp, events, _ = await self._scan(loras)
self.assertEqual(resp.status, 200)
self.assertEqual(events[-1]["type"], "error")
self.assertEqual(
events[-1]["message"],
"LoraManager settings: not usable: the settings file is not valid "
"JSON. Nothing was changed.",
)
kept = self.api.db.search_prompts(category="lora-manager")
self.assertEqual([p["id"] for p in kept], [old_id])
async def test_save_failure_is_counted_as_skipped(self):
self._enable()
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
# Preparing one LoRA's record fails: it is skipped, the import goes on.
with patch.object(
lora_import, "build_record", side_effect=RuntimeError("bad record")
):
resp, events, _ = await self._scan(loras)
self.assertEqual(events[-1]["imported"], 0)
self.assertEqual(events[-1]["skipped"], 1)
class TestLoraTriggerWordRoutes(LoraAPITestCase):
async def test_trigger_words_not_enabled_is_400(self):
resp = await self.client.request(
"GET", "/prompt_manager/lora/trigger-words?name=x"
)
self.assertEqual(resp.status, 400)
async def test_trigger_words_missing_name_is_400(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
resp = await self.client.request("GET", "/prompt_manager/lora/trigger-words")
self.assertEqual(resp.status, 400)
async def test_trigger_words_loads_cache_on_demand(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon", words=["glow", "bright"])
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self.client.request(
"GET", "/prompt_manager/lora/trigger-words?name=NEON"
)
data = await resp.json()
self.assertTrue(data["success"])
self.assertEqual(data["trigger_words"], ["glow", "bright"])
self.assertTrue(get_trigger_cache().is_loaded)
async def test_trigger_words_error_is_500(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
resp = await self.client.request(
"GET", "/prompt_manager/lora/trigger-words?name=x"
)
self.assertEqual(resp.status, 500)
async def test_refresh_not_enabled_is_400(self):
resp = await self.client.request("POST", "/prompt_manager/lora/refresh-cache")
self.assertEqual(resp.status, 400)
async def test_refresh_without_path_is_400(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
resp = await self.client.request("POST", "/prompt_manager/lora/refresh-cache")
self.assertEqual(resp.status, 400)
async def test_refresh_reloads_cache(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
loras = self.tmp_root / "loras"
loras.mkdir()
_write_metadata(loras, "neon", words=["glow"])
with patch.object(lora_utils, "default_lora_roots", return_value=[str(loras)]):
resp = await self.client.request(
"POST", "/prompt_manager/lora/refresh-cache"
)
data = await resp.json()
self.assertTrue(data["success"])
# Keyed by the model name and the sidecar name, which are the same
# for LoraManager's alongside layout.
self.assertEqual(data["loras_with_trigger_words"], 1)
async def test_refresh_error_is_500(self):
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(self.tmp_root)
with patch("py.lora_utils.get_trigger_cache", side_effect=RuntimeError("x")):
resp = await self.client.request(
"POST", "/prompt_manager/lora/refresh-cache"
)
self.assertEqual(resp.status, 500)
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink