With Settings set to the hidden .amb folder on this computer, nine tests read the real api/config.json and looked for side files in the wrong place. tests/isolated_layout.py points model_sidecars at no config, so the tests that write side files always get the default. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
167 lines
8.6 KiB
Python
167 lines
8.6 KiB
Python
"""The user's layer comes first: editor edits, covers and what a scan may touch."""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
PLUGIN = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1])]
|
|
|
|
import isolated_layout # noqa: F401 - side files beside the model, whatever this computer chose
|
|
import scraper
|
|
from aiohttp.test_utils import make_mocked_request
|
|
|
|
from api import metadata as metadata_module, model_media, model_metadata, scan_summary
|
|
from api.metadata import get_metadata
|
|
|
|
CIVITAI = {"id": 7, "modelId": 3, "name": "v2", "model": {"name": "Found"},
|
|
"images": [{"url": "https://example.invalid/cover.png"}]}
|
|
|
|
|
|
def write(path, data):
|
|
Path(path).write_text(json.dumps(data), encoding="utf-8")
|
|
|
|
|
|
class Request(dict): # aiohttp requests also hold values for the middleware
|
|
def __init__(self, body):
|
|
super().__init__()
|
|
self.body = body
|
|
|
|
async def json(self):
|
|
return self.body
|
|
|
|
|
|
class LayerTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.dir = Path(self.temp.name)
|
|
self.model = self.dir / "m.safetensors"
|
|
self.model.write_bytes(b"model")
|
|
self.base = str(self.dir / "m")
|
|
metadata_module.clear_metadata_cache()
|
|
|
|
def tearDown(self):
|
|
self.temp.cleanup()
|
|
|
|
def scan(self, *flags, cover=b"new civitai image"):
|
|
"""Runs the scraper with Civitai and downloads faked; returns what went to the network."""
|
|
calls = {"fetch": 0, "download": 0}
|
|
|
|
def fetch(_hash):
|
|
calls["fetch"] += 1
|
|
return json.loads(json.dumps(CIVITAI))
|
|
|
|
def download(_url, base_path):
|
|
calls["download"] += 1
|
|
Path(base_path + ".png").write_bytes(cover)
|
|
return base_path + ".png"
|
|
|
|
with mock.patch.object(sys, "argv", ["scraper.py", str(self.dir), *flags]), \
|
|
mock.patch.object(scraper, "fetch_civitai_info", side_effect=fetch), \
|
|
mock.patch.object(scraper, "download_media", side_effect=download), \
|
|
mock.patch.object(scraper, "calculate_sha256", side_effect=lambda p: "b" * 64 if p == str(self.model) else open(p, "rb").read().hex()), \
|
|
mock.patch.object(scraper.urllib.request, "urlopen", side_effect=OSError("no network in tests")):
|
|
scraper.main()
|
|
metadata_module.clear_metadata_cache()
|
|
return calls
|
|
|
|
def test_an_edited_but_unscanned_model_still_counts_as_new_and_gets_scanned(self):
|
|
write(self.base + ".civitai.info", {"anomalous_custom_name": "Mine"}) # how older versions saved edits
|
|
with mock.patch.object(scan_summary, "get_active_scan_paths", return_value=[str(self.dir)]):
|
|
self.assertEqual(scan_summary._summarize()["new"], 1)
|
|
self.assertEqual(self.scan("--virtual-rename")["fetch"], 1)
|
|
meta = get_metadata(str(self.model))
|
|
self.assertEqual(meta["info_source"], "civitai")
|
|
self.assertEqual(meta["custom_name"], "Mine") # the old edit still wins over Civitai's name
|
|
|
|
def test_user_layer_wins_and_an_emptied_field_stays_empty(self):
|
|
write(self.base + ".info", {"id": 7, "modelId": 3, "name": "v2", "anomalous_custom_name": "Found_v2"})
|
|
write(self.base + ".anomalous.json", {"format": 1, "custom_name": "Mine", "custom_notes": ""})
|
|
meta = get_metadata(str(self.model))
|
|
self.assertEqual(meta["custom_name"], "Mine")
|
|
self.assertEqual(meta["custom_notes"], "")
|
|
self.assertEqual(sorted(meta["user_fields"]), ["custom_name", "custom_notes"])
|
|
self.assertEqual(meta["info_source"], "civitai")
|
|
|
|
def test_info_source_tells_local_inference_and_unscanned_apart(self):
|
|
self.assertEqual(get_metadata(str(self.model))["info_source"], "")
|
|
write(self.base + ".info", {"id": -1, "modelId": -1, "name": "m"})
|
|
metadata_module.clear_metadata_cache()
|
|
self.assertEqual(get_metadata(str(self.model))["info_source"], "local")
|
|
|
|
def test_editor_writes_only_the_user_layer(self):
|
|
write(self.base + ".civitai.info", {"id": 7, "modelId": 3, "name": "from another tool"})
|
|
before = Path(self.base + ".civitai.info").read_text(encoding="utf-8")
|
|
with mock.patch.object(model_metadata, "resolve_folder_subdir", return_value=(str(self.dir), str(self.dir))):
|
|
asyncio.run(model_metadata.api_update_metadata(Request({
|
|
"type": "loras", "filename": "m.safetensors", "custom_name": "Mine", "custom_notes": "note"})))
|
|
self.assertEqual(Path(self.base + ".civitai.info").read_text(encoding="utf-8"), before)
|
|
user = json.loads(Path(self.base + ".anomalous.json").read_text(encoding="utf-8"))
|
|
self.assertEqual((user["custom_name"], user["custom_notes"]), ("Mine", "note"))
|
|
|
|
def edit(self, **fields):
|
|
with mock.patch.object(model_metadata, "resolve_folder_subdir", return_value=(str(self.dir), str(self.dir))):
|
|
return asyncio.run(model_metadata.api_update_metadata(Request({"type": "loras", "filename": "m.safetensors", **fields})))
|
|
|
|
def test_the_editor_keeps_a_source_link_and_refuses_what_is_no_web_address(self):
|
|
self.assertEqual(self.edit(custom_source_url="https://www.liblib.art/modelinfo/abc").status, 200)
|
|
user = json.loads(Path(self.base + ".anomalous.json").read_text(encoding="utf-8"))
|
|
self.assertEqual(user["source_url"], "https://www.liblib.art/modelinfo/abc")
|
|
for bad in ("javascript:alert(1)", "file:///C:/x.png", "https://"):
|
|
self.assertEqual(self.edit(custom_source_url=bad).status, 400)
|
|
self.assertEqual(json.loads(Path(self.base + ".anomalous.json").read_text(encoding="utf-8")), user)
|
|
self.edit(custom_source_url="") # emptied on purpose
|
|
metadata_module.clear_metadata_cache()
|
|
meta = get_metadata(str(self.model))
|
|
self.assertEqual((meta["source_url"], "source_url" in meta["user_fields"]), ("", True))
|
|
|
|
def test_the_editor_learns_whether_the_cover_shown_is_civitais(self):
|
|
def covers():
|
|
request = make_mocked_request("GET", "/anomalous/model_covers?type=loras&path_idx=0&subfolder=/&filename=m.safetensors")
|
|
with mock.patch.object(model_media, "resolve_folder_subdir", return_value=(str(self.dir), str(self.dir))):
|
|
return json.loads(asyncio.run(model_media.api_model_covers(request)).text)
|
|
|
|
self.assertEqual(covers(), {"civitai": "", "custom": False}) # no cover at all
|
|
Path(self.base + ".preview.png").write_bytes(b"my own picture")
|
|
self.assertEqual(covers(), {"civitai": "", "custom": True})
|
|
Path(self.base + ".civitai_bak.png").write_bytes(b"civitai image")
|
|
found = covers()
|
|
self.assertTrue(found["custom"])
|
|
self.assertIn("filename=m.civitai_bak.png", found["civitai"])
|
|
Path(self.base + ".preview.png").write_bytes(b"civitai image") # a copy of Civitai's
|
|
self.assertFalse(covers()["custom"])
|
|
|
|
def test_a_scan_does_not_refetch_what_is_on_disk(self):
|
|
write(self.base + ".info", {**CIVITAI, "anomalous_custom_name": "Found_v2"})
|
|
Path(self.base + ".civitai_bak.png").write_bytes(b"old civitai image")
|
|
Path(self.base + ".preview.png").write_bytes(b"old civitai image")
|
|
self.assertEqual(self.scan("--virtual-rename"), {"fetch": 0, "download": 0})
|
|
|
|
def test_fetching_again_keeps_the_users_cover(self):
|
|
write(self.base + ".info", CIVITAI)
|
|
Path(self.base + ".civitai_bak.png").write_bytes(b"old civitai image")
|
|
Path(self.base + ".preview.png").write_bytes(b"my own picture")
|
|
self.scan("--force-overwrite", "--virtual-rename")
|
|
self.assertEqual(Path(self.base + ".preview.png").read_bytes(), b"my own picture")
|
|
self.assertEqual(Path(self.base + ".civitai_bak.png").read_bytes(), b"new civitai image")
|
|
|
|
def test_fetching_again_replaces_an_unchanged_civitai_cover(self):
|
|
write(self.base + ".info", CIVITAI)
|
|
Path(self.base + ".civitai_bak.png").write_bytes(b"old civitai image")
|
|
Path(self.base + ".preview.png").write_bytes(b"old civitai image")
|
|
with mock.patch.object(scraper, "move_to_trash") as trash:
|
|
self.scan("--force-overwrite")
|
|
# The cover is Civitai's: the copy kept aside goes, the new picture replaces it in place.
|
|
trash.assert_called_once_with(self.base + ".civitai_bak.png")
|
|
self.assertEqual(Path(self.base + ".preview.png").read_bytes(), b"new civitai image")
|
|
info = json.loads(Path(self.base + ".info").read_text(encoding="utf-8"))
|
|
self.assertEqual(len(info["anomalous_cover_sha256"]), 64) # known as Civitai's from now on
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|