diff --git a/core/routes/civitai_search.py b/core/routes/civitai_search.py index 87e3590..6316eae 100644 --- a/core/routes/civitai_search.py +++ b/core/routes/civitai_search.py @@ -2,12 +2,12 @@ from ..services.model_service import ModelService from .context import RouteContext -from .helpers import register_service_post_route +from .helpers import register_service_route def register_civitai_search_routes(context: RouteContext): model_service = ModelService(context) - register_service_post_route( + register_service_route( context, path="/model_resolver/civitai-search", error_prefix="civitai-search", diff --git a/core/routes/custom_url.py b/core/routes/custom_url.py index 787c735..adbb1db 100644 --- a/core/routes/custom_url.py +++ b/core/routes/custom_url.py @@ -2,12 +2,12 @@ from ..services.model_service import ModelService from .context import RouteContext -from .helpers import register_service_post_route +from .helpers import register_service_route def register_custom_url_routes(context: RouteContext): model_service = ModelService(context) - register_service_post_route( + register_service_route( context, path="/model_resolver/custom-url", error_prefix="custom-url", diff --git a/core/routes/directories.py b/core/routes/directories.py index d6d48e9..d5023a0 100644 --- a/core/routes/directories.py +++ b/core/routes/directories.py @@ -3,6 +3,7 @@ from ..services.directory_service import DirectoryService from ..services.scanner_service import ScannerService from .context import RouteContext +from .helpers import register_service_route def register_directory_routes(context: RouteContext): @@ -14,29 +15,41 @@ def register_directory_routes(context: RouteContext): directory_service = DirectoryService(context) scanner_service = ScannerService(context) - @routes.get("/model_resolver/directories") - @json_api_endpoint("directories") - async def get_directories(request): - return await directory_service.get_directories(request) - - @routes.get("/model_resolver/root-directories") - @json_api_endpoint("root directories") - async def get_root_directories(request): - return await directory_service.get_root_directories(request) - - @routes.get("/model_resolver/path-template-suggestions") - @json_api_endpoint("path template suggestions") - async def get_path_template_suggestions(request): - return await scanner_service.get_path_template_suggestions(request) - - @routes.get("/model_resolver/capabilities") - @json_api_endpoint("capabilities") - async def get_capabilities(request): - return await directory_service.get_capabilities(request) + register_service_route( + context, + path="/model_resolver/directories", + methods=("get",), + error_prefix="directories", + operation=directory_service.get_directories, + ) + register_service_route( + context, + path="/model_resolver/root-directories", + methods=("get",), + error_prefix="root directories", + operation=directory_service.get_root_directories, + ) + register_service_route( + context, + path="/model_resolver/path-template-suggestions", + methods=("get",), + error_prefix="path template suggestions", + operation=scanner_service.get_path_template_suggestions, + ) + register_service_route( + context, + path="/model_resolver/capabilities", + methods=("get",), + error_prefix="capabilities", + operation=directory_service.get_capabilities, + ) register_version_routes(routes, web, json_api_endpoint) - @routes.get("/model_resolver/subfolders/{category}") - @json_api_endpoint("subfolders") - async def get_subfolders(request): - return await directory_service.get_subfolders(request) + register_service_route( + context, + path="/model_resolver/subfolders/{category}", + methods=("get",), + error_prefix="subfolders", + operation=directory_service.get_subfolders, + ) diff --git a/core/routes/downloads.py b/core/routes/downloads.py index 7e6ddf1..e212933 100644 --- a/core/routes/downloads.py +++ b/core/routes/downloads.py @@ -2,67 +2,87 @@ from ..services.download_service import DownloadService from .context import RouteContext +from .helpers import register_service_route def register_download_routes(context: RouteContext): """Register download routes and delegate behavior to :class:`DownloadService`.""" service = DownloadService(context) - json_api_endpoint = context.get("json_api_endpoint") - routes = context.get("routes") - @routes.post("/model_resolver/download") - @json_api_endpoint("download", return_success_on_error=True) - async def download_model(request): - return await service.download_model(request) - - @routes.get("/model_resolver/progress/{download_id}") - @json_api_endpoint("progress") - async def get_download_progress(request): - return await service.get_download_progress(request) - - @routes.get("/model_resolver/progress") - @json_api_endpoint("progress") - async def get_all_downloads_progress(request): - return await service.get_all_downloads_progress(request) - - @routes.post("/model_resolver/cancel/{download_id}") - @json_api_endpoint("cancel", return_success_on_error=True) - async def cancel_download_route(request): - return await service.cancel_download(request) - - @routes.post("/model_resolver/pause/{download_id}") - @json_api_endpoint("pause", return_success_on_error=True) - async def pause_download_route(request): - return await service.pause_download(request) - - @routes.post("/model_resolver/resume/{download_id}") - @json_api_endpoint("resume", return_success_on_error=True) - async def resume_download_route(request): - return await service.resume_download(request) - - @routes.post("/model_resolver/clear_completed_downloads") - @json_api_endpoint("clear_completed_downloads", return_success_on_error=True) - async def clear_completed_downloads_route(request): - return await service.clear_completed_downloads(request) - - @routes.get("/model_resolver/aria2/status") - @routes.post("/model_resolver/aria2/status") - @json_api_endpoint("aria2 status") - async def aria2_status_route(request): - return await service.aria2_status(request) - - @routes.post("/model_resolver/aria2/start") - @json_api_endpoint("aria2 start", return_success_on_error=True) - async def aria2_start_route(request): - return await service.aria2_start(request) - - @routes.get("/model_resolver/aria2/stop") - @routes.post("/model_resolver/aria2/stop") - @json_api_endpoint("aria2 stop", return_success_on_error=True) - async def aria2_stop_route(request): - return await service.aria2_stop(request) - - @routes.post("/model_resolver/aria2/install") - @json_api_endpoint("aria2 install") - async def aria2_install_route(request): - return await service.aria2_install(request) + register_service_route( + context, + path="/model_resolver/download", + error_prefix="download", + operation=service.download_model, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/progress/{download_id}", + methods=("get",), + error_prefix="progress", + operation=service.get_download_progress, + ) + register_service_route( + context, + path="/model_resolver/progress", + methods=("get",), + error_prefix="progress", + operation=service.get_all_downloads_progress, + ) + register_service_route( + context, + path="/model_resolver/cancel/{download_id}", + error_prefix="cancel", + operation=service.cancel_download, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/pause/{download_id}", + error_prefix="pause", + operation=service.pause_download, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/resume/{download_id}", + error_prefix="resume", + operation=service.resume_download, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/clear_completed_downloads", + error_prefix="clear_completed_downloads", + operation=service.clear_completed_downloads, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/aria2/status", + methods=("get", "post"), + error_prefix="aria2 status", + operation=service.aria2_status, + ) + register_service_route( + context, + path="/model_resolver/aria2/start", + error_prefix="aria2 start", + operation=service.aria2_start, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/aria2/stop", + methods=("get", "post"), + error_prefix="aria2 stop", + operation=service.aria2_stop, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/aria2/install", + error_prefix="aria2 install", + operation=service.aria2_install, + ) diff --git a/core/routes/hashes.py b/core/routes/hashes.py index 4e8892d..cdb3d89 100644 --- a/core/routes/hashes.py +++ b/core/routes/hashes.py @@ -2,54 +2,65 @@ from ..services.hash_service import HashService from .context import RouteContext +from .helpers import register_service_route def register_hash_routes(context: RouteContext): """Register hash routes and delegate behavior to :class:`HashService`.""" service = HashService(context) - json_api_endpoint = context.get("json_api_endpoint") routes = context.get("routes") - @routes.post("/model_resolver/local-model-hashes") - @json_api_endpoint("local-model-hashes") - async def local_model_hashes(request): - return await service.local_model_hashes(request) + register_service_route( + context, + path="/model_resolver/local-model-hashes", + error_prefix="local-model-hashes", + operation=service.local_model_hashes, + ) @routes.get("/model_resolver/model-preview") async def get_model_preview(request): return await service.get_model_preview(request) - @routes.post("/model_resolver/workflow-model-hashes") - @json_api_endpoint("workflow-model-hashes") - async def workflow_model_hashes(request): - return await service.workflow_model_hashes(request) - - @routes.post("/model_resolver/local-matches-by-hash") - @json_api_endpoint("local-matches-by-hash") - async def local_matches_by_hash(request): - return await service.local_matches_by_hash(request) - - @routes.post("/model_resolver/open-containing-folder") - @json_api_endpoint("open-containing-folder") - async def open_containing_folder(request): - return await service.open_containing_folder(request) - - @routes.post("/model_resolver/calculate-file-hash") - @json_api_endpoint("calculate-file-hash") - async def calculate_file_hash_route(request): - return await service.calculate_file_hash(request) - - @routes.post("/model_resolver/calculate-file-hash/start") - @json_api_endpoint("calculate-file-hash-start") - async def calculate_file_hash_start_route(request): - return await service.calculate_file_hash_start(request) - - @routes.get("/model_resolver/calculate-file-hash/progress/{progress_id}") - @json_api_endpoint("calculate-file-hash-progress") - async def calculate_file_hash_progress_route(request): - return await service.calculate_file_hash_progress(request) - - @routes.post("/model_resolver/calculate-file-hash/cancel/{progress_id}") - @json_api_endpoint("calculate-file-hash-cancel") - async def calculate_file_hash_cancel_route(request): - return await service.calculate_file_hash_cancel(request) + register_service_route( + context, + path="/model_resolver/workflow-model-hashes", + error_prefix="workflow-model-hashes", + operation=service.workflow_model_hashes, + ) + register_service_route( + context, + path="/model_resolver/local-matches-by-hash", + error_prefix="local-matches-by-hash", + operation=service.local_matches_by_hash, + ) + register_service_route( + context, + path="/model_resolver/open-containing-folder", + error_prefix="open-containing-folder", + operation=service.open_containing_folder, + ) + register_service_route( + context, + path="/model_resolver/calculate-file-hash", + error_prefix="calculate-file-hash", + operation=service.calculate_file_hash, + ) + register_service_route( + context, + path="/model_resolver/calculate-file-hash/start", + error_prefix="calculate-file-hash-start", + operation=service.calculate_file_hash_start, + ) + register_service_route( + context, + path="/model_resolver/calculate-file-hash/progress/{progress_id}", + methods=("get",), + error_prefix="calculate-file-hash-progress", + operation=service.calculate_file_hash_progress, + ) + register_service_route( + context, + path="/model_resolver/calculate-file-hash/cancel/{progress_id}", + error_prefix="calculate-file-hash-cancel", + operation=service.calculate_file_hash_cancel, + ) diff --git a/core/routes/helpers.py b/core/routes/helpers.py index 2f27c47..5fc9973 100644 --- a/core/routes/helpers.py +++ b/core/routes/helpers.py @@ -5,8 +5,16 @@ import threading from functools import wraps -def register_service_post_route(context, *, path, error_prefix, operation): - """Register a JSON POST route for a bound service operation.""" +def register_service_route( + context, + *, + path, + methods=("post",), + error_prefix, + operation, + return_success_on_error=False, +): + """Register JSON routes for a bound service operation.""" json_api_endpoint = context.get("json_api_endpoint") routes = context.get("routes") @@ -14,8 +22,18 @@ def register_service_post_route(context, *, path, error_prefix, operation): return await operation(request) service_route.__name__ = getattr(operation, "__name__", "service_route") - decorated_route = json_api_endpoint(error_prefix)(service_route) - routes.post(path)(decorated_route) + endpoint = ( + json_api_endpoint( + error_prefix, + return_success_on_error=True, + ) + if return_success_on_error + else json_api_endpoint(error_prefix) + ) + decorated_route = endpoint(service_route) + method_names = (methods,) if isinstance(methods, str) else methods + for method in method_names: + getattr(routes, method.lower())(path)(decorated_route) def create_route_helpers(web, logger, load_settings, hash_calculation_cancelled): diff --git a/core/routes/loaded_models.py b/core/routes/loaded_models.py index 21ebaa1..9185c2c 100644 --- a/core/routes/loaded_models.py +++ b/core/routes/loaded_models.py @@ -2,20 +2,23 @@ from ..services.loaded_models_service import LoadedModelsService from .context import RouteContext +from .helpers import register_service_route def register_loaded_model_routes(context: RouteContext): """Register loaded model routes and delegate behavior to a service.""" - json_api_endpoint = context.get("json_api_endpoint") - routes = context.get("routes") service = LoadedModelsService(context) - @routes.post("/model_resolver/loaded") - @json_api_endpoint("get_loaded_models") - async def get_loaded_models(request): - return await service.get_loaded_models(request) - - @routes.get("/model_resolver/loaded-progress/{loaded_id}") - @json_api_endpoint("loaded-progress") - async def get_loaded_models_progress(request): - return await service.get_loaded_models_progress(request) + register_service_route( + context, + path="/model_resolver/loaded", + error_prefix="get_loaded_models", + operation=service.get_loaded_models, + ) + register_service_route( + context, + path="/model_resolver/loaded-progress/{loaded_id}", + methods=("get",), + error_prefix="loaded-progress", + operation=service.get_loaded_models_progress, + ) diff --git a/core/routes/metadata.py b/core/routes/metadata.py index 16dd1fe..b99de85 100644 --- a/core/routes/metadata.py +++ b/core/routes/metadata.py @@ -3,46 +3,49 @@ from ..services.metadata_service import MetadataService from ..services.scanner_service import ScannerService from .context import RouteContext +from .helpers import register_service_route def register_metadata_routes(context: RouteContext): - json_api_endpoint = context.get('json_api_endpoint') - routes = context.get('routes') scanner_service = ScannerService(context) metadata_service = MetadataService(context) - @routes.get("/model_resolver/models") - @json_api_endpoint("get_models") - async def get_models(request): - """Get list of all available models.""" - return await scanner_service.get_models(request) - - @routes.post("/model_resolver/metadata-size-audit") - @json_api_endpoint("metadata-size-audit") - async def metadata_size_audit_route(request): - """Check local model .metadata.json sidecars for stale size values.""" - return await metadata_service.metadata_size_audit(request) - - @routes.get("/model_resolver/metadata-build/capabilities") - @json_api_endpoint("metadata-build-capabilities") - async def metadata_build_capabilities_route(request): - """Return local CPU and worker limits for the metadata builder.""" - return await metadata_service.metadata_build_capabilities(request) - - @routes.post("/model_resolver/metadata-build/start") - @json_api_endpoint("metadata-build-start") - async def metadata_build_start_route(request): - """Create missing local metadata sidecars and fill missing SHA256 values.""" - return await metadata_service.metadata_build_start(request) - - @routes.get("/model_resolver/metadata-build/progress/{progress_id}") - @json_api_endpoint("metadata-build-progress") - async def metadata_build_progress_route(request): - """Return progress for local metadata sidecar creation.""" - return await metadata_service.metadata_build_progress(request) - - @routes.post("/model_resolver/metadata-build/cancel/{progress_id}") - @json_api_endpoint("metadata-build-cancel") - async def metadata_build_cancel_route(request): - """Cancel local metadata sidecar creation.""" - return await metadata_service.metadata_build_cancel(request) + register_service_route( + context, + path="/model_resolver/models", + methods=("get",), + error_prefix="get_models", + operation=scanner_service.get_models, + ) + register_service_route( + context, + path="/model_resolver/metadata-size-audit", + error_prefix="metadata-size-audit", + operation=metadata_service.metadata_size_audit, + ) + register_service_route( + context, + path="/model_resolver/metadata-build/capabilities", + methods=("get",), + error_prefix="metadata-build-capabilities", + operation=metadata_service.metadata_build_capabilities, + ) + register_service_route( + context, + path="/model_resolver/metadata-build/start", + error_prefix="metadata-build-start", + operation=metadata_service.metadata_build_start, + ) + register_service_route( + context, + path="/model_resolver/metadata-build/progress/{progress_id}", + methods=("get",), + error_prefix="metadata-build-progress", + operation=metadata_service.metadata_build_progress, + ) + register_service_route( + context, + path="/model_resolver/metadata-build/cancel/{progress_id}", + error_prefix="metadata-build-cancel", + operation=metadata_service.metadata_build_cancel, + ) diff --git a/core/routes/model_details.py b/core/routes/model_details.py index 80239d5..b4fef3f 100644 --- a/core/routes/model_details.py +++ b/core/routes/model_details.py @@ -2,12 +2,12 @@ from ..services.model_service import ModelService from .context import RouteContext -from .helpers import register_service_post_route +from .helpers import register_service_route def register_model_details_routes(context: RouteContext): model_service = ModelService(context) - register_service_post_route( + register_service_route( context, path="/model_resolver/model-details", error_prefix="model-details", diff --git a/core/routes/search_support.py b/core/routes/search_support.py index e2266fd..03cd332 100644 --- a/core/routes/search_support.py +++ b/core/routes/search_support.py @@ -2,17 +2,19 @@ from ..services.search_support_service import SearchSupportService from .context import RouteContext +from .helpers import register_service_route def register_search_support_routes(context: RouteContext): - 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): - return await service.clear_search_cache(request) + register_service_route( + context, + path="/model_resolver/clear-search-cache", + error_prefix="Clear search cache", + operation=service.clear_search_cache, + ) @routes.post("/model_resolver/civitai/session-token/check") async def civitai_session_token_check_route(request): @@ -30,22 +32,29 @@ def register_search_support_routes(context: RouteContext): async def brave_api_key_check_route(request): 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 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): - 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 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): - return await service.update_model_list(request) + register_service_route( + context, + path="/model_resolver/huggingface/author-index/status", + methods=("get",), + error_prefix="HuggingFace author index status", + operation=service.author_index_status, + ) + register_service_route( + context, + path="/model_resolver/huggingface/author-index/refresh", + error_prefix="HuggingFace author index refresh", + operation=service.refresh_author_index, + ) + register_service_route( + context, + path="/model_resolver/model-list/status", + methods=("get",), + error_prefix="Model list status", + operation=service.model_list_status, + ) + register_service_route( + context, + path="/model_resolver/model-list/update", + error_prefix="Model list update", + operation=service.update_model_list, + ) diff --git a/core/routes/workflow_analysis.py b/core/routes/workflow_analysis.py index fc17239..82205ec 100644 --- a/core/routes/workflow_analysis.py +++ b/core/routes/workflow_analysis.py @@ -2,6 +2,7 @@ from ..services.workflow_service import WorkflowService from .context import RouteContext +from .helpers import register_service_route def register_workflow_analysis_routes(context: RouteContext): @@ -32,12 +33,16 @@ def register_workflow_analysis_routes(context: RouteContext): } ) - @routes.post("/model_resolver/resolve") - @json_api_endpoint("resolve", return_success_on_error=True) - async def resolve_models(request): - return await workflow_service.resolve_models(request) - - @routes.post("/model_resolver/local-matches") - @json_api_endpoint("local-matches") - async def local_matches(request): - return await workflow_service.local_matches(request) + register_service_route( + context, + path="/model_resolver/resolve", + error_prefix="resolve", + operation=workflow_service.resolve_models, + return_success_on_error=True, + ) + register_service_route( + context, + path="/model_resolver/local-matches", + error_prefix="local-matches", + operation=workflow_service.local_matches, + ) diff --git a/tests/test_route_helpers.py b/tests/test_route_helpers.py index 8510b9b..667fafb 100644 --- a/tests/test_route_helpers.py +++ b/tests/test_route_helpers.py @@ -5,13 +5,32 @@ from unittest.mock import MagicMock import pytest from aiohttp import web -from core.routes.helpers import create_route_helpers +from core.routes.helpers import create_route_helpers, register_service_route class _HashCancelled(Exception): pass +class _Routes: + def __init__(self): + self.handlers = {} + + def post(self, path): + def register(handler): + self.handlers[('POST', path)] = handler + return handler + + return register + + def get(self, path): + def register(handler): + self.handlers[('GET', path)] = handler + return handler + + return register + + def _build_helpers(): logger = MagicMock() helpers = create_route_helpers( @@ -23,6 +42,63 @@ def _build_helpers(): return helpers, logger +@pytest.mark.asyncio +async def test_register_service_route_delegates_to_the_operation(): + routes = _Routes() + requests = [] + + async def operation(request): + requests.append(request) + return 'ok' + + register_service_route( + { + 'json_api_endpoint': lambda _error_prefix: lambda func: func, + 'routes': routes, + }, + path='/model_resolver/example', + error_prefix='example', + operation=operation, + ) + + request = object() + response = await routes.handlers[('POST', '/model_resolver/example')](request) + + assert response == 'ok' + assert requests == [request] + + +@pytest.mark.asyncio +async def test_register_service_route_supports_multiple_methods_and_error_flags(): + routes = _Routes() + endpoint_options = [] + + def endpoint(error_prefix, return_success_on_error=False): + endpoint_options.append((error_prefix, return_success_on_error)) + return lambda func: func + + async def operation(_request): + return 'ok' + + register_service_route( + { + 'json_api_endpoint': endpoint, + 'routes': routes, + }, + path='/model_resolver/example', + methods=('get', 'post'), + error_prefix='example', + operation=operation, + return_success_on_error=True, + ) + + assert set(routes.handlers) == { + ('GET', '/model_resolver/example'), + ('POST', '/model_resolver/example'), + } + assert endpoint_options == [('example', True)] + + @pytest.mark.asyncio async def test_json_api_endpoint_returns_500_for_unexpected_exception(): (json_api_endpoint, *_), logger = _build_helpers()