From d247e4c2e74ee532d7f7e10a5a1b9e1fe2d79936 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sun, 25 Oct 2026 08:36:46 -0700 Subject: [PATCH] fix(lora): recheck the files before the replace and stop an abandoned scan 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. --- database/category_import.py | 64 ++++++++++++----- py/api/lora_import_routes.py | 39 ++++++++-- py/api/lora_integration.py | 128 ++++++++++++++++++++++++++++----- py/lora_import.py | 9 +-- py/lora_utils.py | 34 +++++++-- tests/test_api_lora.py | 15 ++++ tests/test_lora_import.py | 136 ++++++++++++++++++++++++++++++++--- tests/test_lora_utils.py | 82 +++++++++++++++++++++ 8 files changed, 444 insertions(+), 63 deletions(-) diff --git a/database/category_import.py b/database/category_import.py index 68ebd58..89e876e 100644 --- a/database/category_import.py +++ b/database/category_import.py @@ -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.""" diff --git a/py/api/lora_import_routes.py b/py/api/lora_import_routes.py index 029b6e6..1531a8c 100644 --- a/py/api/lora_import_routes.py +++ b/py/api/lora_import_routes.py @@ -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')}") diff --git a/py/api/lora_integration.py b/py/api/lora_integration.py index 430d6ce..0d562a0 100644 --- a/py/api/lora_integration.py +++ b/py/api/lora_integration.py @@ -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( { diff --git a/py/lora_import.py b/py/lora_import.py index c1fdca4..c17f631 100644 --- a/py/lora_import.py +++ b/py/lora_import.py @@ -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() diff --git a/py/lora_utils.py b/py/lora_utils.py index 66db922..7a23f68 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -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]]]: diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py index f7809c9..43ebc14 100644 --- a/tests/test_api_lora.py +++ b/tests/test_api_lora.py @@ -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")): diff --git a/tests/test_lora_import.py b/tests/test_lora_import.py index a589d26..3433b49 100644 --- a/tests/test_lora_import.py +++ b/tests/test_lora_import.py @@ -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() diff --git a/tests/test_lora_utils.py b/tests/test_lora_utils.py index aa8eae1..24ebe0e 100644 --- a/tests/test_lora_utils.py +++ b/tests/test_lora_utils.py @@ -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"