nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(lora): recheck the files before the replace and stop an abandoned scan

Browse Source 
Review fixes for the staged LoRA import (#221):

- The metadata files are listed and fingerprinted again right before the
  replace transaction, so a file changed, removed or added while the
  downloads ran aborts the import with nothing deleted (System job:
  plan-changed, Review again; legacy scan: an error frame in its own
  words). The tests now change real files behind a held download.
- A legacy scan whose page closes stops staging before its next LoRA and
  keeps its records hold until staging has stopped; CivitAI images are
  written to a temp file beside the cache entry and moved in place, so a
  half-written image is never seen as cached.
- Start refuses with plan-changed when the integration was turned off or
  moved to another install since the review.
- A failed job keeps the skipped count and the names of the LoRAs that
  could not be prepared.
- The replace binds the category instead of one parameter per prompt
  (SQLite's variable limit), /lora/status says whether a detected install
  can be used, and the legacy scan logs expected refusals as warnings.
This commit is contained in:
Vito Sansevero
2026-10-25 08:36:46 -07:00
parent e3ee34bfd0
commit d247e4c2e7
8 changed files with 444 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
database/category_import.py
+46 -18
View File
@@ -17,9 +17,11 @@ from typing import Any, Dict, List, Optional, Sequence
try:
from .maintenance import PlanChanged
from .pairs import PairDependencyError
from .workflow_documents import documents_of_rows, release_documents
except ImportError: # pragma: no cover - flat test imports
from database.maintenance import PlanChanged
from database.pairs import PairDependencyError
from database.workflow_documents import documents_of_rows, release_documents
@@ -27,6 +29,9 @@ class EmptyReplacement(ValueError):
"""The replacement has no record and would remove every current import."""
# The category's prompt ids, as a subquery (one bound parameter).
_IN_CATEGORY = "SELECT id FROM prompts WHERE category = ?"
# Each record's outcome and the result counter it adds to.
_COUNTERS = {"imported": "imported", "linked": "linked", "duplicate": "duplicates"}
@@ -49,15 +54,12 @@ class CategoryImportOperations:
(category,),
).fetchall()
]
paired = 0
if ids:
marks = _placeholders(ids)
paired = conn.execute(
f"SELECT COUNT(*) FROM prompts WHERE id IN ({marks}) AND id IN "
"(SELECT positive_prompt_id FROM prompt_pairs UNION "
"SELECT negative_prompt_id FROM prompt_pairs)",
ids,
).fetchone()[0]
paired = conn.execute(
"SELECT COUNT(*) FROM prompts WHERE category = ? AND id IN "
"(SELECT positive_prompt_id FROM prompt_pairs UNION "
"SELECT negative_prompt_id FROM prompt_pairs)",
(category,),
).fetchone()[0]
return {"ids": ids, "paired": paired}
def existing_prompt_hashes(
@@ -113,8 +115,8 @@ class CategoryImportOperations:
raise EmptyReplacement(
"Nothing to import; the previous imports were kept"
)
self._ensure_prompts_unpaired(conn, ids)
removed = self._delete_category_rows(conn, ids)
self._ensure_category_unpaired(conn, category)
removed = self._delete_category_rows(conn, category, ids)
result = {
"removed": removed,
"imported": 0,
@@ -128,16 +130,42 @@ class CategoryImportOperations:
self._import_record(conn, category, record, created, result)
return result
def _delete_category_rows(self, conn, ids: List[int]) -> int:
"""Delete prompts ``ids`` with their image links and tags (no commit)."""
def _delete_category_rows(self, conn, category: str, ids: List[int]) -> int:
"""Delete the category's prompts with their image links and tags.
No commit. The rows are selected by category (just verified against
``ids`` under the write lock), so no statement binds one parameter per
prompt and SQLite's variable limit never caps a collection.
"""
if not ids:
return 0
marks = _placeholders(ids)
documents = documents_of_rows(conn, "prompt_id", ids)
conn.execute(f"DELETE FROM generated_images WHERE prompt_id IN ({marks})", ids)
documents = documents_of_rows(conn, "prompt_id", ids) # chunked
conn.execute(
f"DELETE FROM generated_images WHERE prompt_id IN ({_IN_CATEGORY})",
(category,),
)
release_documents(conn, documents)
conn.execute(f"DELETE FROM prompt_tags WHERE prompt_id IN ({marks})", ids)
return conn.execute(f"DELETE FROM prompts WHERE id IN ({marks})", ids).rowcount
conn.execute(
f"DELETE FROM prompt_tags WHERE prompt_id IN ({_IN_CATEGORY})",
(category,),
)
return conn.execute(
"DELETE FROM prompts WHERE category = ?", (category,)
).rowcount
@staticmethod
def _ensure_category_unpaired(conn, category: str) -> None:
"""Raise PairDependencyError when a saved pair uses a category prompt."""
row = conn.execute(
f"SELECT id FROM prompt_pairs WHERE positive_prompt_id IN ({_IN_CATEGORY}) "
f"OR negative_prompt_id IN ({_IN_CATEGORY}) LIMIT 1",
(category, category),
).fetchone()
if row:
raise PairDependencyError(
f"Prompt belongs to saved pair {row[0]}; delete the pair first "
"or save changed text as a new prompt"
)
def _import_record(self, conn, category, record, created, result) -> None:
"""Save one staged record on ``conn``, or link it to the prompt it repeats."""
py/api/lora_import_routes.py
+32 -7
View File
@@ -67,6 +67,10 @@ BLOCKED = (
"Saved pairs use some imported LoRA prompts, so the collection cannot be "
"replaced; the previous imports were kept."
)
INTEGRATION_CHANGED = (
"The LoraManager integration changed since the review; nothing was changed. "
"Review the import again."
)
EMPTY = "No LoRA could be prepared for import; the previous imports were kept."
FAILED = "The import failed; the previous LoRA imports were kept."
CACHE_NOTE = "Images already downloaded stay in the image cache."
@@ -148,6 +152,18 @@ class LoraImportRoutesMixin:
)
)
async def lora_files_unchanged(self, lm_path, expected):
"""Whether the metadata files under ``lm_path`` still fingerprint as
``expected`` (listed again, so an added file counts too)."""
def current():
return lora_import.fingerprint(lora_import.discover(lm_path))
try:
return await self._run_in_executor(current) == expected
except OSError:
return False
async def replace_lora_imports(self, db, staged, *, expected_ids, allow_empty):
"""Swap the collection for ``staged.records`` in one transaction."""
return await self._run_in_executor(
@@ -290,6 +306,11 @@ class LoraImportRoutesMixin:
"storage-changed",
)
lm_path = plan.payload["path"]
# M5: the integration must still be on and at the reviewed install.
current, _refused = await self._lora_source()
if current != lm_path:
store.discard(token)
return _failure(INTEGRATION_CHANGED, 409, "plan-changed")
try:
files = await self._run_in_executor(lora_import.discover, lm_path)
now = await self._run_in_executor(lora_import.fingerprint, files)
@@ -353,6 +374,12 @@ class LoraImportRoutesMixin:
self.logger.exception("LoRA import staging failed")
return _fail_job(job, "import-failed", f"{FAILED} {CACHE_NOTE}")
self._count_downloads(job, staged)
# M4: what staging skipped stays on the job whatever happens next.
job.record_counted(
"skipped",
len(staged.skipped),
[{"path": "", "detail": name} for name in staged.skipped],
)
# Past this point nothing stops halfway: the replace is one transaction.
# Published, so a view that offered Cancel during the downloads drops it.
job.cancellable = False
@@ -362,11 +389,14 @@ class LoraImportRoutesMixin:
job.update(status_line=CANCELLED)
job.finish("cancelled")
return
if staged.fingerprint != plan.fingerprint:
return _fail_job(job, "plan-changed", SOURCE_CHANGED)
ok, _backup = await self._safety_backup_step(job, db, backup_on)
if not ok:
return
# I1: the files are listed and fingerprinted again right before the
# replace, so a file changed, removed or added while the downloads ran
# (LoraManager refreshing its CivitAI info) aborts with nothing deleted.
if not await self.lora_files_unchanged(plan.payload["path"], plan.fingerprint):
return _fail_job(job, "plan-changed", f"{SOURCE_CHANGED} {CACHE_NOTE}")
job.update(status_line="Replacing the LoRA collection", progress=95)
await self._replace_in_job(job, db, plan, staged)
@@ -406,11 +436,6 @@ class LoraImportRoutesMixin:
job.record_counted(
"linked", result["linked"] + result["duplicates"], rows["linked"]
)
job.record_counted(
"skipped",
len(staged.skipped),
[{"path": "", "detail": name} for name in staged.skipped],
)
line = _import_line(result, len(staged.skipped))
if result["removed"]:
job.log_line(f"Replaced {_count(result['removed'], 'previous import')}")
py/api/lora_integration.py
+110 -18
View File
@@ -1,7 +1,9 @@
"""LoraManager integration API routes for PromptManager."""
import asyncio
import contextlib
import json
import threading
from aiohttp import web
@@ -37,6 +39,21 @@ def _lora_settings_mutation(values):
return mutate
# The admin page's own words for the refusals it can meet (it has no review).
_FILES_CHANGED = (
"The LoRA metadata files changed during the scan; the previous imports "
"were kept. Run the scan again."
)
_IMPORTS_CHANGED = (
"The imported LoRA prompts changed during the scan; the previous imports "
"were kept. Run the scan again."
)
class LoraFilesChanged(Exception):
"""A metadata file changed, appeared or went while the legacy scan staged."""
def _lora_error_frame(message):
"""An import stream's error event, as the admin page's modal reads it.
@@ -46,6 +63,16 @@ def _lora_error_frame(message):
return {"type": "error", "status": message, "progress": 0, "message": message}
def _closed(request):
"""A check (safe off the loop) for whether ``request``'s client left."""
def closed():
transport = request.transport
return transport is None or transport.is_closing()
return closed
class LoraIntegrationMixin:
"""Mixin providing LoraManager detection, scanning, and trigger word endpoints."""
@@ -108,7 +135,11 @@ class LoraIntegrationMixin:
"""Get current LoraManager integration status."""
try:
from ..config import IntegrationConfig
from ..lora_utils import detect_lora_manager, get_trigger_cache
from ..lora_utils import (
detect_lora_manager,
get_trigger_cache,
resolve_lora_manager_path,
)
config = IntegrationConfig.get_config()["lora_manager"]
cache = get_trigger_cache()
@@ -117,6 +148,12 @@ class LoraIntegrationMixin:
detected_path = await self._run_in_executor(
detect_lora_manager, config.get("path", "")
)
# A detected install reached through a link outside custom_nodes is
# found but cannot be enabled (S10 review N5): say which.
usable = bool(detected_path) and (
await self._run_in_executor(resolve_lora_manager_path, detected_path)
is not None
)
return web.json_response(
{
@@ -126,6 +163,7 @@ class LoraIntegrationMixin:
"trigger_words_enabled": config["trigger_words_enabled"],
"has_civitai_api_key": bool(config.get("civitai_api_key")),
"detected": detected_path is not None,
"detected_usable": usable,
"detected_path": (
self._public_path(detected_path) if detected_path else ""
),
@@ -316,20 +354,56 @@ class LoraIntegrationMixin:
) as job:
await response.prepare(request)
try:
await self._lora_scan_stream(allowed, send_progress)
await self._lora_scan_stream(
allowed, send_progress, disconnected=_closed(request)
)
except Exception as e:
self.logger.exception("LoRA scan failed")
job.finish("failed")
await send_progress(_lora_error_frame(self._public_error(e)))
await response.write_eof()
await self._send_scan_failure(e, send_progress)
with contextlib.suppress(ConnectionError, RuntimeError):
await response.write_eof()
return response
except Exception as e:
self.logger.error(f"LoRA scan failed: {e}")
return self._error_response(e)
async def _lora_scan_stream(self, lm_path, send_progress):
async def _send_scan_failure(self, error, send_progress):
"""The error frame for a failed legacy scan; quiet when the page left.
Expected refusals (files or imports changed, saved pairs, nothing to
import) are warnings in the admin page's own words; anything else is
logged with its traceback.
"""
from .. import lora_import
from .lora_import_routes import (
EmptyReplacement,
PairDependencyError,
PlanChanged,
)
if isinstance(error, (ConnectionError, lora_import.ImportCancelled)):
self.logger.info("LoRA scan stopped: the page closed its connection")
return
if isinstance(error, PlanChanged):
message = _IMPORTS_CHANGED
elif isinstance(error, LoraFilesChanged):
message = str(error)
else:
message = self._public_error(error)
if isinstance(
error,
(PlanChanged, LoraFilesChanged, EmptyReplacement, PairDependencyError),
):
self.logger.warning(f"LoRA scan refused: {message}")
else:
self.logger.exception("LoRA scan failed")
with contextlib.suppress(ConnectionError, RuntimeError):
await send_progress(_lora_error_frame(message))
async def _lora_scan_stream(
self, lm_path, send_progress, disconnected=lambda: False
):
"""Import every LoRA under ``lm_path``, reporting through ``send_progress``.
The staged path of the System import (#221, ``lora_import_routes``):
@@ -358,6 +432,7 @@ class LoraIntegrationMixin:
}
)
reviewed = await self._run_in_executor(lora_import.fingerprint, files)
events = asyncio.Queue()
state = await self._run_in_executor(
self.db.category_import_state, lora_import.LORA_CATEGORY
@@ -397,20 +472,37 @@ class LoraIntegrationMixin:
# Each record's event is queued (through the loop) before staging
# ends, so reporting until staging is done and then draining the
# queue reports every record in order.
# queue reports every record in order. I2: when the page goes away
# (its connection closes, or a report cannot be sent) staging stops
# before its next LoRA, and this returns only once it has stopped, so
# the caller's records hold covers every download.
stop = threading.Event()
staging = asyncio.ensure_future(
self.stage_lora_import(files, on_record=on_record)
self.stage_lora_import(
files,
on_record=on_record,
cancelled=lambda: stop.is_set() or disconnected(),
)
)
while not staging.done():
getter = asyncio.ensure_future(events.get())
await asyncio.wait({getter, staging}, return_when=asyncio.FIRST_COMPLETED)
if getter.done():
await report(*getter.result())
else:
getter.cancel()
while not events.empty():
await report(*events.get_nowait())
try:
while not staging.done():
getter = asyncio.ensure_future(events.get())
await asyncio.wait(
{getter, staging}, return_when=asyncio.FIRST_COMPLETED
)
if getter.done():
await report(*getter.result())
else:
getter.cancel()
while not events.empty():
await report(*events.get_nowait())
finally:
stop.set()
with contextlib.suppress(Exception):
await asyncio.shield(staging)
staged = await staging
if not await self.lora_files_unchanged(lm_path, reviewed):
raise LoraFilesChanged(_FILES_CHANGED)
await send_progress(
{
py/lora_import.py
+5 -4
View File
@@ -12,8 +12,10 @@ only the last writes to the database:
downloaded into the image cache first (best effort, counted). A cancelled
or failed stage leaves the database untouched; downloaded images stay in
the cache.
3. **replace**: ``PromptDatabase.replace_category_prompts`` swaps the old
collection for the staged records in one transaction (or changes nothing).
3. **replace**: the files are listed and fingerprinted once more (a file
changed, removed or added while staging aborts the import), then
``PromptDatabase.replace_category_prompts`` swaps the old collection for
the staged records in one transaction (or changes nothing).
Both the System job (``POST /lora/import/apply``) and the legacy ``POST
/lora/scan`` run these stages (``py/api/lora_integration.py``).
@@ -126,7 +128,6 @@ class StagedImport:
stats: Dict[str, int] = field(
default_factory=lambda: {"downloaded": 0, "download_errors": 0}
)
fingerprint: str = ""
def stage(
@@ -145,7 +146,7 @@ def stage(
ImportCancelled); ``on_record(index, record_or_None, staged)`` follows
each one.
"""
staged = StagedImport(total=len(files), fingerprint=fingerprint(files))
staged = StagedImport(total=len(files))
for index, path in enumerate(files):
if cancelled():
raise ImportCancelled()
py/lora_utils.py
+29 -5
View File
@@ -514,17 +514,41 @@ def _download_one(url: str, local_path: Path, api_key: str) -> Optional[str]:
from PIL import Image
img = Image.open(BytesIO(raw))
img.thumbnail((_THUMB_MAX_SIZE, _THUMB_MAX_SIZE), Image.LANCZOS)
img = img.convert("RGB")
img.save(str(local_path), "JPEG", quality=85)
with Image.open(BytesIO(raw)) as img:
img.thumbnail((_THUMB_MAX_SIZE, _THUMB_MAX_SIZE), Image.LANCZOS)
rgb = img.convert("RGB")
_save_atomically(rgb, local_path)
return str(local_path.resolve())
except Exception as e:
logger.debug(f"Failed to download {url}: {e}")
return None
def _save_atomically(image, local_path: Path) -> None:
"""Write ``image`` as JPEG so ``local_path`` only ever names a whole file.
The image goes to a temp file in the same folder, which is closed before
``os.replace`` puts it in place (Windows cannot replace an open file); a
failure removes the temp file. A reader (another import's cache check)
never sees a half-written image.
"""
import tempfile
fd, temp = tempfile.mkstemp(
dir=str(local_path.parent), prefix=".pm-", suffix=".part"
)
try:
with os.fdopen(fd, "wb") as handle:
image.save(handle, "JPEG", quality=85)
os.replace(temp, local_path)
except BaseException:
try:
os.remove(temp)
except OSError:
pass
raise
def _civitai_cache_plan(
metadata: Dict, metadata_path: Path, cache_dir: Path
) -> Tuple[Path, List[str], List[Tuple[str, Path]]]:
tests/test_api_lora.py
+15
View File
@@ -252,6 +252,21 @@ class TestLoraStatusRoute(LoraAPITestCase):
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")):
tests/test_lora_import.py
+125 -11
View File
@@ -195,6 +195,31 @@ class CategoryReplacementTests(unittest.TestCase):
["imported", "duplicate", "linked"],
)
# N1: a collection larger than SQLite's bound-variable limit is replaced.
def test_more_imports_than_sqlite_variables(self):
conn = self.db.model.get_connection()
if not hasattr(conn, "setlimit"):
self.skipTest("sqlite3.Connection.setlimit needs Python 3.11")
rows = [
(f"bulk lora {n}", CATEGORY, generate_prompt_hash(f"bulk lora {n}"))
for n in range(1200)
]
with conn:
conn.executemany(
"INSERT INTO prompts(text, category, hash, created_at, updated_at) "
"VALUES(?, ?, ?, '2026-01-01', '2026-01-01')",
rows,
)
conn.setlimit(sqlite3.SQLITE_LIMIT_VARIABLE_NUMBER, 999)
state = self.db.category_import_state(CATEGORY)
self.assertEqual((len(state["ids"]), state["paired"]), (1202, 0))
result = self.db.replace_category_prompts(
CATEGORY, [_record("only one")], expected_ids=state["ids"]
)
self.assertEqual((result["removed"], result["imported"]), (1202, 1))
class LoraImportAPITestCase(LoraAPITestCase):
"""The routes with LoraManager enabled over a disposable LoRA folder."""
@@ -482,6 +507,37 @@ class LoraImportRouteTests(LoraImportAPITestCase):
self.assertEqual(plan["backup"]["policy"], "on")
self.assertEqual(order, [("backup", True), ("replace", None)])
async def test_a_disconnected_legacy_scan_stops_and_holds_until_it_stopped(self):
# I2: the admin page closes its modal while the downloads run.
for number in range(4):
_write_metadata(
self.loras,
f"lora{number}",
images=[f"https://civitai.com/{number}.png"],
)
before = _snapshot(self.api.db)
self.gate = threading.Event()
resp = await self.client.request("POST", SCAN)
await resp.content.readline() # the stream is open
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, self.downloading.wait, 10))
resp.close()
await asyncio.sleep(0.05) # the server sees the closed connection
# Its records hold stays while the in-flight download runs.
self.assertIsNotNone(self.api._jobs.active("lora_import"))
self.gate.set()
for _ in range(500):
if self.api._jobs.active("lora_import") is None:
break
await asyncio.sleep(0.01)
self.assertIsNone(self.api._jobs.active("lora_import"))
# Staging stopped at the next LoRA: no download runs after the hold.
served = len(self.requests)
self.assertLess(served, 4)
await asyncio.sleep(0.1)
self.assertEqual(len(self.requests), served)
self.assertEqual(_snapshot(self.api.db), before)
async def test_legacy_scan_events_unchanged_on_success(self):
self._two_loras()
(self.loras / "broken.metadata.json").write_text("{{{")
@@ -587,6 +643,27 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
self.assertEqual(_snapshot(self.api.db), before)
return job
async def test_start_refused_when_the_integration_changed_since_the_review(self):
# M5: disabled, or moved to another install, after the review.
self._two_loras()
before = _snapshot(self.api.db)
resp, plan = await self._post(PREVIEW)
IntegrationConfig.LORA_MANAGER_ENABLED = False
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "plan-changed"))
self.assertIn("integration changed", body["error"])
IntegrationConfig.LORA_MANAGER_ENABLED = True
resp, plan = await self._post(PREVIEW)
other = _make_lora_manager(self.custom_nodes / "other")
target = self.custom_nodes / "Other-Lora-Manager"
other.rename(target)
IntegrationConfig.LORA_MANAGER_PATH = str(target)
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "plan-changed"))
self.assertEqual(self.api._jobs.list(), [])
self.assertEqual(_snapshot(self.api.db), before)
async def test_imports_changed_after_the_review_fail_the_job(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
@@ -619,20 +696,57 @@ class LoraImportRefusalTests(LoraImportAPITestCase):
with patch.object(lora_import, "build_record", side_effect=ValueError("x")):
job = await self._failed_job()
self.assertEqual(job.scope["failure"], "nothing-to-import")
self.assertEqual(job.counters["skipped"], 0)
# M4: the job keeps what staging skipped, by name, though it failed.
self.assertEqual(job.counters["skipped"], 2)
self.assertEqual(
sorted(row["detail"] for row in job.outcomes["skipped"]), ["neon", "plain"]
)
async def _change_during_download(self, change):
"""Start the reviewed import, run ``change()`` while a download is held."""
self._two_loras()
before = _snapshot(self.api.db)
self.gate = threading.Event()
_plan, job = await self._apply_reviewed()
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, self.downloading.wait, 10))
change()
self.gate.set()
await job.task
self.assertEqual(job.state, "failed", job.log)
self.assertEqual(job.scope["failure"], "plan-changed")
self.assertIn("Review the import again", job.status_line)
self.assertEqual(_snapshot(self.api.db), before)
return job
async def test_files_changed_while_staging_fail_the_job(self):
# I1: LoraManager rewrites a metadata file while the downloads run.
meta = self.loras / "plain.safetensors.metadata.json"
await self._change_during_download(
lambda: meta.write_text(meta.read_text() + " ")
)
async def test_a_file_removed_while_staging_fails_the_job(self):
meta = self.loras / "plain.safetensors.metadata.json"
await self._change_during_download(meta.unlink)
async def test_a_file_added_while_staging_fails_the_job(self):
await self._change_during_download(lambda: _write_metadata(self.loras, "late"))
async def test_legacy_scan_refuses_files_changed_while_it_stages(self):
self._two_loras()
real = lora_import.stage
def changed(*args, **kwargs):
staged = real(*args, **kwargs)
staged.fingerprint = "sha256:other"
return staged
with patch.object(lora_import, "stage", side_effect=changed):
job = await self._failed_job()
self.assertEqual(job.scope["failure"], "plan-changed")
before = _snapshot(self.api.db)
self.gate = threading.Event()
meta = self.loras / "plain.safetensors.metadata.json"
resp = await self.client.request("POST", SCAN)
loop = asyncio.get_running_loop()
self.assertTrue(await loop.run_in_executor(None, self.downloading.wait, 10))
meta.write_text(meta.read_text() + " ")
self.gate.set()
events = self._sse(await resp.text())
self.assertEqual(events[-1]["type"], "error")
self.assertIn("changed during the scan", events[-1]["message"])
self.assertEqual(_snapshot(self.api.db), before)
async def test_a_failed_safety_backup_changes_nothing(self):
self._two_loras()
tests/test_lora_utils.py
+82
View File
@@ -524,6 +524,59 @@ class TestDownloadOne(unittest.TestCase):
self.assertFalse(self.target.exists())
class TestDownloadIsAtomic(unittest.TestCase):
"""I2 (S10 review): a cached image is never seen half written."""
def setUp(self):
self.tmpdir = Path(tempfile.mkdtemp())
self.addCleanup(shutil.rmtree, self.tmpdir, True)
self.target = self.tmpdir / "out.jpg"
def test_the_image_appears_only_once_complete(self):
seen = []
real_replace = os.replace
def replace(source, destination):
# At the moment the name appears, the temp copy is whole and closed.
seen.append((Path(source).parent, self.target.exists()))
with open(source, "rb") as handle:
self.assertEqual(handle.read(2), b"\xff\xd8")
return real_replace(source, destination)
opener = _FakeOpener(_FakeResponse(_png_bytes()))
with (
patch.object(lora_utils, "_open_url", opener),
patch.object(lora_utils.os, "replace", side_effect=replace),
):
result = lora_utils._download_one(
"https://civitai.com/api/download/1", self.target, ""
)
self.assertEqual(seen, [(self.tmpdir, False)])
self.assertEqual(os.path.realpath(result), os.path.realpath(self.target))
self.assertEqual(sorted(p.name for p in self.tmpdir.iterdir()), ["out.jpg"])
def test_a_failed_write_leaves_neither_the_image_nor_a_temp_file(self):
opener = _FakeOpener(_FakeResponse(_png_bytes()))
with (
patch.object(lora_utils, "_open_url", opener),
patch.object(lora_utils.os, "replace", side_effect=OSError("disk full")),
):
result = lora_utils._download_one(
"https://civitai.com/api/download/1", self.target, ""
)
self.assertIsNone(result)
self.assertEqual(list(self.tmpdir.iterdir()), [])
def test_an_unreadable_image_leaves_no_temp_file(self):
opener = _FakeOpener(_FakeResponse(b"not an image"))
with patch.object(lora_utils, "_open_url", opener):
result = lora_utils._download_one(
"https://civitai.com/api/download/1", self.target, ""
)
self.assertIsNone(result)
self.assertEqual(list(self.tmpdir.iterdir()), [])
class _FakeHTTPSHandler(urllib.request.HTTPSHandler):
"""Serves canned responses through a real urllib opener, recording requests.
@@ -861,6 +914,35 @@ class TestDetectLoraManager(unittest.TestCase):
):
self.assertEqual(lora_utils.detect_lora_manager(), str(main.resolve()))
def test_an_entry_that_cannot_be_checked_is_skipped(self):
# S10 review N5: a permission error on one entry, not a missing folder.
root = _make_comfy_root(self.tmp)
blocked = root / "custom_nodes" / "a-lora-manager"
blocked.mkdir()
lm = _make_lora_manager(root / "custom_nodes", "b-Lora-Manager")
real_is_dir = Path.is_dir
def is_dir(path):
if path == blocked:
raise PermissionError(13, "Permission denied", str(path))
return real_is_dir(path)
with (
patch.object(lora_utils, "find_comfyui_root", return_value=root),
patch.object(lora_utils, "custom_nodes_directories", return_value=[]),
patch.object(Path, "is_dir", is_dir),
):
self.assertEqual(lora_utils.detect_lora_manager(), str(lm.resolve()))
def test_a_folder_that_cannot_be_listed_is_skipped(self):
root = _make_comfy_root(self.tmp)
with (
patch.object(lora_utils, "find_comfyui_root", return_value=root),
patch.object(lora_utils, "custom_nodes_directories", return_value=[]),
patch.object(Path, "iterdir", side_effect=PermissionError(13, "denied")),
):
self.assertIsNone(lora_utils.detect_lora_manager())
def test_an_unreadable_extra_folder_is_skipped(self):
root = _make_comfy_root(self.tmp)
missing = self.tmp / "gone"