import ast import os import sys import unittest from pathlib import Path from unittest.mock import MagicMock, call, patch parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if parent_dir not in sys.path: sys.path.insert(0, parent_dir) from core.contracts import SearchResult from core.matcher import MODEL_TITLE_MATCH_THRESHOLD from core.type_utils import ( build_model_result, normalize_lora_manager_type, normalize_model_file_info, ) class TestRefactoringUnification(unittest.TestCase): def tearDown(self): from core.sources.popular import reload_databases reload_databases() def test_thresholds(self): self.assertEqual(MODEL_TITLE_MATCH_THRESHOLD, 82.0) def test_invalidate_model_caches_clears_inventory_then_hash_matches(self): from core import resolver calls = MagicMock() with ( patch.object( resolver, "invalidate_model_files_cache" ) as invalidate_inventory, patch.object( resolver, "invalidate_local_hash_match_cache" ) as invalidate_hash_matches, ): calls.attach_mock(invalidate_inventory, "inventory") calls.attach_mock(invalidate_hash_matches, "hash_matches") resolver.invalidate_model_caches() calls.assert_has_calls([call.inventory(), call.hash_matches()]) invalidate_inventory.assert_called_once_with() invalidate_hash_matches.assert_called_once_with() def test_normalize_lora_manager_type(self): self.assertEqual(normalize_lora_manager_type("loras"), "lora") self.assertEqual(normalize_lora_manager_type("checkpoints"), "checkpoint") self.assertEqual(normalize_lora_manager_type("embedding"), "textualinversion") self.assertEqual(normalize_lora_manager_type("unknown"), "unknown") def test_normalize_model_file_info(self): raw_file = { "id": 1, "name": "test.safetensors", "sizeKB": 100, "primary": True, "hashes": {"sha256": "abc"}, } res = normalize_model_file_info(raw_file, model_id=10, version_id=20) self.assertEqual(res["id"], 1) self.assertEqual(res["name"], "test.safetensors") self.assertEqual(res["size"], 100 * 1024) self.assertTrue(res["primary"]) self.assertEqual(res["sha256"], "abc") self.assertEqual(res["model_id"], 10) self.assertEqual(res["version_id"], 20) def test_build_model_result(self): res = build_model_result( source="civitai", model_id=100, version_id=200, name="Test Model", filename="test.safetensors", extra_field="hello_world", ) self.assertIsInstance(res, SearchResult) serialized = res.to_dict() self.assertEqual(serialized["source"], "civitai") self.assertEqual(serialized["model_id"], 100) self.assertEqual(serialized["version_id"], 200) self.assertEqual(serialized["name"], "Test Model") self.assertEqual(serialized["filename"], "test.safetensors") self.assertEqual(serialized["extra_field"], "hello_world") def test_build_model_result_custom_url_contract(self): res = build_model_result( "civitai", model_id=100, version_id=200, name="Test Model", filename="test.safetensors", url="https://civitai.com/models/100?modelVersionId=200", download_url="https://civitai.com/api/download/models/200", match_type="custom_url", details_source="civitai", version_url="https://civitai.com/models/100?modelVersionId=200", custom_url=True, result_mode="custom_url", ) serialized = res.to_dict() self.assertEqual(serialized["details_source"], "civitai") self.assertEqual(serialized["version_url"], serialized["url"]) self.assertTrue(serialized["custom_url"]) self.assertNotIn("confidence", serialized) def test_build_model_result_compact_custom_url_contract(self): res = build_model_result( "civitai", name="CivitAI", filename="test.safetensors", url="https://civitai.com/api/download/models/200", download_url="https://civitai.com/api/download/models/200", details_source="civitai", version_url="https://civitai.com/api/download/models/200", match_type="custom_url", custom_url=True, result_mode="compact_custom_url", ) self.assertEqual({ "source": "civitai", "details_source": "civitai", "name": "CivitAI", "filename": "test.safetensors", "url": "https://civitai.com/api/download/models/200", "version_url": "https://civitai.com/api/download/models/200", "download_url": "https://civitai.com/api/download/models/200", "match_type": "custom_url", "custom_url": True, }, res.to_dict()) @patch("requests.get") def test_request_page_text_success(self, mock_get): from core.sources.civarchive import _request_page_text mock_response = MagicMock() mock_response.status_code = 200 mock_response.text = "content" mock_get.return_value = mock_response res = _request_page_text("test_path") self.assertEqual(res, "content") mock_get.assert_called_once() @patch("requests.get") def test_request_page_text_fail(self, mock_get): from core.sources.civarchive import _request_page_text mock_response = MagicMock() mock_response.status_code = 404 mock_get.return_value = mock_response res = _request_page_text("test_path") self.assertIsNone(res) mock_get.side_effect = Exception("network error") res2 = _request_page_text("test_path") self.assertIsNone(res2) @patch("requests.get") def test_get_base_models_status_remote(self, mock_get): from core.sources.popular import get_base_models_status mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = {"BaseModel": ["SD 1.5", "SDXL 1.0"]} mock_get.return_value = mock_response res = get_base_models_status(check_remote=True) self.assertIn("local_count", res) self.assertIn("update_available", res) mock_get.assert_called_once_with("https://civitai.com/api/v1/enums", params={}, headers=None, timeout=15) @patch("core.sources.popular.request_source_json") @patch("core.sources.popular._read_base_models_file") @patch("core.sources.popular._read_base_models_meta") @patch("core.sources.popular.base_models_mgr") @patch("core.sources.popular.reload_databases") def test_base_model_alias_matching_is_consistent_for_status_and_update( self, mock_reload_databases, mock_base_models_mgr, mock_read_meta, mock_read_file, mock_request_source_json, ): from core.sources.popular import get_base_models_status, update_base_models_from_remote local_data = { "base_models": [{"name": "Known Model", "aliases": ["Known Alias"]}] } mock_read_file.return_value = local_data mock_read_meta.return_value = {} mock_request_source_json.return_value = { "BaseModel": ["Known Alias", "New Model"] } status = get_base_models_status(check_remote=True) self.assertTrue(status["update_available"]) result = update_base_models_from_remote() self.assertEqual(result["new_models_added"], 1) self.assertEqual(result["new_models_added_list"], ["New Model"]) saved_data = mock_base_models_mgr.sync_catalog.call_args[0][0] self.assertEqual( [model["name"] for model in saved_data["base_models"]], ["Known Model", "New Model"], ) mock_reload_databases.assert_called_once() @patch("core.sources.popular.request_source_json") @patch("core.sources.popular._read_base_models_file") @patch("core.sources.popular._read_base_models_meta") def test_base_model_remote_payload_keeps_status_tolerant_and_update_strict( self, mock_read_meta, mock_read_file, mock_request_source_json, ): from core.sources.popular import ( get_base_models_status, update_base_models_from_remote, ) mock_read_file.return_value = {"base_models": []} mock_read_meta.return_value = {} mock_request_source_json.return_value = {"unexpected": []} status = get_base_models_status(check_remote=True) self.assertFalse(status["update_available"]) self.assertTrue(status["remote_checked_at"]) with self.assertRaisesRegex(ValueError, "valid BaseModel list"): update_base_models_from_remote() self.assertEqual(mock_request_source_json.call_count, 2) self.assertNotIn("raise_on_error", mock_request_source_json.call_args_list[0].kwargs) self.assertTrue(mock_request_source_json.call_args_list[1].kwargs["raise_on_error"]) @patch("core.sources.popular.base_models_mgr") @patch("requests.get") def test_update_base_models_from_remote_success(self, mock_get, mock_base_models_mgr): from core.sources.popular import update_base_models_from_remote mock_response = MagicMock() mock_response.status_code = 200 mock_response.json.return_value = {"BaseModel": ["SD 1.5", "SDXL 1.0", "NewModel"]} mock_get.return_value = mock_response # Mock reading base-models.json mock_base_models_mgr.read_data.return_value = { "base_models": [ {"name": "SD 1.5", "aliases": ["sd1.5"]}, {"name": "SDXL 1.0", "aliases": ["sdxl"]}, ] } mock_base_models_mgr.read_meta.return_value = {} res = update_base_models_from_remote() self.assertTrue(res.get("updated")) self.assertEqual(res.get("new_models_added"), 1) self.assertIn("NewModel", res.get("new_models_added_list", [])) # Verify it saved the updated catalog mock_base_models_mgr.sync_catalog.assert_called_once() saved_data = mock_base_models_mgr.sync_catalog.call_args[0][0] self.assertEqual(len(saved_data["base_models"]), 3) @patch("core.sources.model_list.catalog_mgr") @patch("core.sources.model_list._get_remote_model_list_info") @patch("core.sources.model_list._fetch_json_url") def test_model_list_lifecycle_status_and_update(self, mock_fetch_json, mock_get_remote, mock_catalog_mgr): from core.sources.model_list import get_model_list_update_status, update_model_list_from_remote # Mock responses mock_catalog_mgr.read_data.side_effect = [ {"models": [{"name": "ModelA"}]}, # Data file read {"models": [{"name": "ModelA"}]}, # Reset update {"models": [{"name": "ModelA"}, {"name": "ModelB"}]} # updated check ] mock_catalog_mgr.read_meta.side_effect = [ {"sha": "123", "updated_at": "2026"}, # Meta file read {"sha": "123"}, # Reset update {"sha": "456"} ] mock_get_remote.return_value = { "sha": "123", "size": 100, "download_url": "https://dummy", "html_url": "https://dummy_html" } status = get_model_list_update_status(check_remote=True) self.assertEqual(status["local_count"], 1) self.assertTrue(status["can_compare"]) # Reset mock mock_get_remote.return_value = { "sha": "456", "size": 100, "download_url": "https://dummy", "html_url": "https://dummy_html" } mock_fetch_json.return_value = {"models": [{"name": "ModelA"}, {"name": "ModelB"}]} res = update_model_list_from_remote() self.assertTrue(res.get("updated")) self.assertTrue(mock_catalog_mgr.sync_catalog.called) def test_tracker_progress_updates(self): from core.progress import JobProgressTracker tracker = JobProgressTracker("Test Tracker") tracker.update("job123", status="running", percent=50) state = tracker.get("job123") self.assertEqual(state["status"], "running") self.assertEqual(state["percent"], 50) def test_get_filename_from_path(self): from core.path_utils import get_filename_from_path self.assertEqual(get_filename_from_path("some/path/to/model.safetensors"), "model.safetensors") self.assertEqual(get_filename_from_path("some\\path\\to\\model.safetensors"), "model.safetensors") self.assertEqual(get_filename_from_path("model.safetensors"), "model.safetensors") self.assertEqual(get_filename_from_path(""), "") self.assertEqual(get_filename_from_path(None), "") def test_alphanumeric_normalizers(self): from core.matcher import normalize_base_model from core.type_utils import normalize_alphanumeric_key inputs = ["SD 1.5", "sd-xl_1.0!!", "Flux.1-dev", "", None] for val in inputs: res_matcher = normalize_base_model(val) res_key = normalize_alphanumeric_key(val) # verify the public domain helper uses the canonical implementation self.assertEqual(res_matcher, res_key) # verify correct regex behavior import re expected = re.sub(r"[^a-z0-9]+", "", str(val or "").lower()) self.assertEqual(res_matcher, expected) def test_retired_module_aliases_are_absent(self): project_root = Path(__file__).resolve().parents[1] retired_aliases = { "core/type_utils.py": "normalize_alphanumeric_lower", "core/matcher.py": "MODEL_FILE_EXTENSIONS", } for relative_path, alias_name in retired_aliases.items(): source = (project_root / relative_path).read_text(encoding="utf-8") tree = ast.parse(source, filename=relative_path) module_aliases = { target.id for statement in tree.body if isinstance(statement, ast.Assign) for target in statement.targets if isinstance(target, ast.Name) and isinstance(statement.value, ast.Name) } self.assertNotIn( alias_name, module_aliases, f"Retired module alias {alias_name} was reintroduced in {relative_path}", ) def test_utc_now_iso(self): from core.sources.model_list import _utc_now_iso res = _utc_now_iso() self.assertTrue(isinstance(res, str)) self.assertEqual(len(res), 25) # YYYY-MM-DDTHH:MM:SS+00:00 self.assertIn("T", res) self.assertTrue(res.endswith("+00:00")) # Check no microsecond self.assertNotIn(".", res) def test_listification(self): from core.settings import _listify_tags from core.type_utils import as_list self.assertEqual(as_list("a,b, c"), ["a", "b", "c"]) self.assertEqual(list(_listify_tags("a,b; c")), ["a", "b", "c"]) self.assertEqual(as_list(["a", "", None, "b"]), ["a", "b"]) self.assertEqual(list(_listify_tags(["a", "", None, "b"])), ["a", "b"]) def test_build_model_result_normalizes_hashes(self): from core.type_utils import normalize_hashes_dict valid_sha = "a" * 64 norm_hashes = normalize_hashes_dict({"sha256": valid_sha, "autoV2": "1234567890"}) self.assertEqual(norm_hashes["sha256"], valid_sha.lower()) self.assertEqual(norm_hashes["autoV2"], "1234567890") res = build_model_result( "civitai", model_id=123, version_id=456, name="Test Model", hashes={"sha256": valid_sha}, download_url="https://civitai.com/api/download/123", normalize_hashes=True, ) expected = build_model_result( "civitai", model_id=123, version_id=456, name="Test Model", download_url="https://civitai.com/api/download/123", sha256=valid_sha.lower(), hashes={"sha256": valid_sha.lower()}, ) self.assertEqual(expected, res) serialized = res.to_dict() self.assertEqual(serialized["source"], "civitai") self.assertEqual(serialized["model_id"], 123) self.assertEqual(serialized["version_id"], 456) self.assertEqual(serialized["sha256"], valid_sha.lower()) self.assertEqual(serialized["hashes"]["sha256"], valid_sha.lower()) def test_job_progress_tracker_update_from_payload(self): from core.progress import JobProgressTracker tracker = JobProgressTracker("Test Job") tracker.update_from_payload("job99", {"stage": "processing", "current": 5, "total": 10, "message": "Working..."}) data = tracker.get("job99") self.assertEqual(data["status"], "running") self.assertEqual(data["stage"], "processing") self.assertEqual(data["current"], 5) self.assertEqual(data["message"], "Working...") tracker.update_from_payload("job99", {"stage": "completed", "message": "Done!"}) data_done = tracker.get("job99") self.assertEqual(data_done["status"], "completed") @patch("core.network_utils.request_source_json") def test_fetch_json_from_public_url(self, mock_request_source_json): from core.network_utils import fetch_json_from_public_url mock_request_source_json.return_value = {"key": "value"} data = fetch_json_from_public_url("https://example.com/api.json") self.assertEqual(data, {"key": "value"}) mock_request_source_json.assert_called_once() def test_prepare_remote_size_probe_url_with_allowed_domains(self): from core.type_utils import prepare_remote_size_probe_url # Test default behavior self.assertEqual( prepare_remote_size_probe_url("https://huggingface.co/user/repo/blob/main/model.safetensors"), "https://huggingface.co/user/repo/resolve/main/model.safetensors" ) # Test allowed domains filtering self.assertEqual( prepare_remote_size_probe_url("https://civitai.com/api/download/models/123", ["civitai.com", "civitai.red"]), "https://civitai.com/api/download/models/123" ) self.assertIsNone( prepare_remote_size_probe_url("https://otherdomain.com/model.safetensors", ["civitai.com", "civitai.red"]) ) @patch("core.network_utils.request_source_json") def test_execute_provider_json_request(self, mock_request_source_json): from core.network_utils import execute_provider_json_request mock_request_source_json.return_value = {"response": "ok"} # Test Bearer Authorization and Cookie header injection res = execute_provider_json_request( "Test Provider", "https://example.com/api", api_key="my_secret_token", session_token="my_session" ) self.assertEqual(res, {"response": "ok"}) mock_request_source_json.assert_called_once() call_args = mock_request_source_json.call_args headers = call_args[1]["headers"] self.assertEqual(headers["Authorization"], "Bearer my_secret_token") self.assertIn("__Secure-civ-token=my_session", headers["Cookie"]) if __name__ == "__main__": unittest.main()