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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1403 lines
55 KiB
Python
Raw Blame History

"""The staged LoRA reference import (#221, #249 System S10).
A failed import must leave the previous ``lora-manager`` collection exactly as
it was (prompts, tags and image links); a preview writes nothing and requests
nothing; an apply does what was reviewed or nothing. Discovery, metadata and
CivitAI downloads are disposable fixtures: no network, no LoraManager install.
"""
import asyncio
import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
from unittest.mock import patch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
try:
from tests.test_api_lora import (
SECRET_KEY,
LoraAPITestCase,
_FakeResponse,
_make_lora_manager,
_png_bytes,
_write_metadata,
)
except ImportError: # discovered with tests/ as the top-level directory
from test_api_lora import (
SECRET_KEY,
LoraAPITestCase,
_FakeResponse,
_make_lora_manager,
_png_bytes,
_write_metadata,
)
from database.category_import import EmptyReplacement
from database.maintenance import PlanChanged
from database.operations import PairDependencyError, PromptDatabase
from py import lora_import, lora_utils
from py.api.jobs import RESTORE_JOB_KIND
from py.api.lora_import_routes import LORA_PREVIEW_LEASE
from py.config import IntegrationConfig, PromptManagerConfig
from utils.hashing import generate_prompt_hash
CATEGORY = "lora-manager"
PREVIEW = "/prompt_manager/lora/import/preview"
APPLY = "/prompt_manager/lora/import/apply"
SCAN = "/prompt_manager/lora/scan"
def _snapshot(db):
"""Every row of the tables an import touches, in a fixed order."""
tables = {
"prompts": "SELECT id, text, category, hash FROM prompts ORDER BY id",
"prompt_tags": "SELECT prompt_id, tag_id FROM prompt_tags ORDER BY 1, 2",
"tags": "SELECT id, name FROM tags ORDER BY id",
"generated_images": (
"SELECT id, prompt_id, image_path, file_path FROM generated_images "
"ORDER BY id"
),
}
with db.model.get_connection() as conn:
return {
name: [tuple(row) for row in conn.execute(sql).fetchall()]
for name, sql in tables.items()
}
def _seed_previous_imports(db):
"""Two earlier imports with tags and image links, and one other prompt."""
ids = []
for name in ("old-one", "old-two"):
prompt_id = db.save_prompt(
text=f"{name} prompt",
category=CATEGORY,
tags=[CATEGORY, f"lora:{name}"],
prompt_hash=generate_prompt_hash(f"{name} prompt"),
)
db.link_image_to_prompt(str(prompt_id), f"/cache/{name}/a.jpg")
db.link_image_to_prompt(str(prompt_id), f"/cache/{name}/b.jpg")
ids.append(prompt_id)
db.save_prompt(
text="mine",
category="portraits",
tags=["keep"],
prompt_hash=generate_prompt_hash("mine"),
)
return ids
def _record(text, images=(), name=None):
return {
"name": name or text,
"text": text,
"hash": generate_prompt_hash(text),
"tags": [CATEGORY, f"lora:{name or text}"],
"images": list(images),
}
class CategoryReplacementTests(unittest.TestCase):
"""``replace_category_prompts``: one transaction, all or nothing."""
def setUp(self):
handle = tempfile.NamedTemporaryFile(delete=False, suffix=".db")
handle.close()
self.path = handle.name
self.db = PromptDatabase(self.path)
self.ids = _seed_previous_imports(self.db)
def tearDown(self):
self.db.close_all()
for suffix in ("", "-wal", "-shm"):
if os.path.exists(self.path + suffix):
os.unlink(self.path + suffix)
def test_transaction_failure_rolls_back_everything(self):
before = _snapshot(self.db)
records = [_record("new one", ["/c/n1.jpg"]), _record("new two", ["/c/n2.jpg"])]
real = self.db._insert_image_link
calls = []
def failing(conn, prompt_id, image, metadata=None):
calls.append(image)
if len(calls) == 2: # after the old rows went and one record landed
raise sqlite3.OperationalError("disk I/O error")
return real(conn, prompt_id, image, metadata)
with patch.object(self.db, "_insert_image_link", side_effect=failing):
with self.assertRaises(sqlite3.OperationalError):
self.db.replace_category_prompts(CATEGORY, records)
self.assertEqual(len(calls), 2)
self.assertEqual(_snapshot(self.db), before)
def test_pair_dependency_blocks_without_partial_delete(self):
# A saved pair reuses an imported prompt as its positive component.
self.db.save_prompt_pair("old-two prompt", "blurry")
before = _snapshot(self.db)
self.assertEqual(
self.db.category_import_state(CATEGORY), {"ids": self.ids, "paired": 1}
)
with self.assertRaises(PairDependencyError):
self.db.replace_category_prompts(CATEGORY, [_record("new one")])
self.assertEqual(_snapshot(self.db), before)
def test_explicit_empty_replacement_requires_review_flag(self):
before = _snapshot(self.db)
with self.assertRaises(EmptyReplacement):
self.db.replace_category_prompts(CATEGORY, [])
self.assertEqual(_snapshot(self.db), before)
result = self.db.replace_category_prompts(CATEGORY, [], allow_empty=True)
self.assertEqual(result["removed"], 2)
self.assertEqual(self.db.category_import_state(CATEGORY)["ids"], [])
self.assertEqual(len(self.db.search_prompts(category="portraits")), 1)
def test_changed_imports_since_review_are_plan_changed(self):
before = _snapshot(self.db)
with self.assertRaises(PlanChanged):
self.db.replace_category_prompts(
CATEGORY, [_record("new one")], expected_ids=self.ids[:1]
)
self.assertEqual(_snapshot(self.db), before)
def test_replacement_links_repeats_and_existing_prompts(self):
records = [
_record("fresh", ["/c/f.jpg"]),
_record("Fresh ", ["/c/f2.jpg"], name="fresh-again"), # same hash
_record("mine", ["/c/m.jpg"], name="mine-lora"), # exists elsewhere
]
result = self.db.replace_category_prompts(
CATEGORY, records, expected_ids=self.ids
)
self.assertEqual(
(result["removed"], result["imported"], result["linked"]), (2, 1, 1)
)
self.assertEqual(result["duplicates"], 1)
self.assertEqual(result["images"], 3)
imported = self.db.search_prompts(category=CATEGORY)
self.assertEqual([p["text"] for p in imported], ["fresh"])
self.assertEqual(len(self.db.get_prompt_images(imported[0]["id"])), 2)
mine = self.db.search_prompts(category="portraits")[0]
self.assertEqual(len(self.db.get_prompt_images(mine["id"])), 1)
self.assertEqual(
[row["outcome"] for row in result["rows"]],
["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."""
async def get_application(self):
app = await super().get_application()
self.lm = _make_lora_manager(self.custom_nodes)
self.loras = self.tmp_root / "loras"
self.loras.mkdir()
self.cache_dir = self.tmp_root / "img_cache"
IntegrationConfig.LORA_MANAGER_ENABLED = True
IntegrationConfig.LORA_MANAGER_PATH = str(self.lm)
IntegrationConfig.CIVITAI_API_KEY = SECRET_KEY
self.requests = []
self.gate = None # a threading.Event a download waits on
self.downloading = threading.Event()
self.failing = set() # URLs whose download fails
def fake_open(req, timeout=None):
self.requests.append(req)
self.downloading.set()
if req.full_url in self.failing:
raise OSError("connection reset")
if self.gate is not None:
self.gate.wait(10)
return _FakeResponse(_png_bytes())
self._patches = [
patch.object(
lora_utils, "default_lora_roots", return_value=[str(self.loras)]
),
patch.object(
lora_utils, "get_lora_image_cache_dir", return_value=self.cache_dir
),
patch.object(lora_utils, "_open_url", fake_open),
patch.object(PromptManagerConfig, "SAFETY_BACKUP", False),
]
for item in self._patches:
item.start()
self.old_ids = _seed_previous_imports(self.api.db)
return app
async def tearDownAsync(self):
if self.gate is not None:
self.gate.set()
for job in self.api._jobs.list():
if job.task is not None and not job.task.done():
await job.task
for item in reversed(self._patches):
item.stop()
await super().tearDownAsync()
def _two_loras(self):
"""``neon`` has only a CivitAI image (downloaded); ``plain`` has a
local preview (linked where it is, nothing downloaded)."""
_write_metadata(
self.loras,
"neon",
words=["glow"],
images=["https://civitai.com/a.png"],
prompt="neon street",
)
_write_metadata(self.loras, "plain", images=["https://civitai.com/p.png"])
(self.loras / "plain.png").write_bytes(_png_bytes())
async def _post(self, url, body=None):
resp = await self._post_json(url, {} if body is None else body)
return resp, json.loads(await resp.text())
async def _apply_reviewed(self):
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual(resp.status, 202, body)
job = self.api._jobs.get(body["job"]["id"])
return plan, job
@staticmethod
def _sse(text):
return [
json.loads(line[len("data: ") :])
for line in text.splitlines()
if line.startswith("data: ")
]
class LoraImportRouteTests(LoraImportAPITestCase):
async def test_preview_writes_nothing_and_makes_no_requests(self):
self._two_loras()
(self.loras / "broken.metadata.json").write_text("{{{")
db_path = self.api.db.model.db_path
before = {s: _read(db_path + s) for s in ("", "-wal")}
shm = _size(db_path + "-shm")
rows = _snapshot(self.api.db)
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertEqual(
{
k: plan[k]
for k in (
"lora_files",
"with_metadata",
"skipped",
"with_local_media",
"current_imports",
)
},
{
"lora_files": 3,
"with_metadata": 2,
"skipped": 1,
"with_local_media": 1,
"current_imports": 2,
},
)
self.assertEqual(
plan["layout"],
{
"source": "environment",
"settings_file": str(self.lm_settings / "settings.json"),
"error": None,
"mode": "alongside",
"library": "default",
"example_images": False,
},
)
self.assertEqual(plan["blocked_by_pairs"], 0)
self.assertEqual(plan["planned_downloads"], 1)
self.assertEqual(plan["network"], {"planned_downloads": 1, "api_key": True})
self.assertTrue(plan["token"])
self.assertEqual(plan["scope"]["count"], 2)
self.assertFalse(plan["media"]["deletes_files"])
self.assertEqual(self.requests, [])
self.assertFalse(self.cache_dir.exists())
self.assertEqual({s: _read(db_path + s) for s in ("", "-wal")}, before)
self.assertEqual(_size(db_path + "-shm"), shm)
self.assertEqual(_snapshot(self.api.db), rows)
self.assertNotIn(SECRET_KEY, json.dumps(plan))
# The settings file is the only absolute path the review shows.
shown = dict(plan, layout=dict(plan["layout"], settings_file=None))
self.assertNotIn(str(self.tmp_root), json.dumps(shown))
async def test_unusable_settings_review_offers_no_token(self):
self._two_loras()
(self.lm_settings / "settings.json").write_text("[1, 2]")
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertIsNone(plan["token"])
self.assertEqual(
plan["layout"]["error"], "the settings file is not a JSON object"
)
self.assertEqual((plan["lora_files"], plan["with_metadata"]), (0, 0))
self.assertEqual(self.requests, [])
async def test_centralized_review_counts_what_it_skips(self):
sidecars = self.tmp_root / "sidecars"
examples = self.tmp_root / "examples"
(self.lm_settings / "settings.json").write_text(
json.dumps(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(sidecars),
"example_images_path": str(examples),
"libraries": {"comfyui": {}},
"active_library": "comfyui",
}
)
)
component = sidecars / "loras-0a1b2c3d"
component.mkdir(parents=True)
model = self.loras / "neon.safetensors"
model.write_bytes(b"m")
meta = {"model_name": "neon", "file_path": str(model), "sha256": "AB12"}
(component / "neon.metadata.json").write_text(json.dumps(meta))
(examples / "ab12").mkdir(parents=True)
(examples / "ab12" / "image_1.png").write_bytes(_png_bytes())
gone = dict(meta, file_path=str(self.loras / "gone.safetensors"))
(component / "gone.metadata.json").write_text(json.dumps(gone))
(component / "broken.metadata.json").write_text("{{{")
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertEqual(
(
plan["lora_files"],
plan["with_metadata"],
plan["skipped"],
plan["with_local_media"],
plan["planned_downloads"],
),
(3, 1, 2, 1, 0),
)
self.assertEqual(
(plan["layout"]["mode"], plan["layout"]["library"]),
("centralized", "comfyui"),
)
self.assertTrue(plan["layout"]["example_images"])
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(self.requests, [])
self.assertEqual(
[row["detail"] for row in job.outcomes["skipped"]],
["broken.metadata.json"],
)
imported = self.api.db.search_prompts(category=CATEGORY)
images = self.api.db.get_prompt_images(imported[0]["id"])
self.assertEqual(
[os.path.basename(i["image_path"]) for i in images], ["image_1.png"]
)
async def test_an_unlistable_folder_says_so_in_the_review(self):
self._two_loras()
error = PermissionError(13, "Permission denied", str(self.loras / "lost+found"))
with patch.object(lora_import, "discover_sidecars", side_effect=error):
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (500, "folder-unreadable"))
self.assertEqual(
body["error"],
"A LoRA folder could not be listed (Permission denied: lost+found). "
"Nothing was changed.",
)
self.assertNotIn(str(self.tmp_root), json.dumps(body))
async def test_preview_requires_the_integration(self):
IntegrationConfig.LORA_MANAGER_ENABLED = False
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (400, "lora-disabled"))
async def test_blocked_preview_offers_no_token(self):
self._two_loras()
self.api.db.save_prompt_pair("old-one prompt", "blurry")
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200)
self.assertEqual(plan["blocked_by_pairs"], 1)
self.assertIsNone(plan["token"])
self.assertEqual(plan["protected"]["count"], 1)
async def test_changed_source_after_preview_is_409(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
rows = _snapshot(self.api.db)
meta = self.loras / "plain.metadata.json"
meta.write_text(meta.read_text() + " ")
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(self.requests, [])
self.assertEqual(_snapshot(self.api.db), rows)
async def test_apply_imports_through_the_job(self):
self._two_loras()
plan, job = await self._apply_reviewed()
self.assertEqual(
(job.kind, job.title, job.admission),
("lora_import", plan_title(), "records"),
)
await job.task
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(
{k: job.counters[k] for k in lora_import.GROUPS},
{
"imported": 2,
"linked": 0,
"skipped": 0,
"downloaded": 1,
"download_errors": 0,
"errors": 0,
},
)
self.assertFalse(job.cancellable)
self.assertEqual(
self.requests[0].get_header("Authorization"), f"Bearer {SECRET_KEY}"
)
imported = self.api.db.search_prompts(category=CATEGORY)
self.assertEqual(sorted(p["text"] for p in imported), ["neon street", "plain"])
neon = next(p for p in imported if p["text"] == "neon street")
self.assertEqual(len(self.api.db.get_prompt_images(neon["id"])), 1)
# plain's local preview is linked where it is; its CivitAI image is not
# downloaded (#355).
plain = next(p for p in imported if p["text"] == "plain")
self.assertEqual(
[i["image_path"] for i in self.api.db.get_prompt_images(plain["id"])],
[os.path.realpath(self.loras / "plain.png")],
)
self.assertEqual(
[r.full_url for r in self.requests], ["https://civitai.com/a.png"]
)
self.assertEqual(len(self.api.db.search_prompts(category="portraits")), 1)
# The token is single-use.
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "plan-used"))
async def test_a_finished_import_clears_the_layout_cache(self):
# #362: the import may link files under a folder LoraManager created
# after the image routes cached its layout.
self._two_loras()
with patch("py.lora_utils.clear_layout_cache") as clear:
_plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(job.state, "finished", job.log)
clear.assert_called_once_with()
async def test_discovery_failure_keeps_previous_imports(self):
self._two_loras()
before = _snapshot(self.api.db)
# Legacy scan: discovery raises once the stream is open.
with patch.object(
lora_import, "discover_sidecars", side_effect=OSError("share gone")
):
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
self.assertEqual(events[-1]["type"], "error")
self.assertEqual(_snapshot(self.api.db), before)
# Job: staging fails after the review.
with patch.object(lora_import, "stage", side_effect=OSError("share gone")):
_plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(job.state, "failed")
self.assertIn("previous LoRA imports were kept", job.status_line)
self.assertEqual(job.scope["failure"], "import-failed")
self.assertEqual(_snapshot(self.api.db), before)
async def test_a_failed_replacement_keeps_previous_imports_and_says_so(self):
self._two_loras()
before = _snapshot(self.api.db)
with patch.object(
self.api.db,
"replace_category_prompts",
side_effect=sqlite3.OperationalError("database is locked"),
):
_plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(job.state, "failed")
self.assertEqual(job.counters["errors"], 1)
self.assertEqual(job.counters["downloaded"], 1)
self.assertIn("previous LoRA imports were kept", job.status_line)
self.assertIn("stay in the image cache", job.status_line)
self.assertEqual(_snapshot(self.api.db), before)
async def test_cancel_during_downloads_writes_nothing(self):
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))
self.assertTrue(job.summary()["cancellable"])
self.assertTrue(self.api._jobs.cancel(job.id))
self.gate.set()
await job.task
self.assertEqual(job.state, "cancelled")
self.assertEqual(_snapshot(self.api.db), before)
self.assertTrue(any("stay in the image cache" in line for line in job.log))
async def test_empty_reviewed_import_removes_the_collection(self):
# No metadata at all: the review says so, and its apply is the flag.
resp, plan = await self._post(PREVIEW)
self.assertEqual((plan["with_metadata"], plan["current_imports"]), (0, 2))
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(self.api.db.search_prompts(category=CATEGORY), [])
# The legacy scan has no review: an empty replacement is refused.
_seed_previous_imports(self.api.db)
before = _snapshot(self.api.db)
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
self.assertEqual(events[-1]["type"], "error")
self.assertEqual(_snapshot(self.api.db), before)
async def test_apply_refused_while_a_records_job_holds(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
async with self.api._jobs.hold("scan", "Metadata scan", admission="records"):
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "busy"))
# The token was not used: it still applies.
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual(resp.status, 202)
await self.api._jobs.get(body["job"]["id"]).task
async def test_safety_backup_policy_runs_before_the_replace(self):
self._two_loras()
order = []
real = self.api.db.replace_category_prompts
async def step(job, db, backup_on):
order.append(("backup", backup_on))
return True, "prompts.db.backup_x"
def replace(*args, **kwargs):
order.append(("replace", None))
return real(*args, **kwargs)
with (
patch.object(PromptManagerConfig, "SAFETY_BACKUP", True),
patch.object(self.api, "_safety_backup_step", side_effect=step),
patch.object(self.api.db, "replace_category_prompts", side_effect=replace),
):
plan, job = await self._apply_reviewed()
await job.task
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"))
# R2 (re-review): the user left and nothing failed, so Activity says so.
job = next(j for j in self.api._jobs.list() if j.kind == "lora_import")
self.assertEqual(job.state, "cancelled")
# 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_an_import_sweeps_stale_temp_downloads_from_the_cache(self):
# N7 (re-review): a download killed mid-write leaves a .pm-*.part file;
# the next import removes it once it is an hour old.
self._two_loras()
stale = self.cache_dir / "old_lora" / ".pm-dead.part"
stale.parent.mkdir(parents=True)
stale.write_bytes(b"x")
stamp = time.time() - 2 * 3600
os.utime(stale, (stamp, stamp))
resp = await self.client.request("POST", SCAN)
await resp.text()
self.assertFalse(stale.exists())
async def test_legacy_scan_events_unchanged_on_success(self):
self._two_loras()
(self.loras / "broken.metadata.json").write_text("{{{")
(self.loras / "broken.safetensors").write_bytes(b"model")
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
self.assertEqual(resp.status, 200)
progress_keys = {
"type",
"status",
"progress",
"total",
"processed",
"imported",
"skipped",
}
for event in events[:-1]:
self.assertEqual(event["type"], "progress")
self.assertLessEqual(set(event), progress_keys)
self.assertIsInstance(event["status"], str)
self.assertIsInstance(event["progress"], int)
per_record = [e for e in events if "processed" in e]
self.assertEqual([e["processed"] for e in per_record], [1, 2, 3])
self.assertEqual(
events[-1],
{
"type": "complete",
"progress": 100,
"total": 3,
"imported": 2,
"skipped": 1,
},
)
imported = self.api.db.search_prompts(category=CATEGORY)
self.assertEqual(sorted(p["text"] for p in imported), ["neon street", "plain"])
class LoraImportDownloadOutcomeTests(LoraImportAPITestCase):
"""#337: failed CivitAI downloads show in the job's status line."""
URLS = ["https://civitai.com/a.png", "https://civitai.com/b.png"]
def _lora_with_images(self):
_write_metadata(
self.loras,
"neon",
words=["glow"],
images=list(self.URLS),
prompt="neon street",
)
_write_metadata(self.loras, "plain")
async def _run(self):
_plan, job = await self._apply_reviewed()
await job.task
return job
async def test_status_line_unchanged_when_every_download_succeeds(self):
self._lora_with_images()
job = await self._run()
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(job.status_line, "Imported 2 LoRAs")
self.assertEqual(job.outcomes["download_errors"], [])
async def test_status_line_counts_failed_downloads(self):
self._lora_with_images()
self.failing = {self.URLS[1]}
job = await self._run()
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(
(job.counters["downloaded"], job.counters["download_errors"]), (1, 1)
)
self.assertEqual(
job.status_line,
"Imported 2 LoRAs · 1 example image could not be downloaded",
)
self.assertNotIn("Warning", job.status_line)
self.assertEqual(job.outcomes["download_errors"], [])
async def test_every_download_failing_is_a_warning(self):
self._lora_with_images()
self.failing = set(self.URLS)
job = await self._run()
# The import itself succeeded, so the job still finishes.
self.assertEqual(job.state, "finished", job.log)
self.assertEqual(
(job.counters["downloaded"], job.counters["download_errors"]), (0, 2)
)
warning = "Warning: none of the 2 example images could be downloaded"
self.assertEqual(job.status_line, f"{warning} · Imported 2 LoRAs")
self.assertEqual(
job.outcomes["download_errors"],
[
{
"path": "",
"root": "",
"root_index": None,
"detail": f"{warning}; the log has the details",
}
],
)
self.assertTrue(any(warning in line for line in job.log), job.log)
self.assertEqual(job.summary()["status_line"], job.status_line)
self.assertNotIn(SECRET_KEY, json.dumps(job.report()))
async def test_one_planned_download_failing_reads_in_the_singular(self):
_write_metadata(self.loras, "neon", images=[self.URLS[0]], prompt="neon street")
self.failing = {self.URLS[0]}
job = await self._run()
self.assertEqual(
job.status_line,
"Warning: the example image could not be downloaded · Imported 1 LoRA",
)
async def test_legacy_scan_events_unchanged_when_every_download_fails(self):
self._lora_with_images()
self.failing = set(self.URLS)
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
self.assertEqual(resp.status, 200)
self.assertEqual(
events[-1],
{
"type": "complete",
"progress": 100,
"total": 2,
"imported": 2,
"skipped": 0,
},
)
for event in events[:-1]:
self.assertEqual(event["type"], "progress")
self.assertNotIn("Warning", event.get("status", ""))
class LoraImportRefusalTests(LoraImportAPITestCase):
"""Every refusal leaves the previous imports as they were."""
async def test_preview_refusals(self):
IntegrationConfig.LORA_MANAGER_PATH = ""
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (400, "lora-path"))
IntegrationConfig.LORA_MANAGER_PATH = str(_make_lora_manager(self.tmp_root))
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (400, "lora-path"))
self.assertNotIn(str(self.tmp_root), json.dumps(body))
IntegrationConfig.LORA_MANAGER_PATH = str(self.lm)
resp, body = await self._post(PREVIEW, {"extra": 1})
self.assertEqual((resp.status, body["code"]), (400, "invalid-request"))
with patch.object(lora_import, "preview", side_effect=OSError("gone")):
resp, body = await self._post(PREVIEW)
self.assertEqual(resp.status, 500)
self.assertFalse(body["success"])
generations = iter([(0, "a"), (1, "a")])
with patch.object(
self.api, "_database_identity", side_effect=lambda db: next(generations)
):
resp, body = await self._post(PREVIEW)
self.assertEqual((resp.status, body["code"]), (409, "storage-changed"))
async def test_apply_refusals(self):
self._two_loras()
before = _snapshot(self.api.db)
resp, body = await self._post(APPLY, {"token": ""})
self.assertEqual((resp.status, body["code"]), (400, "invalid-request"))
resp, body = await self._post(APPLY, {"token": "unknown"})
self.assertEqual((resp.status, body["code"]), (410, "plan-expired"))
other = self.api._maintenance_plans.issue(
"optimize", "f", {"backup": "off"}, self.api._database_identity(self.api.db)
)
resp, body = await self._post(APPLY, {"token": other})
self.assertEqual((resp.status, body["code"]), (410, "plan-expired"))
resp, plan = await self._post(PREVIEW)
with patch.object(lora_import, "discover", side_effect=OSError("gone")):
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (500, "apply-failed"))
short = ("database folder", "the safety backup", 10, 1)
with patch("py.api.lora_import_routes.disk_plan", return_value=({}, short, 0)):
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (507, "insufficient-space"))
self.api._db_generation += 1
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual((resp.status, body["code"]), (409, "storage-changed"))
self.assertEqual(self.api._jobs.list(), [])
self.assertEqual(_snapshot(self.api.db), before)
async def _failed_job(self, **patches):
before = _snapshot(self.api.db)
_plan, job = await self._apply_reviewed()
await job.task
self.assertEqual(job.state, "failed", job.log)
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)
self.api.db.save_prompt(
text="late", category=CATEGORY, prompt_hash=generate_prompt_hash("late")
)
before = _snapshot(self.api.db)
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual((job.state, job.scope["failure"]), ("failed", "plan-changed"))
self.assertEqual(_snapshot(self.api.db), before)
async def test_a_refused_import_keeps_the_layout_cache(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
self.api.db.save_prompt(
text="late", category=CATEGORY, prompt_hash=generate_prompt_hash("late")
)
with patch("py.lora_utils.clear_layout_cache") as clear:
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual(job.scope["failure"], "plan-changed")
clear.assert_not_called()
async def test_start_refused_when_a_skipped_sidecar_changed(self):
# #362 CUJ 4: the review counted the orphan as skipped.
self._two_loras()
orphan = self.loras / "orphan.metadata.json"
orphan.write_text(json.dumps({"model_name": "orphan"}))
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertEqual((plan["lora_files"], plan["skipped"]), (3, 1))
before = _snapshot(self.api.db)
orphan.write_text(orphan.read_text() + " ")
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual(resp.status, 409, body)
self.assertEqual(body["code"], "plan-changed")
self.assertEqual(self.api._jobs.list(), [])
self.assertEqual(_snapshot(self.api.db), before)
async def test_a_pair_saved_after_the_review_blocks_the_job(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
self.api.db.save_prompt_pair("old-one prompt", "blurry")
before = _snapshot(self.api.db)
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
await job.task
self.assertEqual(
(job.state, job.scope["failure"]), ("failed", "blocked-by-pairs")
)
self.assertIn("Saved pairs", job.status_line)
self.assertEqual(_snapshot(self.api.db), before)
async def test_nothing_prepared_keeps_the_collection(self):
self._two_loras()
with patch.object(lora_import, "build_record", side_effect=ValueError("x")):
job = await self._failed_job()
self.assertEqual(job.scope["failure"], "nothing-to-import")
# 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.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.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()
before = _snapshot(self.api.db)
self.gate = threading.Event()
meta = self.loras / "plain.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()
from py.api import maintenance
with (
patch.object(PromptManagerConfig, "SAFETY_BACKUP", True),
patch.object(
maintenance,
"take_safety_backup",
side_effect=maintenance.SafetyBackupFailed("no room"),
),
):
job = await self._failed_job()
self.assertEqual(job.scope["failure"], "backup-failed")
async def test_cancel_right_after_the_start_writes_nothing(self):
self._two_loras()
before = _snapshot(self.api.db)
resp, plan = await self._post(PREVIEW)
resp, body = await self._post(APPLY, {"token": plan["token"]})
job = self.api._jobs.get(body["job"]["id"])
self.api._jobs.cancel(job.id) # at once, while it stages
await job.task
self.assertEqual(job.state, "cancelled")
self.assertEqual(_snapshot(self.api.db), before)
def plan_title():
return lora_import.JOB_TITLE
def _read(path):
try:
with open(path, "rb") as handle:
return handle.read()
except FileNotFoundError:
return None
def _size(path):
return os.path.getsize(path) if os.path.exists(path) else None
class LoraImportRestoreTests(LoraImportAPITestCase):
"""The import and a restore (#249 S8) never run beside each other."""
async def test_a_restore_is_refused_while_the_import_runs(self):
self._two_loras()
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))
resp = await self.client.request("POST", "/prompt_manager/restore")
body = json.loads(await resp.text())
self.assertEqual(resp.status, 409)
self.assertIn(lora_import.JOB_TITLE, body["error"])
# The workspace restore job is refused for the same running import.
blocker = self.api._jobs.blocker(RESTORE_JOB_KIND, admission="exclusive")
self.assertIs(blocker, job)
self.gate.set()
await job.task
self.assertEqual(job.state, "finished")
async def test_the_import_is_refused_while_a_restore_runs(self):
self._two_loras()
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
before = _snapshot(self.api.db)
async with self.api._jobs.hold(
RESTORE_JOB_KIND, "Restore library", admission="exclusive"
):
# A review reads the library: refused, not run beside the switch.
resp, body = await self._post(PREVIEW)
self.assertEqual(resp.status, 409, body)
self.assertEqual(body["code"], "busy")
self.assertEqual(body["running"]["title"], "Restore library")
resp, body = await self._post(APPLY, {"token": plan["token"]})
self.assertEqual(resp.status, 409, body)
self.assertEqual(body["code"], "busy")
# The admin page's scan says so in its own error event.
resp = await self.client.request("POST", SCAN)
events = self._sse(await resp.text())
self.assertEqual(events[-1]["type"], "error")
self.assertIn("Restore library", events[-1]["message"])
self.assertEqual(_snapshot(self.api.db), before)
self.assertIsNone(self.api._jobs.active("lora_import"))
async def test_a_running_review_holds_a_lease_a_restore_waits_for(self):
self._two_loras()
seen = []
real_preview = lora_import.preview
def preview(*args):
seen.append(self.api._jobs.leases)
return real_preview(*args)
with patch.object(lora_import, "preview", preview):
resp, plan = await self._post(PREVIEW)
self.assertEqual(resp.status, 200, plan)
self.assertEqual(seen, [{LORA_PREVIEW_LEASE: 1}])
self.assertEqual(self.api._jobs.leases, {})
class LoraImportStageTests(unittest.TestCase):
"""The stages over a disposable LoraManager install (#355)."""
def setUp(self):
self.tmp = Path(tempfile.mkdtemp()).resolve()
self.addCleanup(lambda: __import__("shutil").rmtree(self.tmp, True))
self.lm = _make_lora_manager(self.tmp / "custom_nodes")
self.loras = self.tmp / "loras"
self.loras.mkdir()
self.settings = self.tmp / "settings"
self.settings.mkdir()
self.roots = [str(self.loras)]
patchers = [
patch.object(
lora_utils,
"custom_nodes_directories",
return_value=[self.tmp / "custom_nodes"],
),
patch.object(lora_utils, "default_lora_roots", lambda: list(self.roots)),
patch.dict(os.environ, {"LORA_MANAGER_SETTINGS_DIR": str(self.settings)}),
]
for patcher in patchers:
patcher.start()
self.addCleanup(patcher.stop)
def settings_file(self, data):
(self.settings / "settings.json").write_text(json.dumps(data))
def discover(self):
return lora_import.discover(str(self.lm))
def test_discover_sorts_and_drops_repeats(self):
a = self.loras / "a"
b = self.loras / "a" / "b"
b.mkdir(parents=True)
_write_metadata(b, "zeta")
_write_metadata(a, "alpha")
self.roots = [str(a), str(b)]
discovery = self.discover()
names = [s.metadata_path.name for s in discovery.sidecars]
self.assertEqual(names, ["alpha.metadata.json", "zeta.metadata.json"])
self.assertEqual(discovery.scanned, 2)
def test_discover_without_an_install_finds_nothing(self):
_write_metadata(self.loras, "alpha")
discovery = lora_import.discover(str(self.tmp))
self.assertEqual(discovery.layout.error, "no LoraManager install found")
self.assertEqual(discovery.sidecars, [])
def test_fingerprint_follows_size_and_mtime(self):
meta = _write_metadata(self.loras, "one")
first = lora_import.fingerprint(self.discover())
self.assertEqual(lora_import.fingerprint(self.discover()), first)
meta.write_text(meta.read_text() + " ")
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
changed = lora_import.fingerprint(self.discover())
meta.unlink()
self.assertNotEqual(lora_import.fingerprint(self.discover()), changed)
def test_fingerprint_ignores_unrelated_settings(self):
_write_metadata(self.loras, "one")
self.settings_file({"civitai_api_key": "a"})
first = lora_import.fingerprint(self.discover())
self.settings_file({"civitai_api_key": "b", "theme": "dark"})
self.assertEqual(lora_import.fingerprint(self.discover()), first)
sidecars = self.tmp / "sidecars"
sidecars.mkdir()
self.settings_file(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(sidecars),
}
)
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
# #362: a skipped sidecar is part of what the review counted.
def _assert_fingerprint_follows(self, change):
first = lora_import.fingerprint(self.discover())
self.assertEqual(lora_import.fingerprint(self.discover()), first)
change()
self.assertNotEqual(lora_import.fingerprint(self.discover()), first)
def _assert_add_edit_remove_change_it(self, path, text):
self._assert_fingerprint_follows(lambda: path.write_text(text))
self._assert_fingerprint_follows(
lambda: path.write_text(path.read_text() + " ")
)
self._assert_fingerprint_follows(path.unlink)
def _centralized(self):
sidecars = self.tmp / "sidecars"
sidecars.mkdir()
self.settings_file(
{
"sidecar_storage_mode": "centralized",
"sidecar_storage_path": str(sidecars),
}
)
model = self.loras / "kept.safetensors"
model.write_bytes(b"m")
(sidecars / "kept.metadata.json").write_text(
json.dumps({"file_path": str(model)})
)
return sidecars
def test_fingerprint_follows_a_centralized_sidecar_whose_model_is_missing(self):
sidecars = self._centralized()
gone = self.loras / "gone.safetensors"
self._assert_add_edit_remove_change_it(
sidecars / "gone.metadata.json", json.dumps({"file_path": str(gone)})
)
self.assertEqual(len(self.discover().sidecars), 1)
def test_fingerprint_follows_a_checkpoint_sidecar_outside_the_lora_roots(self):
sidecars = self._centralized()
checkpoints = self.tmp / "checkpoints"
checkpoints.mkdir()
ckpt = checkpoints / "base.safetensors"
ckpt.write_bytes(b"c")
self._assert_add_edit_remove_change_it(
sidecars / "base.metadata.json", json.dumps({"file_path": str(ckpt)})
)
def test_fingerprint_follows_an_alongside_sidecar_without_its_model(self):
_write_metadata(self.loras, "kept")
self._assert_add_edit_remove_change_it(
self.loras / "orphan.metadata.json", json.dumps({"model_name": "orphan"})
)
self.assertEqual(len(self.discover().sidecars), 1)
def test_fingerprint_follows_a_missing_file_between_stat_calls(self):
meta = _write_metadata(self.loras, "one")
discovery = self.discover()
first = lora_import.fingerprint(discovery)
meta.unlink()
self.assertNotEqual(lora_import.fingerprint(discovery), first)
def test_stage_cancel_before_a_file_raises(self):
_write_metadata(self.loras, "one")
with self.assertRaises(lora_import.ImportCancelled):
lora_import.stage(
self.discover(),
self.tmp,
media_extensions=MEDIA,
cancelled=lambda: True,
)
def test_local_media_means_no_download(self):
_write_metadata(self.loras, "neon", images=["https://civitai.com/a.png"])
(self.loras / "neon.mp4").write_bytes(b"video")
def refuse(*args, **kwargs):
raise AssertionError("a LoRA with local media downloads nothing")
with patch.object(lora_utils, "download_civitai_images", refuse):
staged = lora_import.stage(
self.discover(), self.tmp, media_extensions=MEDIA
)
self.assertEqual(staged.skipped, [])
self.assertEqual(
staged.records[0]["images"], [os.path.realpath(self.loras / "neon.mp4")]
)
def test_a_lora_without_local_media_downloads(self):
_write_metadata(self.loras, "neon", images=["https://civitai.com/a.png"])
calls = []
def download(metadata, path, cache_dir, api_key, stats=None):
calls.append(path.name)
return ["/cache/a.jpg"]
with patch.object(lora_utils, "download_civitai_images", download):
staged = lora_import.stage(
self.discover(), self.tmp, media_extensions=MEDIA
)
self.assertEqual(calls, ["neon.metadata.json"])
self.assertEqual(staged.records[0]["images"], ["/cache/a.jpg"])
def test_alongside_records_match_the_previous_import(self):
_write_metadata(
self.loras,
"neon",
words=["glow", "glow"],
images=["https://civitai.com/a.png"],
prompt="neon street",
)
(self.loras / "neon.png").write_bytes(b"x")
staged = lora_import.stage(self.discover(), self.tmp, media_extensions=MEDIA)
self.assertEqual(
staged.records,
[
{
"name": "neon",
"text": "neon street",
"hash": generate_prompt_hash("neon street"),
"tags": ["lora-manager", "lora:neon", "glow"],
"images": [os.path.realpath(self.loras / "neon.png")],
}
],
)
def test_preview_counts_local_media_and_downloads(self):
_write_metadata(self.loras, "local", images=["https://civitai.com/l.png"])
(self.loras / "local.webp").write_bytes(b"x")
_write_metadata(
self.loras,
"remote",
images=["https://civitai.com/r1.png", "https://civitai.com/r2.png"],
)
(self.loras / "orphan.metadata.json").write_text("{}")
found = lora_import.preview(str(self.lm), self.tmp / "cache", MEDIA)
self.assertEqual(
{
k: found[k]
for k in (
"lora_files",
"with_metadata",
"skipped",
"with_local_media",
"planned_downloads",
)
},
{
"lora_files": 3,
"with_metadata": 2,
"skipped": 1,
"with_local_media": 1,
"planned_downloads": 2,
},
)
self.assertEqual(found["names"], ["local", "remote"])
self.assertEqual(found["layout"]["mode"], "alongside")
def test_preview_of_an_unusable_layout(self):
(self.settings / "settings.json").write_text("nope")
_write_metadata(self.loras, "one")
found = lora_import.preview(str(self.lm), self.tmp / "cache", MEDIA)
self.assertEqual(
found["layout"]["error"], "the settings file is not valid JSON"
)
self.assertEqual((found["lora_files"], found["with_metadata"]), (0, 0))
MEDIA = frozenset({".png", ".jpg", ".jpeg", ".webp", ".mp4", ".webm"})
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink