From 68321b77992d9e55b3ca200d79bc97668d86d23f Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sun, 2 Aug 2026 01:50:27 +0200 Subject: [PATCH] refactor(downloads): extract download route handlers into DownloadService and add tests --- core/routes/downloads.py | 318 ++-------------------------- core/services/download_service.py | 334 ++++++++++++++++++++++++++++++ tests/test_download_routes.py | 289 ++++++++++++++++++++++++++ 3 files changed, 640 insertions(+), 301 deletions(-) create mode 100644 core/services/download_service.py create mode 100644 tests/test_download_routes.py diff --git a/core/routes/downloads.py b/core/routes/downloads.py index d42b3e7..7e6ddf1 100644 --- a/core/routes/downloads.py +++ b/core/routes/downloads.py @@ -1,352 +1,68 @@ -"""Downloads route registration.""" +"""HTTP adapters for download operations.""" +from ..services.download_service import DownloadService from .context import RouteContext def register_download_routes(context: RouteContext): - Aria2InstallError = context.get('Aria2InstallError') - UnsafeUrlError = context.get('UnsafeUrlError') - asyncio = context.get('asyncio') - cancel_download = context.get('cancel_download') - clear_completed_downloads = context.get('clear_completed_downloads') - first_non_empty = context.get('first_non_empty') - get_all_progress = context.get('get_all_progress') - get_aria2_status = context.get('get_aria2_status') - get_civarchive_model_details = context.get('get_civarchive_model_details') - get_civitai_model_details = context.get('get_civitai_model_details') - get_default_root_for_category = context.get('get_default_root_for_category') - get_download_directory = context.get('get_download_directory') - get_override_settings_from_request = context.get('get_override_settings_from_request') - get_progress = context.get('get_progress') - host_matches_domain = context.get('host_matches_domain') - install_aria2_engine = context.get('install_aria2_engine') - is_allowed_model_download_filename = context.get('is_allowed_model_download_filename') - json_api_endpoint = context.get('json_api_endpoint') - load_resolver_settings = context.get('load_resolver_settings') - normalize_download_category = context.get('normalize_download_category') - pause_download = context.get('pause_download') - resolve_download_subfolder = context.get('resolve_download_subfolder') - resume_download = context.get('resume_download') - routes = context.get('routes') - sanitize_download_filename = context.get('sanitize_download_filename') - save_resolver_settings = context.get('save_resolver_settings') - self = context.get('self') - split_path_segments = context.get('split_path_segments') - start_aria2_daemon = context.get('start_aria2_daemon') - start_background_download = context.get('start_background_download') - stop_aria2_daemon = context.get('stop_aria2_daemon') - to_bool = context.get('to_bool') - to_int = context.get('to_int') - validate_public_http_url = context.get('validate_public_http_url') - web = context.get('web') + """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): - """Start downloading a model.""" - data = await request.json() - url = data.get("url", "") - filename = data.get("filename", "") - category = data.get("category", "checkpoints") - category = normalize_download_category(category) - subfolder = data.get("subfolder", "") - base_directory = data.get("base_directory", "") - path_metadata = data.get("path_metadata", {}) - if not isinstance(path_metadata, dict): - path_metadata = {} - download_metadata = data.get("download_metadata") or data.get( - "metadata", {} - ) - if not isinstance(download_metadata, dict): - download_metadata = {} - download_metadata = dict(download_metadata) - settings = load_resolver_settings() - if not base_directory: - base_directory = get_default_root_for_category(category, settings) - subfolder = resolve_download_subfolder( - category, - subfolder, - path_metadata, - settings, - ) - - if not url: - return web.json_response( - {"error": "URL is required"}, status=400 - ) - try: - url = await asyncio.to_thread(validate_public_http_url, url) - except UnsafeUrlError as exc: - return web.json_response({"error": str(exc)}, status=400) - - from urllib.parse import urlparse as _download_urlparse - - download_host = _download_urlparse(url).hostname - - if not filename: - # Extract filename from URL - from urllib.parse import unquote, urlparse - - parsed = urlparse(url) - filename = unquote(parsed.path.split("/")[-1]) - filename = sanitize_download_filename(filename) - - if not filename: - return web.json_response( - {"error": "Could not determine filename"}, status=400 - ) - if not is_allowed_model_download_filename(filename): - return web.json_response( - {"error": "Unsupported model file extension"}, status=400 - ) - - # Build headers if needed - headers = {} - if host_matches_domain(download_host, "huggingface.co"): - hf_token = data.get("hf_token", "") - if hf_token: - headers["Authorization"] = f"Bearer {hf_token}" - elif host_matches_domain( - download_host, - "civitai.com", - "civitai.red", - ): - civitai_key = data.get("civitai_key", "") - if civitai_key and "token=" not in url: - url += ( - f"{'&' if '?' in url else '?'}token={civitai_key}" - ) - civitai_session_token = str( - data.get("civitai_session_token", "") or "" - ).strip() - if civitai_session_token: - headers["Cookie"] = ( - f"__Secure-civitai-token={civitai_session_token}" - ) - - def _first_metadata_value(*values): - return first_non_empty(*values) - - def _metadata_int(value): - return to_int(value) - - inferred_source = "" - if host_matches_domain(download_host, "civitai.com", "civitai.red"): - inferred_source = "civitai" - elif host_matches_domain(download_host, "huggingface.co"): - inferred_source = "huggingface" - - download_metadata.setdefault("filename", filename) - download_metadata.setdefault("category", category) - download_metadata.setdefault("download_url", url) - download_metadata.setdefault("source_url", url) - download_metadata.setdefault("path_metadata", path_metadata) - download_metadata.setdefault( - "source", - _first_metadata_value( - download_metadata.get("details_source"), - path_metadata.get("source"), - inferred_source, - ), - ) - - model_id = _metadata_int( - _first_metadata_value( - download_metadata.get("model_id"), - download_metadata.get("modelId"), - path_metadata.get("model_id"), - ) - ) - version_id = _metadata_int( - _first_metadata_value( - download_metadata.get("version_id"), - download_metadata.get("versionId"), - path_metadata.get("version_id"), - ) - ) - source_name = str( - _first_metadata_value( - download_metadata.get("details_source"), - download_metadata.get("source"), - ) - ).lower() - try: - if ( - source_name == "civitai" - and model_id - and not download_metadata.get("civitai_details") - ): - details = await asyncio.to_thread( - get_civitai_model_details, - model_id, - version_id, - data.get("civitai_key", ""), - ) - if details: - download_metadata["civitai_details"] = details - elif ( - source_name == "civarchive" - and model_id - and not download_metadata.get("civitai_details") - ): - details = await asyncio.to_thread( - get_civarchive_model_details, - model_id, - version_id, - ) - if details: - download_metadata["civitai_details"] = details - except Exception as metadata_error: - self.logger.warning( - f"Model metadata lookup failed: {metadata_error}" - ) - - target_directory = "" - target_path = "" - try: - import os as _download_os - - target_directory = ( - get_download_directory(category, base_directory) or "" - ) - if target_directory and subfolder: - target_directory = _download_os.path.join( - target_directory, *split_path_segments(subfolder) - ) - if target_directory: - target_path = _download_os.path.join( - target_directory, filename - ) - except Exception: - target_directory = "" - target_path = "" - - # Start background download - download_id = start_background_download( - url=url, - filename=filename, - category=category, - headers=headers if headers else None, - subfolder=subfolder, - base_directory=base_directory, - metadata=download_metadata, - ) - - return web.json_response( - { - "success": True, - "download_id": download_id, - "filename": filename, - "category": category, - "path": target_path, - "directory": target_directory, - "download_backend": settings.get("download_backend", "python"), - } - ) + return await service.download_model(request) @routes.get("/model_resolver/progress/{download_id}") @json_api_endpoint("progress") async def get_download_progress(request): - """Get progress for a specific download.""" - download_id = request.match_info["download_id"] - progress = get_progress(download_id) - - if progress: - return web.json_response(progress) - else: - return web.json_response( - {"error": "Download not found"}, status=404 - ) + return await service.get_download_progress(request) @routes.get("/model_resolver/progress") @json_api_endpoint("progress") async def get_all_downloads_progress(request): - """Get progress for all downloads.""" - progress = get_all_progress() - return web.json_response(progress) + 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): - """Cancel a download in progress.""" - download_id = request.match_info["download_id"] - cancel_download(download_id) - return web.json_response({"success": True}) + 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): - """Pause an aria2 download.""" - download_id = request.match_info["download_id"] - result = pause_download(download_id) - status = 200 if result.get("success") else 400 - return web.json_response(result, status=status) + 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): - """Resume an aria2 download.""" - download_id = request.match_info["download_id"] - result = resume_download(download_id) - status = 200 if result.get("success") else 400 - return web.json_response(result, status=status) + 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): - """Clear completed and failed downloads from progress memory.""" - clear_completed_downloads() - return web.json_response({"success": True}) + 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): - """Report aria2 availability using saved settings, with optional override via POST.""" - settings = await get_override_settings_from_request(request) - return web.json_response(await asyncio.to_thread(get_aria2_status, settings)) - + 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): - """Start the aria2 daemon without starting a download.""" - settings = await get_override_settings_from_request(request) - result = await asyncio.to_thread(start_aria2_daemon, settings) - status = 200 if result.get("success") else 400 - return web.json_response(result, status=status) + 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): - """Stop the aria2 daemon started by Model Resolver.""" - result = await asyncio.to_thread(stop_aria2_daemon) - status = 200 if result.get("success") else 400 - return web.json_response(result, status=status) + return await service.aria2_stop(request) @routes.post("/model_resolver/aria2/install") @json_api_endpoint("aria2 install") async def aria2_install_route(request): - """Download official aria2c binary and save its path.""" - try: - payload = await request.json() - except Exception: - payload = {} - force = False - if isinstance(payload, dict): - force = to_bool(payload.get("force"), False) - try: - install_result = await asyncio.to_thread(install_aria2_engine, force) - except Aria2InstallError as exc: - self.logger.warning(f"Model Resolver aria2 install failed: {exc}") - return web.json_response({"success": False, "error": str(exc)}, status=500) - settings = await asyncio.to_thread( - save_resolver_settings, - { - "aria2c_path": install_result.get("aria2c_path", ""), - "download_backend": "aria2", - }, - ) - install_result["settings"] = settings - return web.json_response(install_result) + return await service.aria2_install(request) diff --git a/core/services/download_service.py b/core/services/download_service.py new file mode 100644 index 0000000..22e966c --- /dev/null +++ b/core/services/download_service.py @@ -0,0 +1,334 @@ +"""Download orchestration service used by HTTP route adapters.""" + +import os +from urllib.parse import unquote, urlparse + +from ..routes.context import RouteContext + + +class DownloadService: + """Coordinate download requests without embedding business logic in routes.""" + + def __init__(self, context: RouteContext): + self.aria2_install_error = context.get("Aria2InstallError") + self.unsafe_url_error = context.get("UnsafeUrlError") + self.asyncio = context.get("asyncio") + self.cancel_download_fn = context.get("cancel_download") + self.clear_completed_downloads_fn = context.get( + "clear_completed_downloads" + ) + self.first_non_empty = context.get("first_non_empty") + self.get_all_progress_fn = context.get("get_all_progress") + self.get_aria2_status_fn = context.get("get_aria2_status") + self.get_civarchive_model_details = context.get( + "get_civarchive_model_details" + ) + self.get_civitai_model_details = context.get("get_civitai_model_details") + self.get_default_root_for_category = context.get( + "get_default_root_for_category" + ) + self.get_download_directory = context.get("get_download_directory") + self.get_override_settings_from_request = context.get( + "get_override_settings_from_request" + ) + self.get_progress_fn = context.get("get_progress") + self.host_matches_domain = context.get("host_matches_domain") + self.install_aria2_engine = context.get("install_aria2_engine") + self.is_allowed_model_download_filename = context.get( + "is_allowed_model_download_filename" + ) + self.load_resolver_settings = context.get("load_resolver_settings") + self.normalize_download_category = context.get( + "normalize_download_category" + ) + self.pause_download_fn = context.get("pause_download") + self.resolve_download_subfolder = context.get( + "resolve_download_subfolder" + ) + self.resume_download_fn = context.get("resume_download") + self.save_resolver_settings = context.get("save_resolver_settings") + self.sanitize_download_filename = context.get( + "sanitize_download_filename" + ) + self.extension = context.get("self") + self.split_path_segments = context.get("split_path_segments") + self.start_aria2_daemon = context.get("start_aria2_daemon") + self.start_background_download = context.get("start_background_download") + self.stop_aria2_daemon = context.get("stop_aria2_daemon") + self.to_bool = context.get("to_bool") + self.to_int = context.get("to_int") + self.validate_public_http_url = context.get("validate_public_http_url") + self.web = context.get("web") + + @property + def logger(self): + return self.extension.logger + + async def download_model(self, request): + """Validate and start a background model download.""" + data = await request.json() + url = data.get("url", "") + filename = data.get("filename", "") + category = self.normalize_download_category( + data.get("category", "checkpoints") + ) + subfolder = data.get("subfolder", "") + base_directory = data.get("base_directory", "") + path_metadata = data.get("path_metadata", {}) + if not isinstance(path_metadata, dict): + path_metadata = {} + download_metadata = data.get("download_metadata") or data.get( + "metadata", {} + ) + if not isinstance(download_metadata, dict): + download_metadata = {} + download_metadata = dict(download_metadata) + settings = self.load_resolver_settings() + if not base_directory: + base_directory = self.get_default_root_for_category(category, settings) + subfolder = self.resolve_download_subfolder( + category, + subfolder, + path_metadata, + settings, + ) + + if not url: + return self.web.json_response({"error": "URL is required"}, status=400) + try: + url = await self.asyncio.to_thread(self.validate_public_http_url, url) + except self.unsafe_url_error as exc: + return self.web.json_response({"error": str(exc)}, status=400) + + download_host = urlparse(url).hostname + + if not filename: + filename = unquote(urlparse(url).path.split("/")[-1]) + filename = self.sanitize_download_filename(filename) + + if not filename: + return self.web.json_response( + {"error": "Could not determine filename"}, status=400 + ) + if not self.is_allowed_model_download_filename(filename): + return self.web.json_response( + {"error": "Unsupported model file extension"}, status=400 + ) + + headers = {} + if self.host_matches_domain(download_host, "huggingface.co"): + hf_token = data.get("hf_token", "") + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" + elif self.host_matches_domain( + download_host, + "civitai.com", + "civitai.red", + ): + civitai_key = data.get("civitai_key", "") + if civitai_key and "token=" not in url: + url += f"{'&' if '?' in url else '?'}token={civitai_key}" + civitai_session_token = str( + data.get("civitai_session_token", "") or "" + ).strip() + if civitai_session_token: + headers["Cookie"] = ( + f"__Secure-civitai-token={civitai_session_token}" + ) + + inferred_source = "" + if self.host_matches_domain(download_host, "civitai.com", "civitai.red"): + inferred_source = "civitai" + elif self.host_matches_domain(download_host, "huggingface.co"): + inferred_source = "huggingface" + + download_metadata.setdefault("filename", filename) + download_metadata.setdefault("category", category) + download_metadata.setdefault("download_url", url) + download_metadata.setdefault("source_url", url) + download_metadata.setdefault("path_metadata", path_metadata) + download_metadata.setdefault( + "source", + self.first_non_empty( + download_metadata.get("details_source"), + path_metadata.get("source"), + inferred_source, + ), + ) + + model_id = self.to_int( + self.first_non_empty( + download_metadata.get("model_id"), + download_metadata.get("modelId"), + path_metadata.get("model_id"), + ) + ) + version_id = self.to_int( + self.first_non_empty( + download_metadata.get("version_id"), + download_metadata.get("versionId"), + path_metadata.get("version_id"), + ) + ) + source_name = str( + self.first_non_empty( + download_metadata.get("details_source"), + download_metadata.get("source"), + ) + ).lower() + try: + if ( + source_name == "civitai" + and model_id + and not download_metadata.get("civitai_details") + ): + details = await self.asyncio.to_thread( + self.get_civitai_model_details, + model_id, + version_id, + data.get("civitai_key", ""), + ) + if details: + download_metadata["civitai_details"] = details + elif ( + source_name == "civarchive" + and model_id + and not download_metadata.get("civitai_details") + ): + details = await self.asyncio.to_thread( + self.get_civarchive_model_details, + model_id, + version_id, + ) + if details: + download_metadata["civitai_details"] = details + except Exception as metadata_error: + self.logger.warning( + f"Model metadata lookup failed: {metadata_error}" + ) + + target_directory = "" + target_path = "" + try: + target_directory = self.get_download_directory( + category, + base_directory, + ) or "" + if target_directory and subfolder: + target_directory = os.path.join( + target_directory, + *self.split_path_segments(subfolder), + ) + if target_directory: + target_path = os.path.join(target_directory, filename) + except Exception: + target_directory = "" + target_path = "" + + download_id = self.start_background_download( + url=url, + filename=filename, + category=category, + headers=headers if headers else None, + subfolder=subfolder, + base_directory=base_directory, + metadata=download_metadata, + ) + + return self.web.json_response( + { + "success": True, + "download_id": download_id, + "filename": filename, + "category": category, + "path": target_path, + "directory": target_directory, + "download_backend": settings.get("download_backend", "python"), + } + ) + + async def get_download_progress(self, request): + """Return progress for one download.""" + download_id = request.match_info["download_id"] + progress = self.get_progress_fn(download_id) + if progress: + return self.web.json_response(progress) + return self.web.json_response({"error": "Download not found"}, status=404) + + async def get_all_downloads_progress(self, request): + """Return progress for all downloads.""" + return self.web.json_response(self.get_all_progress_fn()) + + async def cancel_download(self, request): + """Cancel a download in progress.""" + download_id = request.match_info["download_id"] + self.cancel_download_fn(download_id) + return self.web.json_response({"success": True}) + + async def pause_download(self, request): + """Pause an aria2 download.""" + download_id = request.match_info["download_id"] + result = self.pause_download_fn(download_id) + status = 200 if result.get("success") else 400 + return self.web.json_response(result, status=status) + + async def resume_download(self, request): + """Resume an aria2 download.""" + download_id = request.match_info["download_id"] + result = self.resume_download_fn(download_id) + status = 200 if result.get("success") else 400 + return self.web.json_response(result, status=status) + + async def clear_completed_downloads(self, request): + """Clear completed and failed downloads from progress memory.""" + self.clear_completed_downloads_fn() + return self.web.json_response({"success": True}) + + async def aria2_status(self, request): + """Report aria2 availability using optional request settings.""" + settings = await self.get_override_settings_from_request(request) + result = await self.asyncio.to_thread(self.get_aria2_status_fn, settings) + return self.web.json_response(result) + + async def aria2_start(self, request): + """Start the aria2 daemon without starting a download.""" + settings = await self.get_override_settings_from_request(request) + result = await self.asyncio.to_thread(self.start_aria2_daemon, settings) + status = 200 if result.get("success") else 400 + return self.web.json_response(result, status=status) + + async def aria2_stop(self, request): + """Stop the aria2 daemon started by Model Resolver.""" + result = await self.asyncio.to_thread(self.stop_aria2_daemon) + status = 200 if result.get("success") else 400 + return self.web.json_response(result, status=status) + + async def aria2_install(self, request): + """Install aria2 and persist the selected download backend.""" + try: + payload = await request.json() + except Exception: + payload = {} + force = False + if isinstance(payload, dict): + force = self.to_bool(payload.get("force"), False) + try: + install_result = await self.asyncio.to_thread( + self.install_aria2_engine, + force, + ) + except self.aria2_install_error as exc: + self.logger.warning(f"Model Resolver aria2 install failed: {exc}") + return self.web.json_response( + {"success": False, "error": str(exc)}, + status=500, + ) + settings = await self.asyncio.to_thread( + self.save_resolver_settings, + { + "aria2c_path": install_result.get("aria2c_path", ""), + "download_backend": "aria2", + }, + ) + install_result["settings"] = settings + return self.web.json_response(install_result) diff --git a/tests/test_download_routes.py b/tests/test_download_routes.py new file mode 100644 index 0000000..2239394 --- /dev/null +++ b/tests/test_download_routes.py @@ -0,0 +1,289 @@ +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.downloads import register_download_routes +from core.type_utils import to_bool, to_int + + +class _UnsafeUrlError(Exception): + pass + + +class _Aria2InstallError(Exception): + pass + + +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 _build_download_routes(overrides=None): + routes = _Routes() + extension = SimpleNamespace(logger=MagicMock()) + sanitize_download_filename = MagicMock( + side_effect=lambda filename: str(filename).strip() + ) + + def host_matches_domain(host, *domains): + return any(host == domain or host.endswith(f".{domain}") for domain in domains) + + def first_non_empty(*values): + return next((value for value in values if value not in (None, "")), "") + + values = { + "Aria2InstallError": _Aria2InstallError, + "UnsafeUrlError": _UnsafeUrlError, + "asyncio": asyncio, + "cancel_download": MagicMock(), + "clear_completed_downloads": MagicMock(), + "first_non_empty": first_non_empty, + "get_all_progress": MagicMock(return_value={"one": {"status": "done"}}), + "get_aria2_status": MagicMock(return_value={"status": "running"}), + "get_civarchive_model_details": MagicMock( + return_value={"name": "Archive model"} + ), + "get_civitai_model_details": MagicMock( + return_value={"name": "CivitAI model"} + ), + "get_default_root_for_category": MagicMock(return_value=r"C:\base"), + "get_download_directory": MagicMock(return_value=r"C:\models"), + "get_override_settings_from_request": AsyncMock( + return_value={"aria2c_path": "aria2c"} + ), + "get_progress": MagicMock(return_value=None), + "host_matches_domain": host_matches_domain, + "install_aria2_engine": MagicMock( + return_value={"success": True, "aria2c_path": r"C:\tools\aria2c.exe"} + ), + "is_allowed_model_download_filename": lambda filename: filename.endswith( + ".safetensors" + ), + "json_api_endpoint": lambda _name, **_kwargs: lambda handler: handler, + "load_resolver_settings": MagicMock( + return_value={"download_backend": "aria2"} + ), + "normalize_download_category": lambda category: str(category).lower(), + "pause_download": MagicMock(return_value={"success": True}), + "resolve_download_subfolder": MagicMock( + side_effect=lambda category, subfolder, metadata, settings: subfolder + ), + "resume_download": MagicMock(return_value={"success": True}), + "routes": routes, + "sanitize_download_filename": sanitize_download_filename, + "save_resolver_settings": MagicMock(return_value={"saved": True}), + "self": extension, + "split_path_segments": lambda value: [ + part for part in str(value).replace("\\", "/").split("/") if part + ], + "start_aria2_daemon": MagicMock(return_value={"success": True}), + "start_background_download": MagicMock(return_value="download-1"), + "stop_aria2_daemon": MagicMock(return_value={"success": True}), + "to_bool": to_bool, + "to_int": to_int, + "validate_public_http_url": MagicMock( + side_effect=lambda url: url + ), + "web": web, + } + if overrides: + values.update(overrides) + + register_download_routes(RouteContext(values)) + return routes.handlers, values + + +def _request(payload=None, *, download_id="download-1", method="POST"): + return SimpleNamespace( + json=AsyncMock(return_value=payload), + match_info={"download_id": download_id}, + method=method, + ) + + +@pytest.mark.asyncio +async def test_download_route_builds_metadata_headers_and_target_path(): + handlers, values = _build_download_routes() + handler = handlers[("POST", "/model_resolver/download")] + payload = { + "url": "https://civitai.com/api/download/models/123", + "filename": "model.safetensors", + "category": "Checkpoints", + "subfolder": "character/pony", + "path_metadata": { + "source": "civitai", + "model_id": 123, + "version_id": 456, + }, + "civitai_key": "api-key", + "civitai_session_token": "session-token", + } + + response = await handler(_request(payload)) + + body = json.loads(response.text) + assert response.status == 200 + assert body["success"] is True + assert body["download_id"] == "download-1" + assert body["category"] == "checkpoints" + assert body["directory"] == r"C:\models\character\pony" + assert body["path"] == r"C:\models\character\pony\model.safetensors" + values["start_background_download"].assert_called_once() + download_call = values["start_background_download"].call_args.kwargs + assert "token=api-key" in download_call["url"] + assert download_call["headers"] == { + "Cookie": "__Secure-civitai-token=session-token" + } + assert download_call["metadata"]["civitai_details"] == { + "name": "CivitAI model" + } + values["get_civitai_model_details"].assert_called_once_with( + 123, + 456, + "api-key", + ) + + +@pytest.mark.asyncio +async def test_download_route_rejects_missing_unsafe_and_unsupported_inputs(): + handlers, values = _build_download_routes() + handler = handlers[("POST", "/model_resolver/download")] + + response = await handler(_request({})) + assert response.status == 400 + assert json.loads(response.text) == {"error": "URL is required"} + + values["validate_public_http_url"].side_effect = _UnsafeUrlError("blocked") + response = await handler(_request({"url": "http://127.0.0.1/model.safetensors"})) + assert response.status == 400 + assert json.loads(response.text) == {"error": "blocked"} + + values["validate_public_http_url"].side_effect = lambda url: url + values["sanitize_download_filename"].side_effect = lambda filename: "" + response = await handler( + _request({"url": "https://example.com/download/model.safetensors"}) + ) + assert response.status == 400 + assert json.loads(response.text) == {"error": "Could not determine filename"} + + values["sanitize_download_filename"].side_effect = lambda filename: "model.ckpt" + response = await handler( + _request( + { + "url": "https://example.com/download/model.ckpt", + "filename": "model.ckpt", + } + ) + ) + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Unsupported model file extension" + } + + +@pytest.mark.asyncio +async def test_download_control_routes_preserve_progress_and_status_codes(): + handlers, values = _build_download_routes() + values["get_progress"].return_value = {"status": "downloading", "progress": 50} + + progress_handler = handlers[("GET", "/model_resolver/progress/{download_id}")] + response = await progress_handler(_request(download_id="progress-1", method="GET")) + assert json.loads(response.text) == {"status": "downloading", "progress": 50} + values["get_progress"].assert_called_once_with("progress-1") + + all_progress_handler = handlers[("GET", "/model_resolver/progress")] + response = await all_progress_handler(_request(method="GET")) + assert json.loads(response.text) == {"one": {"status": "done"}} + + cancel_handler = handlers[("POST", "/model_resolver/cancel/{download_id}")] + response = await cancel_handler(_request(download_id="cancel-1")) + assert json.loads(response.text) == {"success": True} + values["cancel_download"].assert_called_once_with("cancel-1") + + pause_handler = handlers[("POST", "/model_resolver/pause/{download_id}")] + response = await pause_handler(_request(download_id="pause-1")) + assert response.status == 200 + values["pause_download"].assert_called_once_with("pause-1") + + values["pause_download"].return_value = {"success": False, "error": "not aria2"} + response = await pause_handler(_request(download_id="pause-2")) + assert response.status == 400 + + resume_handler = handlers[("POST", "/model_resolver/resume/{download_id}")] + response = await resume_handler(_request(download_id="resume-1")) + assert response.status == 200 + values["resume_download"].assert_called_once_with("resume-1") + + clear_handler = handlers[("POST", "/model_resolver/clear_completed_downloads")] + response = await clear_handler(_request()) + assert json.loads(response.text) == {"success": True} + values["clear_completed_downloads"].assert_called_once_with() + + +@pytest.mark.asyncio +async def test_aria2_routes_use_override_settings_and_map_failures(): + handlers, values = _build_download_routes() + status_get = handlers[("GET", "/model_resolver/aria2/status")] + status_post = handlers[("POST", "/model_resolver/aria2/status")] + start_handler = handlers[("POST", "/model_resolver/aria2/start")] + stop_get = handlers[("GET", "/model_resolver/aria2/stop")] + install_handler = handlers[("POST", "/model_resolver/aria2/install")] + + response = await status_get(_request(method="GET")) + assert json.loads(response.text) == {"status": "running"} + response = await status_post(_request({}, method="POST")) + assert json.loads(response.text) == {"status": "running"} + assert values["get_aria2_status"].call_count == 2 + + values["start_aria2_daemon"].return_value = { + "success": False, + "error": "cannot start", + } + response = await start_handler(_request()) + assert response.status == 400 + + values["stop_aria2_daemon"].return_value = { + "success": False, + "error": "not running", + } + response = await stop_get(_request(method="GET")) + assert response.status == 400 + + response = await install_handler(_request({"force": True})) + assert json.loads(response.text)["success"] is True + values["install_aria2_engine"].assert_called_once_with(True) + values["save_resolver_settings"].assert_called_once_with( + { + "aria2c_path": r"C:\tools\aria2c.exe", + "download_backend": "aria2", + } + ) + + values["install_aria2_engine"].side_effect = _Aria2InstallError("install failed") + response = await install_handler(_request({})) + assert response.status == 500 + assert json.loads(response.text) == { + "success": False, + "error": "install failed", + } + values["self"].logger.warning.assert_called_once()