DemonGatanjieuandClaude Opus 5.5
45353ffea1
A model Civitai has no record of now has its safetensors header checked (model_identity. safetensors_damage: too short, header longer than the file, not JSON, or a tensor ending past the file's end; only the header is read). Such files get the unmatched reason "damaged": a red badge on the card, a line in the details, a warning on the scan page (the summary counts them) and first place in the unmatched list, with advice to download the model again. Tests: header cases, a scan telling a damaged file from a whole one; the scan report fixtures are whole safetensors files now. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
197 lines
10 KiB
Python
197 lines
10 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 sys
|
|
import tempfile
|
|
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}, "w": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}}).encode()
|
|
path.write_bytes(len(header).to_bytes(8, "little") + header + b"\0" * 4) # a whole 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_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("offline")
|
|
|
|
asked, lines = self.scan(down)
|
|
self.assertEqual(len(asked), scraper.CIVITAI_DOWN_AFTER)
|
|
self.assertIn({"event": "civitai_down"}, lines)
|
|
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"},
|
|
{"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"])
|
|
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"}])
|
|
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", "")})
|
|
|
|
|
|
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()
|