refactor(downloads): extract download route handlers into DownloadService and add tests
This commit is contained in:
Dariusz L
+17
-301
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user