""" 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 from core.contracts import ResolvedModel, SearchResult # --------------------------------------------------------------------------- # civarchive internals # --------------------------------------------------------------------------- from core.sources.civarchive import ( CIVARCHIVE_SIZE_PROBE_DOMAINS, REQUEST_HEADERS, _build_result_from_normalized_version, _build_result_from_payload, _collect_archive_download_urls, _extract_archive_size_fields, _merge_top_file_mirrors, _mirror_from_top_file, _normalize_archive_version, _normalize_download_url, _resolve_file_size_bytes, _transform_file_entry, build_civarchive_custom_result, get_civarchive_model_details, parse_civarchive_url, search_civarchive, search_civarchive_for_file, ) # --------------------------------------------------------------------------- # civitai internals # --------------------------------------------------------------------------- from core.sources.civitai import ( _build_civitai_result_from_version, _build_civitai_title_search_queries, _enrich_model_info_with_details, _extract_model_images, _extract_public_api_model_candidates, _extract_trpc_model_candidates, _find_civitai_file_in_model, _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_result_payload, build_civitai_session_cookie, check_civitai_api_key, check_civitai_session_token, clear_search_cache, get_civitai_model_details, get_model_info_by_hash, get_model_info_for_file, parse_civitai_url, resolve_civitai_version_custom_result, resolve_urn, search_civitai, search_civitai_for_file, ) from core.sources.common import ( build_custom_result_fields, is_remote_link_marked_dead, probe_remote_file_size, resolve_file_size, ) 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" class CivArchiveFileSizeAliasTests(unittest.TestCase): def test_extract_archive_size_fields_preserves_alias_precedence(self): self.assertEqual( { "sizeKB": 12, "sizeBytes": 34, "size": 56, }, _extract_archive_size_fields( { "sizeKB": 0, "size_kb": 12, "sizeBytes": 0, "size_bytes": 0, "fileSize": 34, "size": 56, } ), ) def test_all_archive_file_projections_use_the_same_size_aliases(self): file_data = { "id": 7, "name": "model.safetensors", "url": "https://civarchive.com/api/download/models/7", "size_kb": 12, "file_size": 34, } mirror = _mirror_from_top_file(file_data) transformed = _transform_file_entry(file_data) merged = _merge_top_file_mirrors([], [file_data]) for projection in (mirror, transformed, merged[0]): with self.subTest(projection=projection): self.assertEqual(12, projection["sizeKB"]) self.assertEqual(34, projection["sizeBytes"]) 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_payload_preserves_optional_search_fields(self): payload = build_civitai_result_payload( model_id=123, version_id=456, model_name="Example model", model_type="Checkpoint", file_info={"name": "example.safetensors", "hashes": {"SHA256": "abc123"}}, filename="example.safetensors", download_url="https://example.test/download", size=2048, base_model="SDXL", tags=["style"], match_type="exact", version_name="v1", confidence=100.0, ) self.assertEqual("Example model", payload["name"]) self.assertEqual("v1", payload["version_name"]) self.assertEqual("example.safetensors", payload["filename"]) self.assertEqual("abc123", payload["sha256"]) self.assertEqual(100.0, payload["confidence"]) 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.assertIsInstance(result, SearchResult) result = result.to_dict() 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 CivitaiCustomVersionLookupTests(unittest.TestCase): def test_version_lookup_uses_shared_json_request_contract(self): with ( patch( "core.sources.civitai.execute_provider_json_request", return_value={"modelId": 123}, ) as request_json, patch( "core.sources.civitai.get_civitai_model_details", return_value={"model_id": 123}, ) as get_details, patch( "core.sources.civitai.build_civitai_custom_result", return_value=SearchResult( source="civitai", version_id=456, ), ), ): result = resolve_civitai_version_custom_result( 456, expected_filename="model.safetensors", api_key="secret", ) self.assertIsInstance(result, SearchResult) self.assertEqual("civitai", result.source) self.assertEqual(456, result.version_id) request_json.assert_called_once_with( "CivitAI custom URL version lookup", "https://civitai.com/api/v1/model-versions/456", api_key="secret", timeout=20, max_attempts=1, raise_on_error=False, ) get_details.assert_called_once_with(123, 456, "secret") def test_version_lookup_returns_none_without_json_payload(self): with ( patch( "core.sources.civitai.execute_provider_json_request", return_value=None, ) as request_json, patch("core.sources.civitai.get_civitai_model_details") as get_details, ): result = resolve_civitai_version_custom_result(456) self.assertIsNone(result) request_json.assert_called_once() get_details.assert_not_called() def test_version_lookup_handles_adapter_exception(self): with patch( "core.sources.civitai.execute_provider_json_request", side_effect=RuntimeError("network failure"), ): result = resolve_civitai_version_custom_result(456) self.assertIsNone(result) class CivarchiveResultBuilderTests(unittest.TestCase): def test_name_search_sorts_typed_results(self): lower_result = SearchResult( source="civarchive", name="typed-civarchive-lower", filename="typed-civarchive-lower.safetensors", confidence=55.0, match_type="similar", ) higher_result = SearchResult( source="civarchive", name="typed-civarchive-higher", filename="typed-civarchive-higher.safetensors", confidence=90.0, match_type="exact", ) with ( patch( "core.sources.civarchive._search_page", return_value=[ {"id": 1, "url": "/models/1", "name": "lower"}, {"id": 2, "url": "/models/2", "name": "higher"}, ], ), patch( "core.sources.civarchive._resolve_search_candidate", side_effect=[lower_result, higher_result], ), ): results = search_civarchive( "typed-civarchive-sort-regression", limit=2, ) self.assertEqual( [higher_result, lower_result], results, ) def test_details_fallback_serializes_typed_version_result(self): resolved = SearchResult( source="civarchive", model_id=123, version_id=456, name="Typed fallback model", version_name="v1", model_type="LORA", filename="typed-fallback.safetensors", download_url="https://civarchive.com/download/456", base_model="SDXL", ) with ( patch( "core.sources.civarchive._request_model_payload", return_value=None, ), patch( "core.sources.civarchive.resolve_civarchive_model_version", return_value=resolved, ), patch( "core.sources.civarchive._normalize_archive_version", return_value={ "id": 456, "url": "https://civarchive.com/models/123?modelVersionId=456", }, ), ): result = get_civarchive_model_details(123, version_id=456) self.assertEqual("civarchive", result["source"]) self.assertEqual(456, result["version_id"]) self.assertEqual("Typed fallback model", result["name"]) def test_filename_search_keeps_typed_resolved_result(self): filename = "typed-civarchive-model.safetensors" resolved = SearchResult( source="civarchive", name="typed-civarchive-model", version_name="v1", filename=filename, download_url="https://civarchive.com/api/download/models/456", ) with ( patch( "core.sources.civarchive._search_page", return_value=[ { "name": "typed-civarchive-model", "url": "/models/123?modelVersionId=456", } ], ), patch( "core.sources.civarchive._resolve_search_candidate", return_value=resolved, ), ): result = search_civarchive_for_file(filename, limit=1) self.assertIsInstance(result, SearchResult) self.assertEqual(filename, result.filename) self.assertEqual(100.0, result.confidence) self.assertEqual("exact", result.match_type) def test_normalized_version_builder_preserves_search_result_contract(self): result = _build_result_from_normalized_version( model_details={ "model_id": 123, "name": "Archive model", "type": "LORA", "tags": ["style"], "creator": {"username": "archive-user"}, "platform": "civarchive", }, version={ "id": 456, "name": "v1", "base_model": "SDXL", "trained_words": ["trigger"], "images": [], "url": "https://civarchive.com/models/123?modelVersionId=456", }, file_info={ "name": "archive.safetensors", "download_url": "https://civarchive.com/api/download/models/456.safetensors", "size": 2048, }, match_type="similar", ) self.assertIsInstance(result, SearchResult) result = result.to_dict() self.assertEqual("civarchive", result["source"]) self.assertEqual(123, result["model_id"]) self.assertEqual(456, result["version_id"]) self.assertEqual("archive.safetensors", result["filename"]) self.assertEqual(2048, result["size"]) self.assertEqual("SDXL", result["base_model"]) self.assertEqual("archive-user", result["creator"]["username"]) self.assertFalse(result["is_deleted"]) self.assertEqual(0.0, result["confidence"]) self.assertNotIn("hash", result) def test_payload_builder_preserves_confidence_hash_and_deleted_fields(self): sha256 = "a" * 64 result = _build_result_from_payload( { "id": 123, "name": "Archive model", "type": "LORA", "creator_username": "archive-user", "creator_name": "Archive User", "creator_url": "https://civarchive.com/users/archive-user", "deletedAt": "2025-01-01T00:00:00Z", "tags": ["style"], "version": { "id": 456, "name": "v1", "baseModel": "SDXL", "trainedWords": ["trigger"], "files": [ { "name": "archive.safetensors", "downloadUrl": "https://civarchive.com/api/download/models/456.safetensors", "sizeBytes": 4096, "hashes": {"SHA256": sha256}, } ], }, }, query="archive.safetensors", preferred_filename="archive.safetensors", ) self.assertIsInstance(result, SearchResult) result = result.to_dict() self.assertEqual("civarchive", result["source"]) self.assertEqual(123, result["model_id"]) self.assertEqual(456, result["version_id"]) self.assertEqual(4096, result["size"]) self.assertEqual(100.0, result["confidence"]) self.assertEqual("exact", result["match_type"]) self.assertEqual(sha256, result["sha256"]) self.assertEqual(sha256, result["hash"]) self.assertEqual(sha256, result["hashes"]["sha256"]) self.assertTrue(result["is_deleted"]) self.assertEqual("archive-user", result["creator"]["username"]) class CivitaiCredentialCheckTests(unittest.TestCase): def test_empty_credentials_keep_provider_specific_precheck_messages(self): self.assertEqual( { "success": False, "valid": False, "status": "missing", "message": "Paste a CivitAI session token first.", }, check_civitai_session_token(" "), ) self.assertEqual( { "success": False, "valid": False, "status": "missing", "message": "Paste a CivitAI API key first.", }, check_civitai_api_key(" "), ) @patch("core.sources.civitai.check_credential_http") def test_session_token_uses_cookie_transport_and_shared_result(self, mock_check): expected = {"success": True, "valid": True, "status": "valid", "username": "tester"} mock_check.return_value = expected result = check_civitai_session_token(" session-value ") self.assertIs(expected, result) mock_check.assert_called_once() args, kwargs = mock_check.call_args self.assertEqual("https://civitai.com/api/v1/me", args[0]) self.assertEqual( { "accept": "application/json", "Cookie": "__Secure-civ-token=session-value; __Secure-civitai-token=session-value", "user-agent": kwargs["headers"]["user-agent"], }, kwargs["headers"], ) self.assertEqual("Session token is valid.", kwargs["success_message"]) self.assertEqual("Session token is not accepted by CivitAI.", kwargs["error_msg_401_403"]) @patch("core.sources.civitai.check_credential_http") def test_api_key_uses_bearer_transport_and_shared_result(self, mock_check): expected = {"success": False, "valid": False, "status": "invalid"} mock_check.return_value = expected result = check_civitai_api_key(" api-key-value ") self.assertIs(expected, result) mock_check.assert_called_once() args, kwargs = mock_check.call_args self.assertEqual("https://civitai.com/api/v1/me", args[0]) self.assertEqual("application/json", kwargs["headers"]["accept"]) self.assertEqual("Bearer api-key-value", kwargs["headers"]["Authorization"]) self.assertNotIn("Cookie", kwargs["headers"]) self.assertEqual("CivitAI API key is valid.", kwargs["success_message"]) self.assertEqual("CivitAI API key is not accepted.", kwargs["error_msg_401_403"]) class CivitaiPrimaryFileSelectionTests(unittest.TestCase): @patch("core.sources.civitai.execute_provider_json_request") def test_resolve_urn_preserves_first_marked_file_order(self, mock_request): clear_search_cache() mock_request.return_value = { "name": "Example model", "tags": [], "modelVersions": [ { "id": 456, "name": "v1", "baseModel": "SDXL", "files": [ {"name": "first.safetensors", "type": "Model", "sizeKB": 1}, { "name": "primary.safetensors", "primary": True, "type": "Model", "sizeKB": 2, }, ], } ], } result = resolve_urn(123, 456) self.assertEqual("first.safetensors", result["expected_filename"]) self.assertEqual(1024, result["files"][0]["size"]) clear_search_cache() @patch("core.sources.civitai.execute_provider_json_request") def test_general_search_preserves_first_marked_file_order(self, mock_request): clear_search_cache() mock_request.return_value = { "items": [ { "id": 123, "name": "Example model", "type": "Checkpoint", "stats": {"downloadCount": 1}, "modelVersions": [ { "id": 456, "baseModel": "SDXL", "files": [ {"name": "first.safetensors", "type": "Model", "sizeKB": 1}, { "name": "primary.safetensors", "primary": True, "type": "Model", "sizeKB": 2, }, ], } ], } ] } results = search_civitai("example", limit=1) self.assertIsInstance(results[0], SearchResult) self.assertEqual("first.safetensors", results[0].filename) self.assertEqual(1024, results[0].size) clear_search_cache() class CustomUrlResultBuilderTests(unittest.TestCase): def test_build_custom_result_fields_preserves_shared_result_contract(self): result = build_custom_result_fields( source="civitai", details={ "model_id": 123, "name": "Example model", "type": "LORA", "tags": ["style"], "images": [{"url": "https://example.test/image"}], "description": "Model description", }, selected_version={ "id": 456, "name": "v1", "base_model": "SD 1.5", "trained_words": ["example"], "description": "Version notes", }, file_info={ "name": "example.safetensors", "type": "Model", "size": 123, "sha256": "abc123", "hashes": {"SHA256": "abc123"}, }, filename="example.safetensors", download_url="https://civitai.com/api/download/models/456", version_url="https://civitai.com/models/123?modelVersionId=456", ) self.assertEqual("civitai", result["source"]) self.assertEqual(123, result["model_id"]) self.assertEqual(456, result["version_id"]) self.assertEqual("Version notes", result["description"]) self.assertEqual("abc123", result["sha256"]) self.assertEqual({"SHA256": "abc123"}, result["hashes"]) self.assertEqual("custom_url", result["result_mode"]) 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.to_dict()) 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.to_dict()) 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_archive_download_urls (common collector adapter) # =========================================================================== class CollectArchiveDownloadUrlsTests(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_archive_download_urls(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_archive_download_urls(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_archive_download_urls(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_archive_download_urls(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_archive_download_urls(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_archive_download_urls(fi) self.assertEqual(result.count(url), 1) def test_empty_file_info_returns_empty(self): self.assertEqual(_collect_archive_download_urls({}), []) # =========================================================================== # civarchive - archive download URL collection strategy # =========================================================================== class CollectArchiveDownloadUrlsStrategyTests(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_archive_download_urls( 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_archive_download_urls( {}, prioritize_civitai_last=True, download_url_keys=("downloadUrl",), ) self.assertIsInstance(result, list) # =========================================================================== # civarchive - normalized download URL collection strategy # =========================================================================== class CollectNormalizedArchiveDownloadUrlsTests(unittest.TestCase): def test_skips_dead_file(self): fi = { "isDead": True, "downloadUrl": "https://civarchive.com/api/download/models/1.safetensors", } self.assertEqual( _collect_archive_download_urls( 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_archive_download_urls( 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_archive_download_urls( 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.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.model_id, 123) self.assertEqual(result.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): def test_shared_file_size_resolver_prefers_metadata(self): probe = MagicMock(return_value=654321) size = resolve_file_size( {"size": 123456}, ["https://example.com/model.safetensors"], probe=probe, ) self.assertEqual(size, 123456) probe.assert_not_called() def test_shared_file_size_resolver_uses_ordered_probe_until_success(self): probe = MagicMock(side_effect=[None, 654321]) urls = [ "https://example.com/first.safetensors", "https://example.com/second.safetensors", "https://example.com/third.safetensors", ] size = resolve_file_size({}, urls, probe=probe) self.assertEqual(size, 654321) probe.assert_has_calls([ ((urls[0],), {}), ((urls[1],), {}), ]) self.assertEqual(probe.call_count, 2) def test_metadata_size_is_preferred_over_remote_probe(self): with patch("core.sources.civarchive.probe_remote_file_size") as mock_fetch: size = _resolve_file_size_bytes( {"size": 123456}, ["https://huggingface.co/org/repo/resolve/main/model.safetensors"], ) self.assertEqual(size, 123456) mock_fetch.assert_not_called() @patch( "core.sources.civarchive.probe_remote_file_size", return_value=654321, ) def test_zero_metadata_size_falls_back_to_remote_probe(self, mock_fetch): url = "https://huggingface.co/org/repo/resolve/main/model.safetensors" size = _resolve_file_size_bytes({"size": 0}, [url]) self.assertEqual(size, 654321) mock_fetch.assert_called_once_with( url, url_normalizer=_normalize_download_url, allowed_domains=CIVARCHIVE_SIZE_PROBE_DOMAINS, headers=REQUEST_HEADERS, ) @patch( "core.sources.common.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( probe_remote_file_size( url, allowed_domains=CIVARCHIVE_SIZE_PROBE_DOMAINS, headers=REQUEST_HEADERS, ), 364853140, ) mock_fetch.assert_called_once() self.assertEqual(mock_fetch.call_args.args[0], url) @patch( "core.sources.common.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( probe_remote_file_size( url, allowed_domains=CIVARCHIVE_SIZE_PROBE_DOMAINS, headers=REQUEST_HEADERS, ), 364853140, ) mock_fetch.assert_called_once() self.assertEqual(mock_fetch.call_args.args[0], url) @patch("core.sources.civarchive.probe_remote_file_size") 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, **kwargs: 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, url_normalizer=_normalize_download_url, allowed_domains=CIVARCHIVE_SIZE_PROBE_DOMAINS, headers=REQUEST_HEADERS, ) @patch("core.sources.common.fetch_remote_file_size_cached") def test_untrusted_mirror_is_not_probed(self, mock_fetch): self.assertIsNone( probe_remote_file_size( "https://example.com/models/model.safetensors", allowed_domains=CIVARCHIVE_SIZE_PROBE_DOMAINS, headers=REQUEST_HEADERS, ) ) 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.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.model_id, 123456) self.assertEqual(result.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.model_id, 123456) self.assertEqual(result.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.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 = SearchResult( 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.assertIsInstance(result, SearchResult) self.assertEqual(expected.source, result.source) self.assertEqual(expected.model_id, result.model_id) self.assertEqual(expected.version_id, result.version_id) self.assertEqual(expected.filename, result.filename) self.assertEqual(expected.confidence, result.confidence) 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_generic_prefix_does_not_match_longer_model_filename(self): versions = [ { "id": 1, "files": [ {"name": "text.safetensors"}, ], } ] result = _find_matching_file_in_versions( versions, "Text Refusal Reduction Krea2.safetensors", ) self.assertIsNone(result) 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") @patch("core.sources.civitai.resolve_urn") @patch("core.sources.civitai.execute_provider_json_request", return_value=None) def test_resolved_meaningful_partial_filename_keeps_partial_type( self, mock_request, mock_resolve_urn, ): mock_resolve_urn.return_value = { "model_name": "Example model", "version_name": "v1", "expected_filename": "model.safetensors", "base_model": None, "tags": [], "files": [{"name": "model.safetensors", "size": 1024}], } result = _find_civitai_file_in_model( model_id=1, filename="model_v2.safetensors", preferred_version_id=2, ) self.assertIsNotNone(result) self.assertEqual("partial", result.match_type) @patch("core.sources.civitai.resolve_urn") @patch("core.sources.civitai.execute_provider_json_request") def test_resolved_partial_filename_is_not_returned_as_exact( self, mock_request, mock_resolve_urn, ): clear_search_cache() target_filename = "Text Refusal Reduction Krea2.safetensors" resolved = { "model_name": "Text Aid SDXL", "version_name": "v1", "expected_filename": "text.safetensors", "base_model": "SDXL 1.0", "tags": [], "files": [{"name": "text.safetensors", "size": 1024}], } mock_resolve_urn.return_value = resolved mock_request.return_value = { "name": "Text Aid SDXL", "type": "LORA", "tags": [], "modelVersions": [ { "id": 467355, "files": [{"name": "text.safetensors"}], } ], } result = _find_civitai_file_in_model( model_id=419492, filename=target_filename, preferred_version_id=467355, base_model_context="SDXL 1.0", ) self.assertIsNone(result) clear_search_cache() class CivitaiTitleFallbackTests(unittest.TestCase): def test_title_queries_reorder_filename_and_add_lora_type(self): queries = _build_civitai_title_search_queries( "Text Refusal Reduction Krea2.safetensors", model_type="loras", ) self.assertEqual( [ "Krea2 Refusal Reduction LoRA", "Krea2 Refusal Reduction", "Krea2 Text Refusal Reduction LoRA", "Krea2 Text Refusal Reduction", "Text Refusal Reduction Krea2", ], queries, ) @patch("core.sources.civitai._search_civitai_trpc_candidates") @patch("core.sources.civitai._find_civitai_file_in_model") def test_title_fallback_finds_model_when_filename_differs( self, mock_find_model, mock_search_candidates, ): clear_search_cache() expected = SearchResult( source="civitai", model_id=2775340, version_id=3125118, name="Krea2 TextFusion Refusal-Reduction LoRA", filename="Krea2_TextFusion_Refusal_Reduction.safetensors", match_type="model_title", confidence=83.6, ) def search_candidates(query, **kwargs): if query == "Krea2 Refusal Reduction LoRA": return [{"model_id": 2775340, "version_id": 3125118}] return [] def find_model(model_id, filename, **kwargs): if filename == "Krea2 Refusal Reduction LoRA": return expected return None mock_search_candidates.side_effect = search_candidates mock_find_model.side_effect = find_model result = search_civitai_for_file( "Text Refusal Reduction Krea2.safetensors", model_type="loras", use_api_search=False, use_html_fallback=False, ) self.assertIsNotNone(result) self.assertEqual(2775340, result.model_id) self.assertEqual("model_title", result.match_type) self.assertTrue(result.extra_value("title_fallback")) self.assertEqual( "Krea2 Refusal Reduction LoRA", result.extra_value("title_query"), ) clear_search_cache() # =========================================================================== # 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 = ResolvedModel(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(ResolvedModel(), "checkpoints") self.assertIsNone(result) if __name__ == "__main__": unittest.main() class CivitaiRequestOptimizationTests(unittest.TestCase): @patch("core.sources.civitai.resolve_urn") @patch("core.sources.civitai.execute_provider_json_request") def test_reuses_model_data_across_queries_and_skips_unrelated_versions(self, request, resolve): request.return_value = { "name": "Unrelated", "modelVersions": [{"id": n, "files": [{"name": "zzzzzz.safetensors"}]} for n in range(20)], } cache = {} for query in ("target_model.safetensors", "target model", "target model LoRA"): _find_civitai_file_in_model(1, query, model_data_cache=cache, title_only=True) request.assert_called_once() resolve.assert_not_called() @patch("core.sources.civitai.resolve_urn") @patch("core.sources.civitai.execute_provider_json_request") def test_finds_non_primary_file_without_version_requests(self, request, resolve): request.return_value = { "name": "Example", "type": "LORA", "modelVersions": [{"id": 2, "baseModel": "Krea 2", "files": [ {"name": "other.safetensors", "primary": True}, {"name": "target.safetensors", "sizeKB": 12, "hashes": {"SHA256": "a" * 64}}, ]}], } result = _find_civitai_file_in_model(1, "target.safetensors", preferred_version_id=2, exact_only=True) self.assertEqual(result.filename, "target.safetensors") self.assertEqual(result.match_type, "exact") request.assert_called_once() resolve.assert_not_called() @patch("core.sources.civitai.resolve_urn") @patch("core.sources.civitai.execute_provider_json_request") def test_unrelated_complete_file_lists_need_no_version_requests(self, request, resolve): request.return_value = { "name": "Unrelated", "modelVersions": [{"id": n, "files": [{"name": "zzzzzz.safetensors"}]} for n in range(20)], } result = _find_civitai_file_in_model(1, "target.safetensors", exact_only=True) self.assertIsNone(result) resolve.assert_not_called() @patch("core.sources.civitai.resolve_urn") @patch("core.sources.civitai.execute_provider_json_request") def test_missing_file_list_falls_back_to_version_details(self, request, resolve): request.return_value = {"name": "Example", "modelVersions": [{"id": 2}]} resolve.return_value = { "expected_filename": "target.safetensors", "files": [{"name": "target.safetensors", "size": 12}], } result = _find_civitai_file_in_model(1, "target.safetensors", exact_only=True) self.assertEqual(result.filename, "target.safetensors") resolve.assert_called_once_with(1, 2, None)