nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
b50bf862f059737eab9400fae53ce5f8bc6bb361
ComfyAssets-ComfyUI_PromptM…/tests/test_api_lora.py
T
Vito 3b1d00501f fix(config): serialized config store and registry admission (#289) 
* fix(config): report settings persistence failures (#220)

Settings saves now go through one serialized config store
(py/config_store.py). A commit re-reads config.json under a process
lock, applies the change to a copy and replaces the file atomically
(temp file beside it, fsync, close, os.replace), so a failed write
leaves the old file in place.

POST /prompt_manager/settings commits first and only then publishes
the values at runtime. A write failure is 500 config-write-failed, an
unparseable config.json is 409 config-unreadable (the file is kept
instead of replaced) and a stale expected_revision is 409
settings-changed. A save reports persisted, revision, runtime_applied
and restart_required_fields; GET adds the revision. Only submitted keys
are written; other sections and the CivitAI key are kept.

* fix(lora): report integration save failures (#220)

POST /prompt_manager/lora/enable commits integrations.lora_manager
through the config store and publishes the runtime values only after
config.json holds them. Before, it set the runtime values, rewrote the
whole runtime config (dropping sections it did not know) and reported
success even when the write had failed.

A write failure is 500 config-write-failed and changes nothing at
runtime; an unparseable config.json is 409 config-unreadable. The
body goes through the capped JSON reader and the flags must be
booleans. A trigger word cache that fails to load is reported as
persisted:true, cache_loaded:false. The CivitAI key is never echoed.
The LoRA tests now write to a temporary config path.

* fix(jobs): refuse legacy record-changing operations while a job runs

The job registry gains admission classes. A records job is refused
while any records or guarded job runs, a guarded job while a records
job runs, and read jobs share only the one-job-per-kind rule.
record_changing=True stays as the spelling of records. Job summaries
add cancellable and admission, and a non-cancellable job ignores
cancel.

JobRegistry.hold() registers a running, non-cancellable job for a
legacy-page operation and ends it when the request does. The legacy
restore, maintenance, prompt cleanup, duplicate image delete, LoRA
import and batch auto-tag routes hold as records; thumbnail generation
(both routes) and thumbnail clear hold as guarded. A maintenance run
of only read checks holds as read. So a restore during a scan, or a
scan during a legacy restore, is refused with the running job named
instead of interleaving; this closes the gap noted on #269.

Refusals keep each route's shape: the JSON routes return the legacy
409 "<title> already running", and the thumbnail and LoRA streams
send their existing error event.

* test(lora): keep the preserved-section fixture free of secret values

CodeQL flagged the fixture that writes the secret constant into a test
config file as clear-text storage. The test only needs a marker value
to prove the stored key survives a save.

* fix(config): read under the commit lock and publish saves in write order

- ConfigStore.read takes the commit lock, so no read handle is open while
  config.json is replaced (Windows refuses that replace).
- os.replace is retried briefly on a Windows PermissionError (antivirus or
  indexer handles) before the save is reported as not written.
- The temp file is written beside realpath(config.json), so a symlinked
  config keeps its link and a bind mount fails truthfully.
- commit() takes a publish callback that runs under the lock after the
  write; settings and LoRA saves publish there, so two concurrent saves
  reach the runtime in write order and the runtime ends on the last one.
- The concurrency tests now wait until the second caller is blocked on the
  lock (tests/lock_probe.py) before releasing the first.

* fix(jobs): require an admission class, admit holds on entry, true Activity states

- JobRegistry.start raises TypeError when neither admission nor
  record_changing is given, so a new job never runs as "read" beside a
  scan by omission.
- hold() admits when its block is entered, so a hold that is never entered
  leaves no running records job behind; callers that answer a refusal in
  their own way check blocker() first, with no await before the block.
- A held SSE operation that ends in an error event (thumbnail generation,
  legacy batch auto-tag) is listed as failed, and refused input (duplicate
  delete without paths, an unreadable thumbnail body) is answered before
  the hold, so it is not listed as a job.
- Activity keeps a job's cancellable flag (true when an older server omits
  it) and shows Cancel only for a running, cancellable job, in Recent work
  and in the job view. A held legacy job no longer offers a Cancel that
  does nothing.

* fix(admin): legacy pages show why a request was refused

A job running in the workspace refuses the legacy record-changing requests,
and a corrupt config.json refuses a settings save. The legacy pages now
name the reason instead of hiding it:

- LoRA import: the busy and failure frames carry status and progress, and
  the admin modal handles type "error": it shows the message, notifies and
  closes, instead of staying on "Initializing..." at 0%.
- Legacy batch auto-tag: a refusal by another job is a 200 stream with one
  error frame ({type, status, message}), which both pages already render;
  a second batch keeps its exact JSON 409.
- Gallery cache clear: a refused server clear shows the reason and stops
  before the success toast and the "Not generated" reset.
- Admin settings save and maintenance show the reply's error ("config.json
  cannot be read...", "Scan for metadata already running") instead of
  "Failed to save settings" or "HTTP 409: Conflict".
- Tests cover the read-only /maintenance branch: a read-only selection runs
  during a scan as a read hold; a mixed or default selection is refused.

* refactor(api): move the settings routes out of admin.py

GET and POST /prompt_manager/settings, their validation, the config commit
and runtime publication, and the shared commit-failure reply now live in
py/api/settings_routes.py as SettingsRoutesMixin, before S3 adds host
fields there. No behaviour change: the same routes, payloads and status
codes; admin.py drops about 300 lines.

* chore(coverage): raise the floors of the frontend files this round touched

admin.js, gallery.js, activity.mjs and activity-job-view.mjs, to their
measured coverage. No floor is lowered.

* fix(admin): the gallery starts no thumbnail status read while it leaves

The gallery's startup thumbnail check runs one second after load. When that
second ended while the page was leaving (CI, WebKit: the workflow test
navigates from the gallery to the admin page), WebKit failed the new read
with "Fetch API cannot load ... due to access control checks." and
Playwright reported a page error.

Like the admin page (#290), the gallery now has a page lifetime that ends
on beforeunload and pagehide (and starts again on a back/forward restore).
The check starts no read once it has ended and passes it as the read's
signal. A page-clock workflow test that leaves before the second ends
failed on both events, in Chromium and WebKit, before this change.
2026-10-22 21:39:01 -07:00

952 lines
36 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_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"
self._env_patch = patch.dict(
os.environ, {"PROMPT_MANAGER_CONFIG_PATH": str(self.config_file)}
)
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):
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
]
path = directory / f"{stem}.safetensors.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"])
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, "find_lora_directories", 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_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, "delete_prompts_by_category", side_effect=err):
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
body = await resp.text()
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, "find_lora_directories", 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_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",
)
(loras / "neon.png").write_bytes(_png_bytes())
_write_metadata(loras, "plain")
(loras / "broken.metadata.json").write_text("{{{")
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"])), 2)
self.assertEqual(len(self.api.db.get_prompt_images(existing_id)), 0)
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"])
with patch.object(self.api.db, "save_prompt", side_effect=RuntimeError("db")):
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, "find_lora_directories", 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, "find_lora_directories", 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 both the file_name stem and the metadata-file stem
self.assertEqual(data["loras_with_trigger_words"], 2)
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