Vito
863d717435
image.civitai.com now answers every image request with a redirect to blobs-b2.civitai.com, and the downloader refused all redirects, so every LoRA import download failed. Follow a redirect only over HTTPS to civitai.com or a subdomain of it (parsed host name, default port, no credentials in the URL), at most 3 hops, and never send the API key on a redirected request. Every other redirect is still refused; the size cap, timeout and image checks apply to the final response. Logged URLs drop their query, so signed storage URLs stay out of the log. The System import job's status line now counts failed downloads, and when every planned download fails it leads with a warning, which the Download errors outcome lists too. The legacy admin scan's events are unchanged.
1024 lines
40 KiB
Python
1024 lines
40 KiB
Python
"""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, "find_lora_directories", 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):
|
|
_write_metadata(
|
|
self.loras,
|
|
"neon",
|
|
words=["glow"],
|
|
images=["https://civitai.com/a.png"],
|
|
prompt="neon street",
|
|
)
|
|
(self.loras / "neon.png").write_bytes(_png_bytes())
|
|
_write_metadata(self.loras, "plain")
|
|
|
|
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", "current_imports")},
|
|
{"lora_files": 3, "with_metadata": 2, "current_imports": 2},
|
|
)
|
|
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))
|
|
self.assertNotIn(str(self.tmp_root), json.dumps(plan))
|
|
|
|
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.safetensors.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"])), 2)
|
|
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_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_utils, "find_lora_directories", 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("{{{")
|
|
|
|
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_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.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()
|
|
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()
|
|
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):
|
|
|
|
def setUp(self):
|
|
self.tmp = Path(tempfile.mkdtemp())
|
|
self.addCleanup(lambda: __import__("shutil").rmtree(self.tmp, True))
|
|
|
|
def test_discover_sorts_and_drops_repeats(self):
|
|
a = self.tmp / "a"
|
|
b = self.tmp / "a" / "b"
|
|
b.mkdir(parents=True)
|
|
_write_metadata(b, "zeta")
|
|
_write_metadata(a, "alpha")
|
|
with patch.object(
|
|
lora_utils, "find_lora_directories", return_value=[str(a), str(b)]
|
|
):
|
|
files = lora_import.discover(str(self.tmp))
|
|
self.assertEqual([f.name for f in files][0], "alpha.safetensors.metadata.json")
|
|
self.assertEqual(len(files), 2)
|
|
|
|
def test_fingerprint_follows_size_and_mtime(self):
|
|
meta = _write_metadata(self.tmp, "one")
|
|
first = lora_import.fingerprint([meta])
|
|
self.assertEqual(lora_import.fingerprint([meta]), first)
|
|
meta.write_text(meta.read_text() + " ")
|
|
self.assertNotEqual(lora_import.fingerprint([meta]), first)
|
|
meta.unlink()
|
|
self.assertNotEqual(lora_import.fingerprint([meta]), first)
|
|
|
|
def test_stage_cancel_before_a_file_raises(self):
|
|
meta = _write_metadata(self.tmp, "one")
|
|
with self.assertRaises(lora_import.ImportCancelled):
|
|
lora_import.stage([meta], self.tmp, cancelled=lambda: True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|