1453 lines
55 KiB
Python
1453 lines
55 KiB
Python
"""
|
|
test_source_internals.py
|
|
|
|
Integration-smoke tests for pure (network-free) internal helper functions
|
|
in civarchive.py, civitai.py, and workflow_updater.py.
|
|
|
|
Goal: catch NameError / ImportError style regressions where a helper is
|
|
renamed or prefixed incorrectly (e.g. _extract_trained_words vs
|
|
extract_trained_words), or where a refactored unified function changes
|
|
its output contract.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import requests
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# civarchive internals
|
|
# ---------------------------------------------------------------------------
|
|
from core.sources.civarchive import (
|
|
_collect_download_urls_unified,
|
|
_fetch_remote_file_size_bytes,
|
|
_normalize_archive_version,
|
|
_normalize_download_url,
|
|
_resolve_file_size_bytes,
|
|
build_civarchive_custom_result,
|
|
parse_civarchive_url,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# civitai internals
|
|
# ---------------------------------------------------------------------------
|
|
from core.sources.civitai import (
|
|
_build_civitai_result_from_version,
|
|
_enrich_model_info_with_details,
|
|
_extract_model_images,
|
|
_extract_public_api_model_candidates,
|
|
_extract_trpc_model_candidates,
|
|
_find_matching_file_in_versions,
|
|
_normalize_civitai_file,
|
|
_search_civitai_public_api_candidates,
|
|
_search_civitai_red_candidates,
|
|
_search_civitai_trpc_candidates,
|
|
build_civitai_custom_result,
|
|
build_civitai_session_cookie,
|
|
clear_search_cache,
|
|
get_civitai_model_details,
|
|
get_model_info_by_hash,
|
|
get_model_info_for_file,
|
|
parse_civitai_url,
|
|
search_civitai_for_file,
|
|
)
|
|
from core.sources.common import is_remote_link_marked_dead
|
|
from core.type_utils import prepare_remote_size_probe_url
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# workflow_updater
|
|
# ---------------------------------------------------------------------------
|
|
from core.workflow_updater import (
|
|
convert_to_relative_path,
|
|
get_base_directory_for_model,
|
|
)
|
|
|
|
CIVARCHIVE_BASE = "https://civarchive.com"
|
|
|
|
|
|
def write_safetensors_stub(file_path, header):
|
|
header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
|
|
with open(file_path, "wb") as handle:
|
|
handle.write(len(header_bytes).to_bytes(8, byteorder="little"))
|
|
handle.write(header_bytes)
|
|
handle.write(b"tensor payload")
|
|
|
|
|
|
class CivitaiResultBuilderTests(unittest.TestCase):
|
|
|
|
def test_build_result_from_version_has_unified_builder_available(self):
|
|
result = _build_civitai_result_from_version(
|
|
model_id=123,
|
|
model_name="Aisha Greyrat",
|
|
model_type="LORA",
|
|
version={"id": 456, "baseModel": "Anima"},
|
|
file_info={
|
|
"name": "aisha-greyrat.safetensors",
|
|
"downloadUrl": "https://civitai.com/api/download/models/456",
|
|
"sizeKB": 10,
|
|
"hashes": {"SHA256": "abc123"},
|
|
},
|
|
tags=["character"],
|
|
)
|
|
|
|
self.assertEqual("civitai", result["source"])
|
|
self.assertEqual(123, result["model_id"])
|
|
self.assertEqual(456, result["version_id"])
|
|
self.assertEqual("Anima", result["base_model"])
|
|
self.assertEqual("abc123", result["sha256"])
|
|
self.assertEqual(10 * 1024, result["size"])
|
|
|
|
|
|
class CustomUrlResultBuilderTests(unittest.TestCase):
|
|
def test_civitai_custom_result_preserves_all_result_fields(self):
|
|
result = build_civitai_custom_result({
|
|
"model_id": 123,
|
|
"version_id": 456,
|
|
"name": "Example model",
|
|
"type": "LORA",
|
|
"tags": ["style"],
|
|
"images": [{"url": "https://example.test/image"}],
|
|
"description": "Model description",
|
|
"version_url": "https://civitai.com/models/123?modelVersionId=456",
|
|
"selected_version": {
|
|
"id": 456,
|
|
"name": "v1",
|
|
"base_model": "SD 1.5",
|
|
"trained_words": ["example"],
|
|
"description": "Version notes",
|
|
"files": [{
|
|
"name": "example.safetensors",
|
|
"type": "Model",
|
|
"download_url": "https://civitai.com/api/download/models/456",
|
|
"size": 123,
|
|
"sha256": "abc123",
|
|
"hashes": {"SHA256": "abc123"},
|
|
}],
|
|
},
|
|
})
|
|
|
|
self.assertEqual({
|
|
"source": "civitai",
|
|
"details_source": "civitai",
|
|
"model_id": 123,
|
|
"version_id": 456,
|
|
"name": "Example model",
|
|
"version_name": "v1",
|
|
"type": "LORA",
|
|
"filename": "example.safetensors",
|
|
"url": "https://civitai.com/models/123?modelVersionId=456",
|
|
"version_url": "https://civitai.com/models/123?modelVersionId=456",
|
|
"download_url": "https://civitai.com/api/download/models/456",
|
|
"size": 123,
|
|
"base_model": "SD 1.5",
|
|
"tags": ["style"],
|
|
"trained_words": ["example"],
|
|
"images": [{"url": "https://example.test/image"}],
|
|
"description": "Version notes",
|
|
"model_description": "Model description",
|
|
"sha256": "abc123",
|
|
"hashes": {"SHA256": "abc123"},
|
|
"match_type": "custom_url",
|
|
"custom_url": True,
|
|
}, result)
|
|
|
|
def test_civarchive_custom_result_preserves_provider_fields(self):
|
|
result = build_civarchive_custom_result({
|
|
"model_id": 789,
|
|
"version_id": 987,
|
|
"name": "Archived model",
|
|
"type": "Checkpoint",
|
|
"tags": ["archive"],
|
|
"platform": "civitai",
|
|
"platform_url": "https://civitai.com/models/789",
|
|
"description": "Archive description",
|
|
"version_url": "https://civarchive.com/models/789?modelVersionId=987",
|
|
"selected_version": {
|
|
"id": 987,
|
|
"name": "archive-v1",
|
|
"base_model": "SDXL",
|
|
"trained_words": ["archived"],
|
|
"description": "Archive notes",
|
|
"platform_url": "https://civitai.com/models/789?modelVersionId=987",
|
|
"files": [{
|
|
"name": "archived.safetensors",
|
|
"type": "Model",
|
|
"download_url": "https://civarchive.com/download/987",
|
|
"download_urls": [
|
|
"https://civarchive.com/download/987",
|
|
"https://mirror.example.test/987",
|
|
],
|
|
"size": 456,
|
|
"sha256": "def456",
|
|
"hashes": {"SHA256": "def456"},
|
|
}],
|
|
},
|
|
})
|
|
|
|
self.assertEqual({
|
|
"source": "civarchive",
|
|
"details_source": "civarchive",
|
|
"model_id": 789,
|
|
"version_id": 987,
|
|
"name": "Archived model",
|
|
"version_name": "archive-v1",
|
|
"type": "Checkpoint",
|
|
"filename": "archived.safetensors",
|
|
"url": "https://civarchive.com/models/789?modelVersionId=987",
|
|
"version_url": "https://civarchive.com/models/789?modelVersionId=987",
|
|
"platform_url": "https://civitai.com/models/789",
|
|
"download_url": "https://civarchive.com/download/987",
|
|
"download_urls": [
|
|
"https://civarchive.com/download/987",
|
|
"https://mirror.example.test/987",
|
|
],
|
|
"size": 456,
|
|
"base_model": "SDXL",
|
|
"tags": ["archive"],
|
|
"trained_words": ["archived"],
|
|
"images": [],
|
|
"description": "Archive notes",
|
|
"sha256": "def456",
|
|
"hash": "def456",
|
|
"hashes": {"SHA256": "def456"},
|
|
"platform": "civitai",
|
|
"match_type": "custom_url",
|
|
"custom_url": True,
|
|
}, result)
|
|
|
|
|
|
class CivitaiModelDetailsTests(unittest.TestCase):
|
|
|
|
def test_preserves_full_selected_version_for_metadata_sidecar(self):
|
|
response = {
|
|
"id": 123,
|
|
"name": "Alt-Girl/E-Girl",
|
|
"type": "LORA",
|
|
"description": "Model description",
|
|
"tags": ["style", "e-girl"],
|
|
"allowNoCredit": True,
|
|
"allowCommercialUse": ["Image", "Sell"],
|
|
"creator": {"username": "TheAINovice", "image": None},
|
|
"stats": {"downloadCount": 1000},
|
|
"modelVersions": [
|
|
{
|
|
"id": 456,
|
|
"name": "Old version",
|
|
"baseModel": "Krea 1",
|
|
"images": [],
|
|
"files": [],
|
|
},
|
|
{
|
|
"id": 789,
|
|
"name": "Krea2 v1.0",
|
|
"baseModel": "Krea 2",
|
|
"createdAt": "2026-07-14T05:52:54.321Z",
|
|
"stats": {"downloadCount": 1705, "thumbsUpCount": 106},
|
|
"images": [
|
|
{
|
|
"url": "https://image.civitai.com/example.jpeg",
|
|
"nsfwLevel": 4,
|
|
"hash": "blurhash",
|
|
"meta": {
|
|
"seed": 884670593557080,
|
|
"prompt": "AltGirl prompt",
|
|
"steps": 8,
|
|
"additionalResources": [
|
|
{
|
|
"name": "support-style.safetensors",
|
|
"type": "lora",
|
|
"strength": 0.5,
|
|
}
|
|
],
|
|
},
|
|
}
|
|
],
|
|
"files": [
|
|
{
|
|
"id": 321,
|
|
"name": "AltGirlKrea.safetensors",
|
|
"type": "Model",
|
|
"pickleScanResult": "Success",
|
|
"virusScanResult": "Success",
|
|
"hashes": {"SHA256": "A" * 64},
|
|
"primary": True,
|
|
}
|
|
],
|
|
},
|
|
],
|
|
}
|
|
|
|
with patch(
|
|
"core.sources.civitai.execute_provider_json_request",
|
|
return_value=response,
|
|
):
|
|
details = get_civitai_model_details(123, 789)
|
|
|
|
self.assertIsNotNone(details)
|
|
civitai = details["civitai"]
|
|
self.assertEqual(789, civitai["id"])
|
|
self.assertEqual(123, civitai["modelId"])
|
|
self.assertEqual(4, civitai["images"][0]["nsfwLevel"])
|
|
self.assertEqual("AltGirl prompt", civitai["images"][0]["meta"]["prompt"])
|
|
self.assertEqual(
|
|
{
|
|
"name": "Alt-Girl/E-Girl",
|
|
"versionName": "Krea2 v1.0",
|
|
"type": "LORA",
|
|
"modelId": 123,
|
|
"modelVersionId": 789,
|
|
"url": "https://civitai.com/models/123?modelVersionId=789",
|
|
"primary": True,
|
|
},
|
|
civitai["images"][0]["resources"][0],
|
|
)
|
|
self.assertEqual(
|
|
{
|
|
"name": "support-style.safetensors",
|
|
"type": "lora",
|
|
"strength": 0.5,
|
|
},
|
|
civitai["images"][0]["resources"][1],
|
|
)
|
|
self.assertEqual("Success", civitai["files"][0]["virusScanResult"])
|
|
self.assertEqual(1705, civitai["stats"]["downloadCount"])
|
|
self.assertEqual(
|
|
["Image", "Sell"],
|
|
civitai["model"]["allowCommercialUse"],
|
|
)
|
|
self.assertEqual("TheAINovice", civitai["creator"]["username"])
|
|
self.assertNotIn("modelVersions", civitai["model"])
|
|
|
|
|
|
class CivitaiHashLookupCacheTests(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
clear_search_cache()
|
|
|
|
def tearDown(self):
|
|
clear_search_cache()
|
|
|
|
@staticmethod
|
|
def _http_error(status_code):
|
|
response = requests.Response()
|
|
response.status_code = status_code
|
|
return requests.HTTPError(response=response)
|
|
|
|
def test_only_not_found_response_is_cached(self):
|
|
with patch(
|
|
"core.sources.civitai.execute_provider_json_request",
|
|
side_effect=self._http_error(404),
|
|
) as request_json:
|
|
self.assertIsNone(get_model_info_by_hash("a" * 64))
|
|
self.assertIsNone(get_model_info_by_hash("a" * 64))
|
|
|
|
request_json.assert_called_once()
|
|
self.assertTrue(request_json.call_args.kwargs["raise_on_error"])
|
|
|
|
def test_transient_http_error_is_not_cached(self):
|
|
with patch(
|
|
"core.sources.civitai.execute_provider_json_request",
|
|
side_effect=self._http_error(503),
|
|
) as request_json:
|
|
self.assertIsNone(get_model_info_by_hash("b" * 64))
|
|
self.assertIsNone(get_model_info_by_hash("b" * 64))
|
|
|
|
self.assertEqual(request_json.call_count, 2)
|
|
|
|
def test_unexpected_json_shape_is_not_cached(self):
|
|
with patch(
|
|
"core.sources.civitai.execute_provider_json_request",
|
|
return_value=[],
|
|
) as request_json:
|
|
self.assertIsNone(get_model_info_by_hash("c" * 64))
|
|
self.assertIsNone(get_model_info_by_hash("c" * 64))
|
|
|
|
self.assertEqual(request_json.call_count, 2)
|
|
|
|
|
|
# ===========================================================================
|
|
# civitai - local safetensors fallback
|
|
# ===========================================================================
|
|
class CivitaiLocalFallbackTests(unittest.TestCase):
|
|
|
|
def test_get_model_info_for_file_infers_base_model_when_civitai_misses(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
file_path = os.path.join(tmpdir, "orphan_flux.safetensors")
|
|
write_safetensors_stub(
|
|
file_path,
|
|
{
|
|
"__metadata__": {},
|
|
"double_blocks.0.img_attn.qkv.weight": {
|
|
"dtype": "F16",
|
|
"shape": [1],
|
|
},
|
|
},
|
|
)
|
|
|
|
with patch("core.sources.civitai.get_model_info_by_hash", return_value=None):
|
|
result = get_model_info_for_file(file_path)
|
|
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result["source"], "local")
|
|
self.assertEqual(result["base_model"], "Flux.1 D")
|
|
self.assertTrue(result["base_model_inferred"])
|
|
self.assertEqual(result["base_model_source"], "safetensors_header")
|
|
|
|
def test_get_model_info_for_file_infers_base_model_when_civitai_errors(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
file_path = os.path.join(tmpdir, "offline_sd3.safetensors")
|
|
write_safetensors_stub(
|
|
file_path,
|
|
{
|
|
"__metadata__": {},
|
|
"joint_blocks.0.x_block.attn.qkv.weight": {
|
|
"dtype": "F16",
|
|
"shape": [1],
|
|
},
|
|
},
|
|
)
|
|
|
|
with patch(
|
|
"core.sources.civitai.get_model_info_by_hash",
|
|
side_effect=OSError("offline"),
|
|
):
|
|
result = get_model_info_for_file(file_path)
|
|
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result["source"], "local")
|
|
self.assertEqual(result["base_model"], "SD 3")
|
|
self.assertTrue(result["base_model_inferred"])
|
|
|
|
def test_get_model_info_for_file_uses_safetensors_header_metadata_on_miss(self):
|
|
expected_hash = "d" * 64
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
file_path = os.path.join(tmpdir, "header_lora.safetensors")
|
|
write_safetensors_stub(
|
|
file_path,
|
|
{
|
|
"__metadata__": {
|
|
"modelspec.title": "Header LoRA",
|
|
"modelspec.author": "Header Maker",
|
|
"modelspec.description": "Local header description",
|
|
"modelspec.hash.sha256": expected_hash,
|
|
"modelspec.tags": json.dumps(["style"]),
|
|
"modelspec.trigger_phrase": "header trigger",
|
|
"ss_base_model_version": "pony",
|
|
"ss_network_module": "networks.lora",
|
|
"sshs_model_hash": "e" * 64,
|
|
},
|
|
},
|
|
)
|
|
|
|
with patch("core.sources.civitai.get_model_info_by_hash", return_value=None):
|
|
result = get_model_info_for_file(file_path)
|
|
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result["source"], "local")
|
|
self.assertEqual(result["details_source"], "safetensors_header")
|
|
self.assertEqual(result["model_name"], "Header LoRA")
|
|
self.assertEqual(result["creator"]["username"], "Header Maker")
|
|
self.assertEqual(result["description"], "Local header description")
|
|
self.assertEqual(result["base_model"], "Pony")
|
|
self.assertEqual(result["sha256"], expected_hash)
|
|
self.assertEqual(result["tags"], ["style"])
|
|
self.assertEqual(result["trained_words"], ["header trigger"])
|
|
self.assertTrue(result["local_metadata_available"])
|
|
|
|
|
|
# ===========================================================================
|
|
# civarchive - _normalize_download_url
|
|
# ===========================================================================
|
|
class NormalizeDownloadUrlTests(unittest.TestCase):
|
|
|
|
def test_absolute_https_url_passes_through(self):
|
|
url = "https://civarchive.com/api/download/models/123"
|
|
self.assertEqual(_normalize_download_url(url), url)
|
|
|
|
def test_absolute_http_url_passes_through(self):
|
|
url = "http://example.com/model.safetensors"
|
|
self.assertEqual(_normalize_download_url(url), url)
|
|
|
|
def test_protocol_relative_url_gets_https(self):
|
|
result = _normalize_download_url("//civarchive.com/api/download/models/99")
|
|
self.assertEqual(result, "https://civarchive.com/api/download/models/99")
|
|
|
|
def test_relative_api_path_gets_base_url(self):
|
|
result = _normalize_download_url("/api/download/models/42")
|
|
self.assertTrue(result.startswith(CIVARCHIVE_BASE))
|
|
self.assertIn("/api/download/models/42", result)
|
|
|
|
def test_none_returns_none(self):
|
|
self.assertIsNone(_normalize_download_url(None))
|
|
|
|
def test_empty_string_returns_none(self):
|
|
self.assertIsNone(_normalize_download_url(""))
|
|
|
|
def test_whitespace_only_returns_none(self):
|
|
self.assertIsNone(_normalize_download_url(" "))
|
|
|
|
def test_non_string_returns_none(self):
|
|
self.assertIsNone(_normalize_download_url(12345))
|
|
|
|
def test_non_http_relative_path_returns_none(self):
|
|
self.assertIsNone(_normalize_download_url("/some/other/path"))
|
|
|
|
|
|
# ===========================================================================
|
|
# shared remote link dead marker
|
|
# ===========================================================================
|
|
class ArchiveLinkIsDeadTests(unittest.TestCase):
|
|
|
|
def test_not_dead_when_no_flags(self):
|
|
self.assertFalse(is_remote_link_marked_dead({"url": "https://example.com/m.safetensors"}))
|
|
|
|
def test_dead_when_deletedAt_set(self):
|
|
self.assertTrue(is_remote_link_marked_dead({"deletedAt": "2024-01-01"}))
|
|
|
|
def test_dead_when_deleted_at_set(self):
|
|
self.assertTrue(is_remote_link_marked_dead({"deleted_at": "2024-01-01"}))
|
|
|
|
def test_dead_when_isDead_true(self):
|
|
self.assertTrue(is_remote_link_marked_dead({"isDead": True}))
|
|
|
|
def test_dead_when_is_dead_true(self):
|
|
self.assertTrue(is_remote_link_marked_dead({"is_dead": True}))
|
|
|
|
def test_dead_when_likelyDead_true(self):
|
|
self.assertTrue(is_remote_link_marked_dead({"likelyDead": True}))
|
|
|
|
def test_dead_when_status_dead(self):
|
|
self.assertTrue(is_remote_link_marked_dead({"status": "dead"}))
|
|
|
|
def test_dead_when_status_deleted_case_insensitive(self):
|
|
self.assertTrue(is_remote_link_marked_dead({"status": "Deleted"}))
|
|
|
|
def test_dead_when_status_unavailable(self):
|
|
self.assertTrue(is_remote_link_marked_dead({"status": "unavailable"}))
|
|
|
|
def test_not_dead_for_non_dict(self):
|
|
self.assertFalse(is_remote_link_marked_dead("dead"))
|
|
self.assertFalse(is_remote_link_marked_dead(None))
|
|
self.assertFalse(is_remote_link_marked_dead([]))
|
|
|
|
|
|
class RemoteLinkIsDeadTests(unittest.TestCase):
|
|
def test_common_predicate_matches_all_supported_dead_link_markers(self):
|
|
dead_items = (
|
|
{"deletedAt": "2024-01-01"},
|
|
{"deleted_at": "2024-01-01"},
|
|
{"isDead": True},
|
|
{"is_dead": True},
|
|
{"likelyDead": True},
|
|
{"likely_dead": True},
|
|
{"dead": True},
|
|
{"status": "dead"},
|
|
{"status": "Deleted"},
|
|
{"status": "unavailable"},
|
|
{"status": "missing"},
|
|
)
|
|
for item in dead_items:
|
|
self.assertTrue(is_remote_link_marked_dead(item))
|
|
|
|
def test_common_predicate_keeps_live_and_non_mapping_items_alive(self):
|
|
self.assertFalse(is_remote_link_marked_dead({"url": "https://example.com/model.safetensors"}))
|
|
self.assertFalse(is_remote_link_marked_dead("dead"))
|
|
self.assertFalse(is_remote_link_marked_dead(None))
|
|
self.assertFalse(is_remote_link_marked_dead([]))
|
|
|
|
|
|
# ===========================================================================
|
|
# civarchive - _collect_download_urls_unified (new consolidated helper)
|
|
# ===========================================================================
|
|
class CollectDownloadUrlsUnifiedTests(unittest.TestCase):
|
|
|
|
def _make_file_info(self, mirrors=None, download_url=None,
|
|
download_urls=None, dead=False):
|
|
info = {}
|
|
if mirrors is not None:
|
|
info["mirrors"] = mirrors
|
|
if download_url is not None:
|
|
info["downloadUrl"] = download_url
|
|
if download_urls is not None:
|
|
info["download_urls"] = download_urls
|
|
if dead:
|
|
info["isDead"] = True
|
|
return info
|
|
|
|
def test_extracts_mirror_urls(self):
|
|
fi = self._make_file_info(mirrors=[
|
|
{"url": "https://civarchive.com/api/download/models/1.safetensors"},
|
|
])
|
|
result = _collect_download_urls_unified(fi)
|
|
self.assertIn("https://civarchive.com/api/download/models/1.safetensors", result)
|
|
|
|
def test_skips_dead_mirrors(self):
|
|
fi = self._make_file_info(mirrors=[
|
|
{"url": "https://civarchive.com/api/download/models/1.safetensors", "isDead": True},
|
|
])
|
|
result = _collect_download_urls_unified(fi)
|
|
self.assertEqual(result, [])
|
|
|
|
def test_skip_file_if_dead_returns_empty(self):
|
|
fi = self._make_file_info(
|
|
download_url="https://civarchive.com/api/download/models/1.safetensors",
|
|
dead=True,
|
|
)
|
|
result = _collect_download_urls_unified(fi, skip_file_if_dead=True)
|
|
self.assertEqual(result, [])
|
|
|
|
def test_check_download_urls_list_includes_extra_urls(self):
|
|
fi = self._make_file_info(download_urls=[
|
|
"https://civarchive.com/api/download/models/2.safetensors"
|
|
])
|
|
result = _collect_download_urls_unified(fi, check_download_urls_list=True,
|
|
skip_file_if_dead=True)
|
|
self.assertIn("https://civarchive.com/api/download/models/2.safetensors", result)
|
|
|
|
def test_prioritize_civitai_last_sorts_correctly(self):
|
|
civitai_url = "https://civitai.com/api/download/models/99.safetensors"
|
|
other_url = "https://civarchive.com/api/download/models/77.safetensors"
|
|
fi = self._make_file_info(mirrors=[
|
|
{"url": civitai_url},
|
|
{"url": other_url},
|
|
])
|
|
result = _collect_download_urls_unified(fi, prioritize_civitai_last=True)
|
|
self.assertEqual(result[-1], civitai_url)
|
|
self.assertEqual(result[0], other_url)
|
|
|
|
def test_deduplication(self):
|
|
url = "https://civarchive.com/api/download/models/1.safetensors"
|
|
fi = self._make_file_info(
|
|
mirrors=[{"url": url}],
|
|
download_url=url,
|
|
)
|
|
result = _collect_download_urls_unified(fi)
|
|
self.assertEqual(result.count(url), 1)
|
|
|
|
def test_empty_file_info_returns_empty(self):
|
|
self.assertEqual(_collect_download_urls_unified({}), [])
|
|
|
|
|
|
# ===========================================================================
|
|
# civarchive - archive download URL collection strategy
|
|
# ===========================================================================
|
|
class CollectDownloadUrlsStrategyTests(unittest.TestCase):
|
|
|
|
def test_civitai_url_sorted_last(self):
|
|
civitai = "https://civitai.com/api/download/models/99.safetensors"
|
|
archive = "https://civarchive.com/api/download/models/77.safetensors"
|
|
fi = {"mirrors": [{"url": civitai}, {"url": archive}]}
|
|
result = _collect_download_urls_unified(
|
|
fi,
|
|
prioritize_civitai_last=True,
|
|
download_url_keys=("downloadUrl",),
|
|
)
|
|
self.assertEqual(result[-1], civitai)
|
|
self.assertEqual(result[0], archive)
|
|
|
|
def test_returns_list(self):
|
|
result = _collect_download_urls_unified(
|
|
{},
|
|
prioritize_civitai_last=True,
|
|
download_url_keys=("downloadUrl",),
|
|
)
|
|
self.assertIsInstance(result, list)
|
|
|
|
|
|
# ===========================================================================
|
|
# civarchive - normalized download URL collection strategy
|
|
# ===========================================================================
|
|
class CollectNormalizedDownloadUrlsTests(unittest.TestCase):
|
|
|
|
def test_skips_dead_file(self):
|
|
fi = {
|
|
"isDead": True,
|
|
"downloadUrl": "https://civarchive.com/api/download/models/1.safetensors",
|
|
}
|
|
self.assertEqual(
|
|
_collect_download_urls_unified(
|
|
fi,
|
|
skip_file_if_dead=True,
|
|
check_download_urls_list=True,
|
|
download_url_keys=("download_url", "downloadUrl"),
|
|
),
|
|
[],
|
|
)
|
|
|
|
def test_collects_mirror_and_extra_urls(self):
|
|
fi = {
|
|
"mirrors": [{"url": "https://civarchive.com/api/download/models/1.safetensors"}],
|
|
"download_urls": ["https://civarchive.com/api/download/models/2.safetensors"],
|
|
}
|
|
result = _collect_download_urls_unified(
|
|
fi,
|
|
skip_file_if_dead=True,
|
|
check_download_urls_list=True,
|
|
download_url_keys=("download_url", "downloadUrl"),
|
|
)
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_dead_mirror_url_excluded_from_download_urls(self):
|
|
dead_url = "https://civarchive.com/api/download/models/1.safetensors"
|
|
fi = {
|
|
"mirrors": [{"url": dead_url, "isDead": True}],
|
|
"download_urls": [dead_url],
|
|
}
|
|
result = _collect_download_urls_unified(
|
|
fi,
|
|
skip_file_if_dead=True,
|
|
check_download_urls_list=True,
|
|
download_url_keys=("download_url", "downloadUrl"),
|
|
)
|
|
self.assertNotIn(dead_url, result)
|
|
|
|
|
|
# ===========================================================================
|
|
# civarchive - _normalize_archive_version (calls extract_trained_words)
|
|
# ===========================================================================
|
|
class NormalizeArchiveVersionTests(unittest.TestCase):
|
|
|
|
def _make_version(self, **kwargs):
|
|
return {"id": 1, "name": "v1", **kwargs}
|
|
|
|
def _make_context(self, model_id=100):
|
|
return {"id": model_id, "name": "TestModel", "type": "LORA"}
|
|
|
|
def test_returns_dict_with_expected_keys(self):
|
|
version = self._make_version()
|
|
result = _normalize_archive_version(version, self._make_context())
|
|
for key in ("id", "name", "trained_words", "files", "images", "stats"):
|
|
self.assertIn(key, result)
|
|
|
|
def test_trained_words_from_trainedWords_key(self):
|
|
"""Regression: extract_trained_words must be called correctly (not _extract_trained_words)."""
|
|
version = self._make_version(trainedWords=["lora_trigger", "style_word"])
|
|
result = _normalize_archive_version(version, self._make_context())
|
|
self.assertIn("lora_trigger", result["trained_words"])
|
|
self.assertIn("style_word", result["trained_words"])
|
|
|
|
def test_trained_words_from_trigger_key(self):
|
|
version = self._make_version(trigger="my_trigger_word")
|
|
result = _normalize_archive_version(version, self._make_context())
|
|
self.assertIn("my_trigger_word", result["trained_words"])
|
|
|
|
def test_trained_words_empty_when_not_present(self):
|
|
version = self._make_version()
|
|
result = _normalize_archive_version(version, self._make_context())
|
|
self.assertEqual(result["trained_words"], [])
|
|
|
|
def test_base_model_extracted(self):
|
|
version = self._make_version(baseModel="SDXL 1.0")
|
|
result = _normalize_archive_version(version, self._make_context())
|
|
self.assertEqual(result["base_model"], "SDXL 1.0")
|
|
|
|
def test_images_list_when_no_images(self):
|
|
version = self._make_version()
|
|
result = _normalize_archive_version(version, self._make_context())
|
|
self.assertIsInstance(result["images"], list)
|
|
|
|
def test_files_list_when_no_files(self):
|
|
version = self._make_version()
|
|
result = _normalize_archive_version(version, self._make_context())
|
|
self.assertIsInstance(result["files"], list)
|
|
|
|
|
|
# ===========================================================================
|
|
# civarchive - parse_civarchive_url
|
|
# ===========================================================================
|
|
class ParseCivarchiveUrlTests(unittest.TestCase):
|
|
|
|
def test_model_url_parsed(self):
|
|
result = parse_civarchive_url("https://civarchive.com/models/123")
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.get("model_id"), 123)
|
|
|
|
def test_model_version_url_parsed(self):
|
|
result = parse_civarchive_url("https://civarchive.com/models/123?modelVersionId=456")
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.get("model_id"), 123)
|
|
self.assertEqual(result.get("version_id"), 456)
|
|
|
|
def test_non_civarchive_url_returns_none(self):
|
|
self.assertIsNone(parse_civarchive_url("https://huggingface.co/user/repo"))
|
|
|
|
def test_empty_url_returns_none(self):
|
|
self.assertIsNone(parse_civarchive_url(""))
|
|
|
|
def test_none_returns_none(self):
|
|
self.assertIsNone(parse_civarchive_url(None))
|
|
|
|
|
|
# ===========================================================================
|
|
# civarchive - _prepare_size_probe_url
|
|
# ===========================================================================
|
|
class PrepareSizeProbeUrlTests(unittest.TestCase):
|
|
|
|
def _prepare_size_probe_url(self, url):
|
|
return prepare_remote_size_probe_url(url, ["civarchive.com", "civitai.com", "civitai.red"])
|
|
|
|
def test_civarchive_api_download_url_passes(self):
|
|
url = "https://civarchive.com/api/download/models/77.safetensors"
|
|
self.assertEqual(self._prepare_size_probe_url(url), url)
|
|
|
|
def test_civarchive_non_download_url_blocked(self):
|
|
url = "https://civarchive.com/models/77"
|
|
self.assertIsNone(self._prepare_size_probe_url(url))
|
|
|
|
def test_civitai_non_download_url_blocked(self):
|
|
url = "https://civitai.com/models/123"
|
|
self.assertIsNone(self._prepare_size_probe_url(url))
|
|
|
|
def test_civitai_api_download_url_passes(self):
|
|
url = "https://civitai.com/api/download/models/99.safetensors"
|
|
self.assertEqual(self._prepare_size_probe_url(url), url)
|
|
|
|
def test_none_returns_none(self):
|
|
self.assertIsNone(self._prepare_size_probe_url(None))
|
|
|
|
|
|
class CivArchiveMirrorSizeProbeTests(unittest.TestCase):
|
|
|
|
@patch(
|
|
"core.sources.civarchive.fetch_remote_file_size_cached",
|
|
return_value=364853140,
|
|
)
|
|
def test_huggingface_mirror_is_allowed(self, mock_fetch):
|
|
url = (
|
|
"https://huggingface.co/unsloth/LTX-2.3-GGUF/resolve/main/vae/"
|
|
"ltx-2.3-22b-distilled_audio_vae.safetensors"
|
|
)
|
|
|
|
self.assertEqual(_fetch_remote_file_size_bytes(url), 364853140)
|
|
mock_fetch.assert_called_once()
|
|
self.assertEqual(mock_fetch.call_args.args[0], url)
|
|
|
|
@patch(
|
|
"core.sources.civarchive.fetch_remote_file_size_cached",
|
|
return_value=364853140,
|
|
)
|
|
def test_modelscope_mirror_is_allowed(self, mock_fetch):
|
|
url = (
|
|
"https://www.modelscope.ai/unsloth/LTX-2.3-GGUF/resolve/master/vae/"
|
|
"ltx-2.3-22b-distilled_audio_vae.safetensors"
|
|
)
|
|
|
|
self.assertEqual(_fetch_remote_file_size_bytes(url), 364853140)
|
|
mock_fetch.assert_called_once()
|
|
self.assertEqual(mock_fetch.call_args.args[0], url)
|
|
|
|
@patch("core.sources.civarchive._fetch_remote_file_size_bytes")
|
|
def test_huggingface_is_probed_before_modelscope(self, mock_fetch):
|
|
modelscope_url = "https://www.modelscope.ai/org/repo/resolve/master/model.safetensors"
|
|
huggingface_url = "https://huggingface.co/org/repo/resolve/main/model.safetensors"
|
|
mock_fetch.side_effect = lambda url: 123456 if url == huggingface_url else None
|
|
|
|
size = _resolve_file_size_bytes(
|
|
{},
|
|
[modelscope_url, huggingface_url],
|
|
)
|
|
|
|
self.assertEqual(size, 123456)
|
|
mock_fetch.assert_called_once_with(huggingface_url)
|
|
|
|
@patch("core.sources.civarchive.fetch_remote_file_size_cached")
|
|
def test_untrusted_mirror_is_not_probed(self, mock_fetch):
|
|
self.assertIsNone(
|
|
_fetch_remote_file_size_bytes(
|
|
"https://example.com/models/model.safetensors"
|
|
)
|
|
)
|
|
mock_fetch.assert_not_called()
|
|
|
|
|
|
|
|
# ===========================================================================
|
|
# civitai - parse_civitai_url
|
|
# ===========================================================================
|
|
class ParseCivitaiUrlTests(unittest.TestCase):
|
|
|
|
def test_standard_model_url(self):
|
|
result = parse_civitai_url("https://civitai.com/models/123456")
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.get("model_id"), 123456)
|
|
|
|
def test_model_url_with_version(self):
|
|
result = parse_civitai_url("https://civitai.com/models/123456?modelVersionId=789")
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.get("model_id"), 123456)
|
|
self.assertEqual(result.get("version_id"), 789)
|
|
|
|
def test_civitai_red_model_url_with_version(self):
|
|
result = parse_civitai_url("https://civitai.red/models/123456?modelVersionId=789")
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.get("model_id"), 123456)
|
|
self.assertEqual(result.get("version_id"), 789)
|
|
|
|
def test_civitai_red_download_url(self):
|
|
result = parse_civitai_url("https://civitai.red/api/download/models/789")
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.get("version_id"), 789)
|
|
|
|
def test_non_civitai_url_returns_none(self):
|
|
self.assertIsNone(parse_civitai_url("https://civarchive.com/models/123"))
|
|
|
|
def test_none_returns_none(self):
|
|
self.assertIsNone(parse_civitai_url(None))
|
|
|
|
def test_empty_string_returns_none(self):
|
|
self.assertIsNone(parse_civitai_url(""))
|
|
|
|
|
|
# ===========================================================================
|
|
# civitai - civitai.red HTML fallback
|
|
# ===========================================================================
|
|
class CivitaiRedSearchTests(unittest.TestCase):
|
|
|
|
def test_session_cookie_header_includes_current_and_legacy_names(self):
|
|
cookie = build_civitai_session_cookie("valid-session-token")
|
|
|
|
self.assertIn("__Secure-civ-token=valid-session-token", cookie)
|
|
self.assertIn("__Secure-civitai-token=valid-session-token", cookie)
|
|
|
|
@patch("core.sources.civitai.log")
|
|
@patch("core.sources.civitai.requests.get")
|
|
def test_unauthorized_html_fallback_is_silent_not_warning(self, mock_get, mock_log):
|
|
response = MagicMock()
|
|
response.status_code = 403
|
|
mock_get.return_value = response
|
|
|
|
result = _search_civitai_red_candidates("ae.safetensors")
|
|
|
|
self.assertEqual(result, [])
|
|
mock_log.debug.assert_not_called()
|
|
mock_log.warning.assert_not_called()
|
|
|
|
@patch("core.sources.civitai.requests.get")
|
|
def test_html_fallback_sends_session_token_cookie(self, mock_get):
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
response.text = ""
|
|
mock_get.return_value = response
|
|
|
|
_search_civitai_red_candidates(
|
|
"ae.safetensors",
|
|
session_token="valid-session-token",
|
|
)
|
|
|
|
headers = mock_get.call_args.kwargs.get("headers", {})
|
|
self.assertIn("__Secure-civ-token=valid-session-token", headers.get("Cookie", ""))
|
|
self.assertIn("__Secure-civitai-token=valid-session-token", headers.get("Cookie", ""))
|
|
self.assertIn("user-agent", headers)
|
|
|
|
@patch("core.sources.civitai.requests.get")
|
|
def test_html_fallback_uses_current_search_url_and_model_type(self, mock_get):
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
response.text = ""
|
|
mock_get.return_value = response
|
|
|
|
_search_civitai_red_candidates(
|
|
r"models\diffusion_models\Anima_baseV10.safetensors",
|
|
model_type="diffusion_models",
|
|
)
|
|
|
|
url = mock_get.call_args.args[0]
|
|
self.assertTrue(url.startswith("https://civitai.red/search/models?"))
|
|
self.assertIn("sortBy=models_v9", url)
|
|
self.assertIn("query=Anima_baseV10", url)
|
|
self.assertIn("modelType=Checkpoint", url)
|
|
|
|
|
|
class CivitaiTrpcSearchTests(unittest.TestCase):
|
|
|
|
def test_extracts_candidates_from_devalue_serialized_data(self):
|
|
payload = {
|
|
"result": {
|
|
"data": json.dumps(
|
|
[
|
|
{"items": 1, "nextCursor": -1},
|
|
[2],
|
|
{"id": 3, "version": 4},
|
|
2725430,
|
|
{"id": 5},
|
|
3071760,
|
|
],
|
|
separators=(",", ":"),
|
|
)
|
|
}
|
|
}
|
|
|
|
result = _extract_trpc_model_candidates(payload)
|
|
|
|
self.assertEqual(
|
|
result,
|
|
[{"model_id": 2725430, "version_id": 3071760}],
|
|
)
|
|
|
|
def test_unrecognized_trpc_shapes_return_no_candidates(self):
|
|
payloads = [
|
|
None,
|
|
"not-an-object",
|
|
{"result": "not-an-object"},
|
|
{"result": {"data": "not-json"}},
|
|
{"result": {"data": {"json": "not-json"}}},
|
|
{"result": {"data": {"json": {"items": "not-a-list"}}}},
|
|
]
|
|
|
|
for payload in payloads:
|
|
with self.subTest(payload=payload):
|
|
self.assertEqual(_extract_trpc_model_candidates(payload), [])
|
|
|
|
@patch("core.sources.civitai.requests.get")
|
|
def test_devalue_serialized_empty_items_does_not_abort_search(self, mock_get):
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
response.headers = {"content-type": "application/json"}
|
|
response.text = (
|
|
'{"result":{"data":"[{\\"items\\":1,\\"nextCursor\\":-1},[]]"}}'
|
|
)
|
|
response.json.return_value = {
|
|
"result": {"data": '[{"items":1,"nextCursor":-1},[]]'}
|
|
}
|
|
mock_get.return_value = response
|
|
|
|
result = _search_civitai_trpc_candidates(
|
|
"KNPV3_1.safetensors",
|
|
model_type="loras",
|
|
limit=5,
|
|
)
|
|
|
|
self.assertEqual(result, [])
|
|
mock_get.assert_called_once()
|
|
|
|
@patch("core.sources.civitai.requests.get")
|
|
def test_does_not_retry_without_type_filter_when_typed_search_is_empty(self, mock_get):
|
|
empty_response = MagicMock()
|
|
empty_response.status_code = 200
|
|
empty_response.headers = {"content-type": "application/json"}
|
|
empty_response.text = '{"result":{"data":{"json":{"items":[]}}}}'
|
|
empty_response.json.return_value = {
|
|
"result": {"data": {"json": {"items": []}}}
|
|
}
|
|
|
|
mock_get.return_value = empty_response
|
|
|
|
result = _search_civitai_trpc_candidates(
|
|
"Anima_baseV10.safetensors",
|
|
model_type="diffusion_models",
|
|
limit=5,
|
|
)
|
|
|
|
self.assertEqual(result, [])
|
|
mock_get.assert_called_once()
|
|
url = mock_get.call_args.args[0]
|
|
self.assertIn("Anima_baseV10", url)
|
|
self.assertIn("Checkpoint", url)
|
|
|
|
def test_html_fallback_not_used_to_top_up_trpc_candidates(self):
|
|
clear_search_cache()
|
|
with patch(
|
|
"core.sources.civitai._search_civitai_trpc_candidates",
|
|
return_value=[{"model_id": 10, "version_id": 20}],
|
|
), patch(
|
|
"core.sources.civitai._search_civitai_red_candidates"
|
|
) as html_search, patch(
|
|
"core.sources.civitai._find_civitai_file_in_model",
|
|
return_value=None,
|
|
):
|
|
result = search_civitai_for_file(
|
|
"ae.safetensors",
|
|
model_type="vae",
|
|
candidate_limit=5,
|
|
use_trpc_search=True,
|
|
use_html_fallback=True,
|
|
)
|
|
|
|
self.assertIsNone(result)
|
|
html_search.assert_not_called()
|
|
clear_search_cache()
|
|
|
|
def test_public_api_candidates_used_when_trpc_disabled_and_html_empty(self):
|
|
clear_search_cache()
|
|
expected = {
|
|
"source": "civitai",
|
|
"model_id": 2676616,
|
|
"version_id": 3005748,
|
|
"filename": "sickOllie_v1.safetensors",
|
|
"confidence": 100.0,
|
|
}
|
|
with patch(
|
|
"core.sources.civitai._search_civitai_red_candidates",
|
|
return_value=[],
|
|
) as html_search, patch(
|
|
"core.sources.civitai._search_civitai_public_api_candidates",
|
|
return_value=[{"model_id": 2676616, "version_id": 3005748}],
|
|
) as api_search, patch(
|
|
"core.sources.civitai._find_civitai_file_in_model",
|
|
return_value=expected,
|
|
):
|
|
result = search_civitai_for_file(
|
|
"sickOllie_v1.safetensors",
|
|
model_type="diffusion_models",
|
|
candidate_limit=5,
|
|
use_trpc_search=False,
|
|
use_html_fallback=True,
|
|
)
|
|
|
|
self.assertEqual(result, expected)
|
|
html_search.assert_called_once()
|
|
api_search.assert_called_once()
|
|
clear_search_cache()
|
|
|
|
def test_public_api_candidate_extraction_uses_first_version(self):
|
|
payload = {
|
|
"items": [
|
|
{
|
|
"id": 101,
|
|
"modelVersions": [
|
|
{"id": 202},
|
|
{"id": 303},
|
|
],
|
|
}
|
|
]
|
|
}
|
|
|
|
result = _extract_public_api_model_candidates(payload, limit=5)
|
|
|
|
self.assertEqual(result, [{"model_id": 101, "version_id": 202}])
|
|
|
|
@patch("core.sources.civitai.requests.get")
|
|
def test_public_api_search_sends_auth_and_type_filter(self, mock_get):
|
|
response = MagicMock()
|
|
response.status_code = 200
|
|
response.headers = {"content-type": "application/json"}
|
|
response.text = '{"items":[{"id":101,"modelVersions":[{"id":202}]}]}'
|
|
response.json.return_value = {
|
|
"items": [{"id": 101, "modelVersions": [{"id": 202}]}]
|
|
}
|
|
mock_get.return_value = response
|
|
|
|
result = _search_civitai_public_api_candidates(
|
|
"realistic.safetensors",
|
|
model_type="checkpoints",
|
|
api_key="api-key",
|
|
session_token="session-token",
|
|
limit=5,
|
|
)
|
|
|
|
self.assertEqual(result, [{"model_id": 101, "version_id": 202}])
|
|
params = mock_get.call_args.kwargs.get("params", {})
|
|
headers = mock_get.call_args.kwargs.get("headers", {})
|
|
self.assertEqual(params.get("query"), "realistic")
|
|
self.assertEqual(params.get("types"), "Checkpoint")
|
|
self.assertEqual(headers.get("Authorization"), "Bearer api-key")
|
|
self.assertIn("__Secure-civitai-token=session-token", headers.get("Cookie", ""))
|
|
|
|
def test_public_api_search_can_be_disabled(self):
|
|
clear_search_cache()
|
|
with patch(
|
|
"core.sources.civitai._search_civitai_trpc_candidates",
|
|
return_value=[],
|
|
), patch(
|
|
"core.sources.civitai._search_civitai_red_candidates",
|
|
return_value=[],
|
|
), patch(
|
|
"core.sources.civitai._search_civitai_public_api_candidates",
|
|
return_value=[{"model_id": 2676616, "version_id": 3005748}],
|
|
) as api_search:
|
|
result = search_civitai_for_file(
|
|
"sickOllie_v1.safetensors",
|
|
candidate_limit=5,
|
|
use_trpc_search=True,
|
|
use_api_search=False,
|
|
use_html_fallback=True,
|
|
)
|
|
|
|
self.assertIsNone(result)
|
|
api_search.assert_not_called()
|
|
clear_search_cache()
|
|
|
|
|
|
# ===========================================================================
|
|
# civitai - _extract_model_images
|
|
# ===========================================================================
|
|
class ExtractModelImagesTests(unittest.TestCase):
|
|
|
|
def test_returns_list_for_empty_input(self):
|
|
self.assertIsInstance(_extract_model_images({}), list)
|
|
|
|
def test_extracts_images_list(self):
|
|
version_info = {
|
|
"images": [
|
|
{"url": "https://image.civitai.com/img1.jpg", "nsfw": False},
|
|
{"url": "https://image.civitai.com/img2.jpg", "nsfw": True},
|
|
]
|
|
}
|
|
result = _extract_model_images(version_info)
|
|
self.assertIsInstance(result, list)
|
|
self.assertGreater(len(result), 0)
|
|
|
|
def test_all_results_have_url(self):
|
|
version_info = {
|
|
"images": [
|
|
{"url": "https://image.civitai.com/img1.jpg"},
|
|
]
|
|
}
|
|
result = _extract_model_images(version_info)
|
|
for item in result:
|
|
self.assertIn("url", item)
|
|
|
|
def test_no_images_key_returns_empty(self):
|
|
result = _extract_model_images({"name": "v1"})
|
|
self.assertEqual(result, [])
|
|
|
|
def test_adds_linked_parent_model_without_replacing_additional_resources(self):
|
|
version_info = {
|
|
"id": 3084537,
|
|
"name": "Krea2 (v0.5)",
|
|
"modelId": 2268008,
|
|
"model": {
|
|
"name": "Realistic Snapshot (Z-Image-Turbo + Krea 2)",
|
|
"type": "LORA",
|
|
},
|
|
"images": [
|
|
{
|
|
"url": "https://image.civitai.com/135642695.jpeg",
|
|
"meta": {
|
|
"additionalResources": [
|
|
{
|
|
"name": (
|
|
"dwdwadawdadw__realisticsnapshotideogram__"
|
|
"5th%20training6000.safetensors"
|
|
),
|
|
"type": "lora",
|
|
"strength": 2,
|
|
},
|
|
{
|
|
"name": "lenovo_krea2.safetensors",
|
|
"type": "lora",
|
|
"strength": 0.5,
|
|
}
|
|
]
|
|
},
|
|
}
|
|
],
|
|
}
|
|
|
|
result = _extract_model_images(version_info)
|
|
|
|
resources = result[0]["resources"]
|
|
self.assertEqual(
|
|
resources[0],
|
|
{
|
|
"name": "Realistic Snapshot (Z-Image-Turbo + Krea 2)",
|
|
"versionName": "Krea2 (v0.5)",
|
|
"type": "LORA",
|
|
"modelId": 2268008,
|
|
"modelVersionId": 3084537,
|
|
"url": (
|
|
"https://civitai.com/models/2268008"
|
|
"?modelVersionId=3084537"
|
|
),
|
|
"primary": True,
|
|
"strength": 2,
|
|
},
|
|
)
|
|
self.assertEqual(
|
|
resources[1],
|
|
{
|
|
"name": "lenovo_krea2.safetensors",
|
|
"type": "lora",
|
|
"strength": 0.5,
|
|
},
|
|
)
|
|
|
|
def test_does_not_duplicate_an_existing_parent_model_resource(self):
|
|
version_info = {
|
|
"id": 3084537,
|
|
"name": "Krea2 (v0.5)",
|
|
"modelId": 2268008,
|
|
"model": {"name": "Realistic Snapshot", "type": "LORA"},
|
|
"images": [
|
|
{
|
|
"url": "https://image.civitai.com/135642695.jpeg",
|
|
"resources": [
|
|
{
|
|
"name": "Realistic Snapshot",
|
|
"modelId": 2268008,
|
|
"modelVersionId": 3084537,
|
|
}
|
|
],
|
|
}
|
|
],
|
|
}
|
|
|
|
result = _extract_model_images(version_info)
|
|
|
|
self.assertEqual(len(result[0]["resources"]), 1)
|
|
|
|
|
|
# ===========================================================================
|
|
# civitai - _normalize_civitai_file
|
|
# ===========================================================================
|
|
class NormalizeCivitaiFileTests(unittest.TestCase):
|
|
|
|
def _make_file(self, **kwargs):
|
|
return {
|
|
"id": 1,
|
|
"name": "model.safetensors",
|
|
"downloadUrl": "https://civitai.com/api/download/models/1",
|
|
"sizeKB": 2048,
|
|
"hashes": {"SHA256": "abc123def456"},
|
|
"primary": True,
|
|
**kwargs,
|
|
}
|
|
|
|
def test_returns_dict_with_expected_keys(self):
|
|
result = _normalize_civitai_file(self._make_file(), model_id=10, version_id=20)
|
|
for key in ("name", "download_url", "sha256", "primary", "size"):
|
|
self.assertIn(key, result)
|
|
|
|
def test_sha256_extracted_from_hashes(self):
|
|
result = _normalize_civitai_file(
|
|
self._make_file(hashes={"SHA256": "deadbeef1234"}),
|
|
model_id=10, version_id=20,
|
|
)
|
|
self.assertEqual(result["sha256"], "deadbeef1234")
|
|
|
|
def test_primary_flag_preserved(self):
|
|
result = _normalize_civitai_file(self._make_file(primary=True), model_id=10, version_id=20)
|
|
self.assertTrue(result["primary"])
|
|
result2 = _normalize_civitai_file(self._make_file(primary=False), model_id=10, version_id=20)
|
|
self.assertFalse(result2["primary"])
|
|
|
|
def test_empty_file_info_does_not_raise_nameerror(self):
|
|
try:
|
|
result = _normalize_civitai_file({}, model_id=10, version_id=20)
|
|
self.assertIsInstance(result, dict)
|
|
except (TypeError, AttributeError, KeyError):
|
|
pass # acceptable - no silent NameError allowed
|
|
|
|
|
|
class CivitaiFilenameMatchTests(unittest.TestCase):
|
|
|
|
def test_short_filename_base_does_not_partial_match_longer_name(self):
|
|
versions = [
|
|
{
|
|
"id": 1,
|
|
"files": [
|
|
{"name": "titanae-akima.safetensors"},
|
|
],
|
|
}
|
|
]
|
|
|
|
result = _find_matching_file_in_versions(versions, "ae.safetensors")
|
|
|
|
self.assertIsNone(result)
|
|
|
|
def test_short_filename_base_still_matches_exact_filename(self):
|
|
versions = [
|
|
{
|
|
"id": 1,
|
|
"files": [
|
|
{"name": "ae.safetensors"},
|
|
],
|
|
}
|
|
]
|
|
|
|
result = _find_matching_file_in_versions(versions, "ae.safetensors")
|
|
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result["match_type"], "exact")
|
|
|
|
|
|
# ===========================================================================
|
|
# civitai - hash lookup details enrichment
|
|
# ===========================================================================
|
|
class EnrichModelInfoWithDetailsTests(unittest.TestCase):
|
|
|
|
def test_selected_version_enrichment_does_not_raise_nameerror(self):
|
|
details = {
|
|
"selected_version": {
|
|
"description": "version description",
|
|
"url": "https://civitai.com/models/10?modelVersionId=20",
|
|
},
|
|
"name": "Resolved Model",
|
|
"type": "Checkpoint",
|
|
"description": "model description",
|
|
"tags": ["tag-a"],
|
|
"images": [{"url": "https://image.civitai.com/example.jpeg"}],
|
|
"version_url": "https://civitai.com/models/10?modelVersionId=20",
|
|
"url": "https://civitai.com/models/10",
|
|
}
|
|
|
|
with patch("core.sources.civitai.get_civitai_model_details", return_value=details):
|
|
result = _enrich_model_info_with_details(
|
|
{"source": "civitai", "sha256": "a" * 64},
|
|
model_id=10,
|
|
version_id=20,
|
|
)
|
|
|
|
self.assertEqual(result["description"], "version description")
|
|
self.assertEqual(result["model_name"], "Resolved Model")
|
|
self.assertEqual(result["model_type"], "Checkpoint")
|
|
self.assertEqual(result["tags"], ["tag-a"])
|
|
self.assertEqual(result["images"], details["images"])
|
|
|
|
|
|
# ===========================================================================
|
|
# workflow_updater - convert_to_relative_path
|
|
# ===========================================================================
|
|
class ConvertToRelativePathTests(unittest.TestCase):
|
|
|
|
def test_returns_filename_when_in_base_dir(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
model_path = os.path.join(tmpdir, "model.safetensors")
|
|
result = convert_to_relative_path(model_path, "checkpoints", tmpdir)
|
|
self.assertEqual(result, "model.safetensors")
|
|
|
|
def test_returns_relative_path_with_subfolder(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
subdir = os.path.join(tmpdir, "sd15")
|
|
os.makedirs(subdir)
|
|
model_path = os.path.join(subdir, "model.safetensors")
|
|
result = convert_to_relative_path(model_path, "checkpoints", tmpdir)
|
|
self.assertIn("sd15", result)
|
|
self.assertIn("model.safetensors", result)
|
|
|
|
def test_uses_forward_slashes(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
subdir = os.path.join(tmpdir, "sub")
|
|
os.makedirs(subdir)
|
|
model_path = os.path.join(subdir, "model.safetensors")
|
|
result = convert_to_relative_path(model_path, "loras", tmpdir)
|
|
# Result should use forward slashes (platform-normalised)
|
|
result_normalised = result.replace("\\", "/")
|
|
self.assertIn("sub", result_normalised)
|
|
self.assertIn("model.safetensors", result_normalised)
|
|
|
|
|
|
# ===========================================================================
|
|
# workflow_updater - get_base_directory_for_model
|
|
# ===========================================================================
|
|
class GetBaseDirectoryForModelTests(unittest.TestCase):
|
|
|
|
def test_returns_base_directory_if_provided(self):
|
|
model = {"base_directory": "/some/path/models"}
|
|
result = get_base_directory_for_model(model, "checkpoints")
|
|
self.assertEqual(result, "/some/path/models")
|
|
|
|
def test_returns_none_without_path_or_base_directory(self):
|
|
result = get_base_directory_for_model({}, "checkpoints")
|
|
self.assertIsNone(result)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|