refactor: centralize CivitAI enum fetching
This commit is contained in:
+16
-2
@@ -267,6 +267,20 @@ def _collect_normalized_base_model_tokens(base_models: List[Dict[str, Any]]) ->
|
||||
return normalized_tokens
|
||||
|
||||
|
||||
def _fetch_civitai_enums(*, raise_on_error: bool = False):
|
||||
"""Fetch CivitAI enum data using the shared source request configuration."""
|
||||
request_kwargs = {
|
||||
"timeout": 15,
|
||||
"log_name": "CivitAI Enums",
|
||||
}
|
||||
if raise_on_error:
|
||||
request_kwargs["raise_on_error"] = True
|
||||
return request_source_json(
|
||||
"https://civitai.com/api/v1/enums",
|
||||
**request_kwargs,
|
||||
)
|
||||
|
||||
|
||||
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()
|
||||
@@ -286,7 +300,7 @@ def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]:
|
||||
if check_remote:
|
||||
try:
|
||||
# Fetch from CivitAI API using unified helper
|
||||
enums = request_source_json("https://civitai.com/api/v1/enums", timeout=15, log_name="CivitAI Enums")
|
||||
enums = _fetch_civitai_enums()
|
||||
if enums:
|
||||
remote_models = enums.get("BaseModel", [])
|
||||
|
||||
@@ -313,7 +327,7 @@ def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]:
|
||||
|
||||
def update_base_models_from_remote() -> Dict[str, Any]:
|
||||
"""Fetch live BaseModel enums from CivitAI and merge into base-models.json."""
|
||||
enums = request_source_json("https://civitai.com/api/v1/enums", timeout=15, log_name="CivitAI Enums", raise_on_error=True)
|
||||
enums = _fetch_civitai_enums(raise_on_error=True)
|
||||
if not enums:
|
||||
raise ValueError("CivitAI enums did not return a valid response")
|
||||
remote_names = enums.get("BaseModel", [])
|
||||
|
||||
@@ -189,6 +189,35 @@ class TestRefactoringUnification(unittest.TestCase):
|
||||
)
|
||||
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):
|
||||
@@ -436,4 +465,3 @@ class TestRefactoringUnification(unittest.TestCase):
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user