diff --git a/core/routes/search_support.py b/core/routes/search_support.py index 4ba06e8..e2266fd 100644 --- a/core/routes/search_support.py +++ b/core/routes/search_support.py @@ -1,117 +1,51 @@ """Search cache, credentials, and index route registration.""" +from ..services.search_support_service import SearchSupportService from .context import RouteContext def register_search_support_routes(context: RouteContext): - asyncio = context.get('asyncio') - check_brave_search_api_key = context.get('check_brave_search_api_key') - check_civitai_api_key = context.get('check_civitai_api_key') - check_civitai_session_token = context.get('check_civitai_session_token') - check_huggingface_token = context.get('check_huggingface_token') - clear_all_search_caches = context.get('clear_all_search_caches') - clear_huggingface_search_cache = context.get('clear_huggingface_search_cache') - get_known_author_fallback_indexes_status = context.get('get_known_author_fallback_indexes_status') - get_model_list_update_status = context.get('get_model_list_update_status') - invalidate_local_hash_match_cache = context.get('invalidate_local_hash_match_cache') - invalidate_model_files_cache = context.get('invalidate_model_files_cache') - json_api_endpoint = context.get('json_api_endpoint') - refresh_known_author_fallback_indexes = context.get('refresh_known_author_fallback_indexes') - reload_model_list = context.get('reload_model_list') - reload_popular_databases = context.get('reload_popular_databases') - routes = context.get('routes') - self = context.get('self') - update_model_list_from_remote = context.get('update_model_list_from_remote') - web = context.get('web') + json_api_endpoint = context.require("json_api_endpoint") + routes = context.require("routes") + service = SearchSupportService(context) @routes.post("/model_resolver/clear-search-cache") @json_api_endpoint("Clear search cache") async def clear_search_cache_route(request): - """Clear backend search caches after token/settings changes.""" - clear_all_search_caches() - reload_popular_databases() - reload_model_list() - invalidate_model_files_cache() - invalidate_local_hash_match_cache() - self.search_result_timestamps.clear() - self.logger.info("Cleared backend search caches") - return web.json_response({"success": True, "cleared": "all"}) - - async def _check_credential_helper(request, payload_key, check_func, log_name): - try: - data = await request.json() - val = data.get(payload_key, "") - result = await asyncio.to_thread(check_func, val) - return web.json_response(result) - except Exception as e: - self.logger.exception(f"{log_name} check error: {e}") - return web.json_response({"error": str(e)}, status=500) + return await service.clear_search_cache(request) @routes.post("/model_resolver/civitai/session-token/check") async def civitai_session_token_check_route(request): - """Check whether a CivitAI browser session token is valid.""" - return await _check_credential_helper( - request, "civitai_session_token", check_civitai_session_token, "CivitAI session token" - ) + return await service.check_civitai_session_token_route(request) @routes.post("/model_resolver/civitai/api-key/check") async def civitai_api_key_check_route(request): - """Check whether a CivitAI API key is valid.""" - return await _check_credential_helper( - request, "civitai_key", check_civitai_api_key, "CivitAI API key" - ) + return await service.check_civitai_api_key_route(request) @routes.post("/model_resolver/huggingface/token/check") async def huggingface_token_check_route(request): - """Check whether a HuggingFace token is valid.""" - return await _check_credential_helper( - request, "hf_token", check_huggingface_token, "HuggingFace token" - ) + return await service.check_huggingface_token_route(request) @routes.post("/model_resolver/brave/api-key/check") async def brave_api_key_check_route(request): - """Check whether a Brave Search API key is valid.""" - return await _check_credential_helper( - request, "brave_search_api_key", check_brave_search_api_key, "Brave Search API key" - ) + return await service.check_brave_search_api_key_route(request) @routes.get("/model_resolver/huggingface/author-index/status") @json_api_endpoint("HuggingFace author index status") async def huggingface_author_index_status_route(request): - """Return local HuggingFace author fallback index status.""" - return web.json_response( - get_known_author_fallback_indexes_status() - ) + return await service.author_index_status(request) @routes.post("/model_resolver/huggingface/author-index/refresh") @json_api_endpoint("HuggingFace author index refresh") async def huggingface_author_index_refresh_route(request): - """Refresh HuggingFace author fallback index.""" - data = await request.json() - hf_token = data.get("hf_token", "") - result = await asyncio.to_thread( - refresh_known_author_fallback_indexes, hf_token or None - ) - clear_huggingface_search_cache() - return web.json_response(result) + return await service.refresh_author_index(request) @routes.get("/model_resolver/model-list/status") @json_api_endpoint("Model list status") async def model_list_status_route(request): - """Return local model-list status and optionally compare with GitHub.""" - check_remote = ( - str(request.query.get("check_remote", "")).lower() - in {"1", "true", "yes"} - ) - return web.json_response( - get_model_list_update_status(check_remote=check_remote) - ) + return await service.model_list_status(request) @routes.post("/model_resolver/model-list/update") @json_api_endpoint("Model list update") async def model_list_update_route(request): - """Download latest ComfyUI-Manager model-list.json.""" - result = await asyncio.to_thread(update_model_list_from_remote) - clear_all_search_caches() - self.search_result_timestamps.clear() - return web.json_response(result) + return await service.update_model_list(request) diff --git a/core/services/search_support_service.py b/core/services/search_support_service.py new file mode 100644 index 0000000..a7787de --- /dev/null +++ b/core/services/search_support_service.py @@ -0,0 +1,142 @@ +"""Search cache, credentials, and index operations for HTTP routes.""" + +from ..routes.context import RouteContext + + +class SearchSupportService: + """Coordinate search support operations independently of route decorators.""" + + def __init__(self, context: RouteContext): + self.asyncio = context.require("asyncio") + self.check_brave_search_api_key = context.require( + "check_brave_search_api_key" + ) + self.check_civitai_api_key = context.require("check_civitai_api_key") + self.check_civitai_session_token = context.require( + "check_civitai_session_token" + ) + self.check_huggingface_token = context.require("check_huggingface_token") + self.clear_all_search_caches = context.require("clear_all_search_caches") + self.clear_huggingface_search_cache = context.require( + "clear_huggingface_search_cache" + ) + self.get_known_author_fallback_indexes_status = context.require( + "get_known_author_fallback_indexes_status" + ) + self.get_model_list_update_status = context.require( + "get_model_list_update_status" + ) + self.invalidate_local_hash_match_cache = context.require( + "invalidate_local_hash_match_cache" + ) + self.invalidate_model_files_cache = context.require( + "invalidate_model_files_cache" + ) + self.refresh_known_author_fallback_indexes = context.require( + "refresh_known_author_fallback_indexes" + ) + self.reload_model_list = context.require("reload_model_list") + self.reload_popular_databases = context.require("reload_popular_databases") + self.extension = context.require("self") + self.update_model_list_from_remote = context.require( + "update_model_list_from_remote" + ) + self.web = context.require("web") + + async def clear_search_cache(self, request): + """Clear backend search caches after token/settings changes.""" + self.clear_all_search_caches() + self.reload_popular_databases() + self.reload_model_list() + self.invalidate_model_files_cache() + self.invalidate_local_hash_match_cache() + self.extension.search_result_timestamps.clear() + self.extension.logger.info("Cleared backend search caches") + return self.web.json_response({"success": True, "cleared": "all"}) + + async def _check_credential_helper( + self, + request, + payload_key, + check_func, + log_name, + ): + try: + data = await request.json() + value = data.get(payload_key, "") + result = await self.asyncio.to_thread(check_func, value) + return self.web.json_response(result) + except Exception as error: + self.extension.logger.exception(f"{log_name} check error: {error}") + return self.web.json_response({"error": str(error)}, status=500) + + async def check_civitai_session_token_route(self, request): + """Check whether a CivitAI browser session token is valid.""" + return await self._check_credential_helper( + request, + "civitai_session_token", + self.check_civitai_session_token, + "CivitAI session token", + ) + + async def check_civitai_api_key_route(self, request): + """Check whether a CivitAI API key is valid.""" + return await self._check_credential_helper( + request, + "civitai_key", + self.check_civitai_api_key, + "CivitAI API key", + ) + + async def check_huggingface_token_route(self, request): + """Check whether a HuggingFace token is valid.""" + return await self._check_credential_helper( + request, + "hf_token", + self.check_huggingface_token, + "HuggingFace token", + ) + + async def check_brave_search_api_key_route(self, request): + """Check whether a Brave Search API key is valid.""" + return await self._check_credential_helper( + request, + "brave_search_api_key", + self.check_brave_search_api_key, + "Brave Search API key", + ) + + async def author_index_status(self, request): + """Return local HuggingFace author fallback index status.""" + return self.web.json_response( + self.get_known_author_fallback_indexes_status() + ) + + async def refresh_author_index(self, request): + """Refresh HuggingFace author fallback index.""" + data = await request.json() + hf_token = data.get("hf_token", "") + result = await self.asyncio.to_thread( + self.refresh_known_author_fallback_indexes, + hf_token or None, + ) + self.clear_huggingface_search_cache() + return self.web.json_response(result) + + async def model_list_status(self, request): + """Return local model-list status and optionally compare with GitHub.""" + check_remote = str(request.query.get("check_remote", "")).lower() in { + "1", + "true", + "yes", + } + return self.web.json_response( + self.get_model_list_update_status(check_remote=check_remote) + ) + + async def update_model_list(self, request): + """Download latest ComfyUI-Manager model-list.json.""" + result = await self.asyncio.to_thread(self.update_model_list_from_remote) + self.clear_all_search_caches() + self.extension.search_result_timestamps.clear() + return self.web.json_response(result) diff --git a/tests/test_route_context_contract.py b/tests/test_route_context_contract.py index 010c64a..3092b5a 100644 --- a/tests/test_route_context_contract.py +++ b/tests/test_route_context_contract.py @@ -7,6 +7,7 @@ from core.services.hash_service import HashService from core.services.loaded_models_service import LoadedModelsService from core.services.metadata_service import MetadataService from core.services.scanner_service import ScannerService +from core.services.search_support_service import SearchSupportService from core.services.workflow_service import WorkflowService @@ -78,3 +79,8 @@ def test_download_service_requires_mandatory_route_dependencies(): match="Missing route dependency: Aria2InstallError", ): DownloadService(RouteContext({})) + + +def test_search_support_service_requires_mandatory_route_dependencies(): + with pytest.raises(KeyError, match="Missing route dependency: asyncio"): + SearchSupportService(RouteContext({})) diff --git a/tests/test_search_support_routes.py b/tests/test_search_support_routes.py new file mode 100644 index 0000000..beb82e1 --- /dev/null +++ b/tests/test_search_support_routes.py @@ -0,0 +1,210 @@ +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from aiohttp import web + +from core.routes.context import RouteContext +from core.routes.search_support import register_search_support_routes + + +class _Routes: + def __init__(self): + self.handlers = {} + + def get(self, path): + def decorator(handler): + self.handlers[("GET", path)] = handler + return handler + + return decorator + + def post(self, path): + def decorator(handler): + self.handlers[("POST", path)] = handler + return handler + + return decorator + + +def _json_api_endpoint(_name, **_kwargs): + return lambda handler: handler + + +def _request(payload=None, query=None): + return SimpleNamespace( + json=AsyncMock(return_value=payload), + query=query or {}, + ) + + +def _response_json(response): + return json.loads(response.text) + + +def _build_routes(overrides=None): + routes = _Routes() + extension = SimpleNamespace( + logger=MagicMock(), + search_result_timestamps={"old": "timestamp"}, + ) + values = { + "asyncio": asyncio, + "check_brave_search_api_key": MagicMock( + return_value={"valid": True, "provider": "brave"} + ), + "check_civitai_api_key": MagicMock( + return_value={"valid": True, "provider": "civitai"} + ), + "check_civitai_session_token": MagicMock( + return_value={"valid": True, "provider": "civitai-session"} + ), + "check_huggingface_token": MagicMock( + return_value={"valid": True, "provider": "huggingface"} + ), + "clear_all_search_caches": MagicMock(), + "clear_huggingface_search_cache": MagicMock(), + "get_known_author_fallback_indexes_status": MagicMock( + return_value={"ready": True} + ), + "get_model_list_update_status": MagicMock( + return_value={"loaded": True} + ), + "invalidate_local_hash_match_cache": MagicMock(), + "invalidate_model_files_cache": MagicMock(), + "json_api_endpoint": _json_api_endpoint, + "refresh_known_author_fallback_indexes": MagicMock( + return_value={"refreshed": True} + ), + "reload_model_list": MagicMock(), + "reload_popular_databases": MagicMock(), + "routes": routes, + "self": extension, + "update_model_list_from_remote": MagicMock( + return_value={"updated": True} + ), + "web": web, + } + if overrides: + values.update(overrides) + register_search_support_routes(RouteContext(values)) + return routes.handlers, values + + +@pytest.mark.asyncio +async def test_clear_search_cache_route_clears_all_state(): + handlers, values = _build_routes() + + response = await handlers[("POST", "/model_resolver/clear-search-cache")]( + _request() + ) + + assert response.status == 200 + assert _response_json(response) == {"success": True, "cleared": "all"} + values["clear_all_search_caches"].assert_called_once_with() + values["reload_popular_databases"].assert_called_once_with() + values["reload_model_list"].assert_called_once_with() + values["invalidate_model_files_cache"].assert_called_once_with() + values["invalidate_local_hash_match_cache"].assert_called_once_with() + assert values["self"].search_result_timestamps == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("path", "payload_key", "value", "checker"), + [ + ( + "/model_resolver/civitai/session-token/check", + "civitai_session_token", + "session-token", + "check_civitai_session_token", + ), + ( + "/model_resolver/civitai/api-key/check", + "civitai_key", + "api-key", + "check_civitai_api_key", + ), + ( + "/model_resolver/huggingface/token/check", + "hf_token", + "hf-token", + "check_huggingface_token", + ), + ( + "/model_resolver/brave/api-key/check", + "brave_search_api_key", + "brave-key", + "check_brave_search_api_key", + ), + ], +) +async def test_credential_routes_validate_credentials( + path, payload_key, value, checker +): + handlers, values = _build_routes() + + response = await handlers[("POST", path)]( + _request({payload_key: value}) + ) + + assert response.status == 200 + assert _response_json(response) == values[checker].return_value + values[checker].assert_called_once_with(value) + + +@pytest.mark.asyncio +async def test_credential_route_returns_error_when_checker_fails(): + handlers, values = _build_routes( + {"check_civitai_api_key": MagicMock(side_effect=RuntimeError("boom"))} + ) + + response = await handlers[ + ("POST", "/model_resolver/civitai/api-key/check") + ](_request({"civitai_key": "bad"})) + + assert response.status == 500 + assert _response_json(response) == {"error": "boom"} + values["self"].logger.exception.assert_called_once() + + +@pytest.mark.asyncio +async def test_author_index_routes_refresh_and_clear_huggingface_cache(): + handlers, values = _build_routes() + + status_response = await handlers[ + ("GET", "/model_resolver/huggingface/author-index/status") + ](_request()) + refresh_response = await handlers[ + ("POST", "/model_resolver/huggingface/author-index/refresh") + ](_request({"hf_token": "token"})) + + assert _response_json(status_response) == {"ready": True} + assert _response_json(refresh_response) == {"refreshed": True} + values["refresh_known_author_fallback_indexes"].assert_called_once_with( + "token" + ) + values["clear_huggingface_search_cache"].assert_called_once_with() + + +@pytest.mark.asyncio +async def test_model_list_routes_forward_remote_check_and_update_caches(): + handlers, values = _build_routes() + + status_response = await handlers[ + ("GET", "/model_resolver/model-list/status") + ](_request(query={"check_remote": "true"})) + update_response = await handlers[ + ("POST", "/model_resolver/model-list/update") + ](_request()) + + assert _response_json(status_response) == {"loaded": True} + assert _response_json(update_response) == {"updated": True} + values["get_model_list_update_status"].assert_called_once_with( + check_remote=True + ) + values["update_model_list_from_remote"].assert_called_once_with() + assert values["self"].search_result_timestamps == {} + assert values["clear_all_search_caches"].call_count == 1