refactor(search): extract search support service
This commit is contained in:
Dariusz L
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
@@ -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({}))
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user