On an AutoDL server civitai.com timed out while civitai.red, Civitai's other address with the same API, answered in about a second, so scans there fell back to guessing from the files. civitai_client.api_hosts gives both addresses, the one that answered last first; the scan's lookups, import's and Model Check's (download_sources._civitai_json) and the network test ask the other one when the first does not answer at all. An HTTP error is Civitai's own answer and is not asked again. The test answers which host answered, and Settings says so when it was civitai.red. The pictures stay on image.civitai.com, which such networks often cannot reach either: after COVERS_DOWN_AFTER (2) covers in a row that would not download, a scan fetches no more covers instead of waiting on each. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
290 lines
16 KiB
Python
290 lines
16 KiB
Python
"""What a scan reports per model, why a model is unmatched, and the scan's result and log entry."""
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from aiohttp.test_utils import make_mocked_request
|
|
|
|
PLUGIN = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1])]
|
|
|
|
import scraper
|
|
from api import activity_log, model_metadata, scan_report, scan_summary, scanner
|
|
from model_identity import computed_file_identity
|
|
|
|
def digest(name):
|
|
return hashlib.sha256(name.encode()).hexdigest()
|
|
|
|
|
|
FOUND = {"id": 7, "modelId": 3, "name": "v2", "baseModel": "SDXL 1.0", "model": {"name": "Found"}, "images": []}
|
|
|
|
|
|
class ScanReportTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.dir = Path(self.temp.name)
|
|
self.report = self.dir / "report.jsonl"
|
|
|
|
def tearDown(self):
|
|
self.temp.cleanup()
|
|
|
|
def model(self, name, info=None):
|
|
path = self.dir / f"{name}.safetensors"
|
|
header = json.dumps({"__metadata__": {"name": name}, "lora_unet_a.lora_up.weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}}).encode()
|
|
path.write_bytes(len(header).to_bytes(8, "little") + header + b"\0" * 4) # a whole LoRA file (not damaged)
|
|
if info is not None:
|
|
Path(str(path)[:-12] + ".info").write_text(json.dumps(info), encoding="utf-8")
|
|
return path
|
|
|
|
def info(self, name):
|
|
return json.loads((self.dir / f"{name}.info").read_text(encoding="utf-8"))
|
|
|
|
def scan(self, fetch, *flags):
|
|
"""Runs the scraper (no real network); returns (models Civitai was asked about, report lines)."""
|
|
asked = []
|
|
|
|
def ask(file_hash):
|
|
asked.append(file_hash)
|
|
return fetch(file_hash)
|
|
|
|
argv = ["scraper.py", str(self.dir), "--report-file", str(self.report), *flags]
|
|
with mock.patch.object(sys, "argv", argv), \
|
|
mock.patch.object(scraper, "fetch_civitai_info", side_effect=ask), \
|
|
mock.patch.object(scraper, "calculate_sha256", side_effect=lambda p: digest(Path(p).stem)), \
|
|
mock.patch.object(scraper.urllib.request, "urlopen", side_effect=OSError("no network in tests")):
|
|
scraper.main()
|
|
lines = [json.loads(line) for line in self.report.read_text(encoding="utf-8").splitlines()]
|
|
self.report.unlink()
|
|
return asked, lines
|
|
|
|
def test_each_model_reports_its_outcome_and_unmatched_ones_keep_why(self):
|
|
self.model("aaa")
|
|
self.model("bbb")
|
|
asked, lines = self.scan(lambda h: dict(FOUND) if h == digest("aaa") else None)
|
|
self.assertEqual(len(asked), 2)
|
|
by_file = {Path(line["file"]).stem: line for line in lines if "file" in line}
|
|
self.assertEqual(by_file["aaa"]["status"], "matched")
|
|
self.assertEqual(by_file["aaa"]["name"], "Found · v2")
|
|
self.assertEqual((by_file["bbb"]["status"], by_file["bbb"]["reason"]), ("inferred", "not_found"))
|
|
self.assertEqual(self.info("bbb")["anomalous_unmatched_reason"], "not_found")
|
|
self.assertEqual(lines[-1], {"event": "done", "unchanged": 0, "civitai_down": False})
|
|
|
|
def test_when_covers_will_not_download_the_rest_fetch_none(self):
|
|
# Civitai answered on civitai.red, but its pictures (image.civitai.com) cannot be reached.
|
|
for name in ("c1", "c2", "c3", "c4"):
|
|
self.model(name)
|
|
tried = []
|
|
pictured = {**FOUND, "images": [{"url": "https://image.civitai.com/a/b/orig"}]}
|
|
with mock.patch.object(scraper, "download_media", side_effect=lambda url, base: tried.append(base) or None):
|
|
asked, _lines = self.scan(lambda h: dict(pictured))
|
|
self.assertEqual(len(asked), 4) # names and trigger words still come
|
|
self.assertEqual(len(tried), scraper.COVERS_DOWN_AFTER)
|
|
|
|
def test_when_civitai_stops_answering_the_rest_is_inferred_without_asking(self):
|
|
for name in ("m1", "m2", "m3", "m4"):
|
|
self.model(name)
|
|
|
|
def down(_hash):
|
|
raise scraper.CivitaiUnreachable("timeout")
|
|
|
|
asked, lines = self.scan(down)
|
|
self.assertEqual(len(asked), scraper.CIVITAI_DOWN_AFTER)
|
|
self.assertIn({"event": "civitai_down", "reason": "timeout"}, lines) # why, for the scan page
|
|
reasons = [line.get("reason") for line in lines if "file" in line]
|
|
self.assertEqual(reasons, ["network"] * 4)
|
|
|
|
def test_online_scans_look_up_models_civitai_was_never_asked_about(self):
|
|
for name, reason in (("off", "offline"), ("net", "network"), ("gone", "not_found")):
|
|
path = self.model(name)
|
|
Path(str(path)[:-12] + ".info").write_text(json.dumps({
|
|
"id": -1, "modelId": -1, "anomalous_unmatched_reason": reason,
|
|
"anomalous_file_identity": computed_file_identity(str(path), digest(name))}), encoding="utf-8")
|
|
asked, lines = self.scan(lambda h: dict(FOUND))
|
|
self.assertEqual(sorted(asked), sorted([digest("net"), digest("off")]))
|
|
self.assertEqual(self.info("gone")["id"], -1) # waits for "look up again"
|
|
self.assertEqual(self.info("off")["id"], 7)
|
|
# An offline scan leaves them alone.
|
|
asked, _ = self.scan(lambda h: dict(FOUND), "--offline-only")
|
|
self.assertEqual(asked, [])
|
|
|
|
def test_one_models_failure_does_not_end_the_scan(self):
|
|
self.model("bad")
|
|
self.model("good")
|
|
|
|
def fetch(file_hash):
|
|
if file_hash == digest("bad"):
|
|
raise PermissionError("locked")
|
|
return dict(FOUND)
|
|
|
|
_asked, lines = self.scan(fetch)
|
|
by_file = {Path(line["file"]).stem: line for line in lines if "file" in line}
|
|
self.assertEqual(by_file["bad"]["status"], "failed")
|
|
self.assertIn("locked", by_file["bad"]["error"])
|
|
self.assertEqual(by_file["good"]["status"], "matched")
|
|
|
|
def test_job_result_locates_models_and_is_kept_and_logged(self):
|
|
sub = self.dir / "sub"
|
|
sub.mkdir()
|
|
self.report.write_text("\n".join(json.dumps(line) for line in (
|
|
{"file": str(sub / "a.safetensors"), "status": "matched", "name": "A · v1", "cover": "new"},
|
|
{"file": str(self.dir / "b.safetensors"), "status": "inferred", "reason": "network", "base": "Flux.1 D"},
|
|
{"file": str(self.dir / "old.safetensors"), "status": "unchanged", "renamed_to": str(self.dir / "New_v1.safetensors")},
|
|
{"event": "civitai_down", "reason": "refused"},
|
|
{"event": "done", "unchanged": 5},
|
|
)), encoding="utf-8")
|
|
log = self.dir / "activity_log.json"
|
|
with mock.patch.object(scan_report, "log_path", return_value=str(log)), \
|
|
mock.patch.object(activity_log, "log_path", return_value=str(log)):
|
|
activity_log._entries = None
|
|
job = scan_report.ScanJob("all", {"offline_only": False})
|
|
job.add_folder(str(self.report), str(self.dir), ("loras", 0))
|
|
result = job.finish()
|
|
entry = activity_log.list_entries()[0]
|
|
activity_log._entries = None
|
|
self.assertFalse(self.report.exists())
|
|
self.assertEqual(result["counts"], {"matched": 1, "inferred": 1, "failed": 0, "covers": 1, "renamed": 1, "unchanged": 5})
|
|
self.assertTrue(result["civitai_down"])
|
|
self.assertEqual(result["civitai_error"], "refused")
|
|
first, _, renamed = result["files"]
|
|
self.assertEqual((first["type"], first["rel"], first["filename"]), ("loras", "sub/a.safetensors", "a.safetensors"))
|
|
self.assertEqual((renamed["filename"], renamed["renamed_from"]), ("New_v1.safetensors", "old.safetensors"))
|
|
self.assertEqual(json.loads((self.dir / "last_scan.json").read_text(encoding="utf-8"))["counts"]["matched"], 1)
|
|
self.assertEqual(entry["action"], "scan_done")
|
|
self.assertEqual(entry["detail"]["scan"]["counts"]["inferred"], 1)
|
|
self.assertEqual(len(entry["detail"]["files"]), 3)
|
|
|
|
def test_summary_lists_unmatched_and_new_models_with_why(self):
|
|
self.model("new")
|
|
self.model("gone", {"id": -1, "anomalous_unmatched_reason": "not_found", "baseModel": "SD 1.5"})
|
|
self.model("off", {"id": -1, "anomalous_unmatched_reason": "offline"})
|
|
self.model("found", {"id": 3})
|
|
with mock.patch.object(scan_summary, "get_active_scan_paths", return_value=[str(self.dir)]), \
|
|
mock.patch.object(scan_summary, "locate_root", return_value=("loras", 1)):
|
|
summary = scan_summary._summarize()
|
|
self.assertEqual((summary["total"], summary["matched"], summary["unmatched"], summary["new"], summary["pending"]), (4, 1, 2, 1, 1))
|
|
self.assertEqual(summary["new_models"], [{"type": "loras", "path_idx": 1, "filename": "new.safetensors", "rel": "new.safetensors", "state": "new"}])
|
|
reasons = {item["filename"]: (item["reason"], item["base"]) for item in summary["unmatched_models"]}
|
|
self.assertEqual(reasons, {"gone.safetensors": ("not_found", "SD 1.5"), "off.safetensors": ("offline", "")})
|
|
|
|
def test_components_are_counted_and_listed_on_their_own(self):
|
|
loras, vae = self.dir / "loras", self.dir / "vae"
|
|
for folder in (loras, vae):
|
|
folder.mkdir()
|
|
top = self.dir
|
|
self.dir = loras
|
|
self.model("style", {"id": 3})
|
|
self.dir = vae
|
|
self.model("found_vae", {"id": 4})
|
|
self.model("odd_vae", {"id": -1, "anomalous_unmatched_reason": "not_found"})
|
|
self.model("new_vae")
|
|
self.dir = top
|
|
roots = {str(loras): ("loras", 0), str(vae): ("vae", 0)}
|
|
with mock.patch.object(scan_summary, "get_active_scan_paths", return_value=[str(loras), str(vae)]), \
|
|
mock.patch.object(scan_summary, "locate_root", side_effect=lambda path: roots[path]):
|
|
summary = scan_summary._summarize()
|
|
self.assertEqual((summary["total"], summary["matched"], summary["unmatched"], summary["new"]), (4, 2, 1, 1))
|
|
self.assertEqual(summary["components"], {"total": 3, "matched": 1, "unmatched": 1, "new": 1, "pending": 0})
|
|
states = {item["filename"]: (item["state"], item.get("reason", "")) for item in summary["component_models"]}
|
|
self.assertEqual(states, {"found_vae.safetensors": ("matched", ""), "odd_vae.safetensors": ("unmatched", "not_found"),
|
|
"new_vae.safetensors": ("new", "")})
|
|
self.assertTrue(all(item.get("component") for item in summary["unmatched_models"] + summary["new_models"]))
|
|
|
|
def test_summary_has_the_scanned_models_covers_newest_first(self):
|
|
loras, vae = self.dir / "loras", self.dir / "vae"
|
|
top, now = self.dir, time.time()
|
|
for folder, name, info, age, cover in (
|
|
(loras, "old", {"id": 3}, 300, True), (loras, "fresh", {"id": 4}, 100, True),
|
|
(loras, "gone", {"id": -1, "anomalous_unmatched_reason": "not_found"}, 200, True),
|
|
(loras, "bare", {"id": 5}, 50, False), (loras, "new", None, 10, True), (vae, "part", {"id": 6}, 5, True)):
|
|
folder.mkdir(exist_ok=True)
|
|
self.dir = folder
|
|
path = self.model(name, info)
|
|
if cover:
|
|
(folder / f"{name}.preview.png").write_bytes(b"png")
|
|
os.utime(path, (now - age, now - age))
|
|
self.dir = top
|
|
roots = {str(loras): ("loras", 0), str(vae): ("vae", 0)}
|
|
with mock.patch.object(scan_summary, "get_active_scan_paths", return_value=[str(loras), str(vae)]), \
|
|
mock.patch.object(scan_summary, "locate_root", side_effect=lambda path: roots[path]):
|
|
summary = scan_summary._summarize()
|
|
with mock.patch.object(scan_summary, "COVER_SAMPLE", 1):
|
|
capped = scan_summary._summarize()
|
|
# Scanned ones with a cover, found or not; not the new one, the one without a cover or the component.
|
|
self.assertEqual([cover["name"] for cover in summary["covers"]], ["fresh", "gone", "old"])
|
|
self.assertIn("filename=fresh.preview.png", summary["covers"][0]["preview_url"])
|
|
self.assertEqual([cover["name"] for cover in capped["covers"]], ["fresh"])
|
|
|
|
def test_picker_lists_every_model_a_scan_reads_with_its_badge_fields(self):
|
|
self.model("new")
|
|
self.model("found", {**FOUND, "description": "<p>long text</p>"})
|
|
(self.dir / "found.preview.png").write_bytes(b"png")
|
|
self.model("skipped")
|
|
(self.dir / "skipped.anomalous.json").write_text(json.dumps({"format": 1, "scan": "skip"}), encoding="utf-8")
|
|
(self.dir / "sub").mkdir()
|
|
self.dir, top = self.dir / "sub", self.dir
|
|
self.model("deep", {"id": -1, "anomalous_unmatched_reason": "damaged"})
|
|
self.dir = top
|
|
(self.dir / "old.ckpt").write_bytes(b"never read")
|
|
with mock.patch.object(scan_summary, "get_active_scan_paths", return_value=[str(self.dir)]), \
|
|
mock.patch.object(scan_summary, "locate_root", return_value=("loras", 1)):
|
|
listed = scan_summary._pickable()
|
|
with mock.patch.object(scan_summary, "MAX_PICKABLE", 2):
|
|
capped = scan_summary._pickable()
|
|
models = {item["rel"]: item for item in listed["models"]}
|
|
self.assertEqual(sorted(models), ["found.safetensors", "new.safetensors", "skipped.safetensors", "sub/deep.safetensors"])
|
|
self.assertEqual(listed["total"], 4)
|
|
self.assertEqual(listed["roots"], [{"type": "loras", "path_idx": 1, "name": self.dir.name, "parent": self.dir.parent.name}])
|
|
found = models["found.safetensors"]
|
|
self.assertEqual((found["type"], found["path_idx"], found["filename"]), ("loras", 1, "found.safetensors"))
|
|
self.assertIn("filename=found.preview.png", found["preview_url"])
|
|
self.assertEqual(found["size"], (self.dir / "found.safetensors").stat().st_size)
|
|
self.assertEqual(set(found["metadata"]), set(scan_summary.PICKER_FIELDS)) # no description
|
|
self.assertEqual((found["metadata"]["info_source"], found["metadata"]["baseModel"]), ("civitai", "SDXL 1.0"))
|
|
self.assertEqual(models["skipped.safetensors"]["metadata"]["scan_choice"], "skip")
|
|
self.assertEqual(models["sub/deep.safetensors"]["metadata"]["unmatched_reason"], "damaged")
|
|
self.assertEqual(models["new.safetensors"]["metadata"]["info_source"], "")
|
|
self.assertEqual((len(capped["models"]), capped["total"]), (2, 4))
|
|
|
|
|
|
class EditDetailTests(unittest.TestCase):
|
|
def test_editing_a_model_logs_what_changed_from_what(self):
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
model = Path(temp) / "m.safetensors"
|
|
model.write_bytes(b"m")
|
|
(Path(temp) / "m.anomalous.json").write_text(json.dumps({"custom_name": "Old"}), encoding="utf-8")
|
|
request = make_mocked_request("POST", "/anomalous/update_metadata")
|
|
body = {"type": "loras", "filename": "m.safetensors", "custom_name": "New", "custom_notes": ""}
|
|
request.json = lambda: asyncio.sleep(0, body)
|
|
with mock.patch.object(model_metadata, "resolve_folder_subdir", return_value=(temp, temp)):
|
|
asyncio.run(model_metadata.api_update_metadata(request))
|
|
self.assertEqual(request[activity_log.ACTIVITY_DETAIL], {"fields": [{"field": "custom_name", "before": "Old", "after": "New"}]})
|
|
self.assertEqual(activity_log._clean_detail({"fields": [{"field": "custom_notes", "after": "x" * 900}]})["fields"][0]["after"][-1], "…")
|
|
|
|
|
|
class ScanPlanTests(unittest.TestCase):
|
|
def test_picked_models_scan_their_own_folders_and_bad_targets_are_refused(self):
|
|
with tempfile.TemporaryDirectory() as temp:
|
|
sub = str(Path(temp) / "sub")
|
|
Path(sub).mkdir()
|
|
resolve = lambda _type, _idx, subfolder: (temp, sub if subfolder == "/sub" else temp)
|
|
with mock.patch.object(scanner, "resolve_folder_subdir", side_effect=resolve):
|
|
plan = scanner._global_scan_plan({"targets": [
|
|
{"type": "loras", "path_idx": 0, "subfolder": "/sub", "files": ["a.safetensors"]},
|
|
{"type": "loras", "path_idx": 0, "subfolder": "/", "files": []}, # nothing picked there
|
|
]})
|
|
self.assertEqual(plan, [(sub, ("loras", 0), temp, ["a.safetensors"], "loras")])
|
|
for bad in ({"targets": "loras"}, {"targets": [{"type": "loras", "files": "a.safetensors"}]}):
|
|
with self.assertRaises(ValueError):
|
|
scanner._global_scan_plan(bad)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|