diff --git a/core/sources/popular.py b/core/sources/popular.py index c26b96c..bb2116b 100644 --- a/core/sources/popular.py +++ b/core/sources/popular.py @@ -6,7 +6,7 @@ Curated list of common models with known download URLs. import os import re -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Set from ..log_system import create_module_logger @@ -257,6 +257,16 @@ def generate_aliases(name: str) -> List[str]: return sorted(list({a for a in aliases if len(a) > 1})) +def _collect_normalized_base_model_tokens(base_models: List[Dict[str, Any]]) -> Set[str]: + """Collect normalized names and aliases from the base-model catalog.""" + normalized_tokens = set() + for model in base_models: + normalized_tokens.add(normalize_alphanumeric_key(model.get("name", ""))) + for alias in model.get("aliases", []): + normalized_tokens.add(normalize_alphanumeric_key(alias)) + return normalized_tokens + + def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]: """Return local base-models metadata and optionally compare with CivitAI.""" local_data = _read_base_models_file() @@ -281,11 +291,7 @@ def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]: remote_models = enums.get("BaseModel", []) # Check if there are new models by comparing normalized names and aliases - all_existing_normalized = set() - for m in base_models: - all_existing_normalized.add(normalize_alphanumeric_key(m.get("name", ""))) - for alias in m.get("aliases", []): - all_existing_normalized.add(normalize_alphanumeric_key(alias)) + all_existing_normalized = _collect_normalized_base_model_tokens(base_models) new_models_found = False for remote_name in remote_models: @@ -318,11 +324,7 @@ def update_base_models_from_remote() -> Dict[str, Any]: base_models = local_data.get("base_models", []) # 1. Build a set of all normalized aliases currently in base-models.json - all_known_normalized = set() - for m in base_models: - for alias in m.get("aliases", []): - all_known_normalized.add(normalize_alphanumeric_key(alias)) - all_known_normalized.add(normalize_alphanumeric_key(m.get("name", ""))) + all_known_normalized = _collect_normalized_base_model_tokens(base_models) updated_models = list(base_models) new_added_count = 0 diff --git a/tests/test_refactoring_unification.py b/tests/test_refactoring_unification.py index 105a974..77c8ff6 100644 --- a/tests/test_refactoring_unification.py +++ b/tests/test_refactoring_unification.py @@ -151,6 +151,44 @@ class TestRefactoringUnification(unittest.TestCase): 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.base_models_mgr") @patch("requests.get") def test_update_base_models_from_remote_success(self, mock_get, mock_base_models_mgr):