mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-12 04:22:13 -03:00
Will Miao
c16de68640
The new informational Windows job surfaced 91 failures: four real bugs (fixed in
the previous two commits) and a long tail of POSIX assumptions in the tests
themselves. None of them assert behaviour that differs on Windows, so they are
fixed rather than skipped wherever the behaviour is the same:
* Path spelling: expectations built with `str(tmp_path)` are compared through an
identity helper (`normalize_root_identity`) or the platform's separator, since
the code hands out forward-slashed business paths while `str()` is native.
* Root labels: `_root_display_labels` prefixes the drive on Windows ("C: loras"),
so label assertions compare the final path component or derive the expected
label from the same helper.
* `~` expansion reads USERPROFILE (ntpath), not HOME; timer granularity rounds
0.1s down; a case-only rename counts as "already named" via `os.path.normcase`.
* Genuinely POSIX-only strategies are skipped with a reason instead of failing:
the headless clipboard fallback, dbus file-manager selection, xdg-open, and
the case-fold index test that simulates a case-sensitive filesystem.
2217 lines
72 KiB
Python
2217 lines
72 KiB
Python
"""Error handling and execution tests for DownloadManager."""
|
|
|
|
import asyncio
|
|
import os
|
|
import zipfile
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Optional
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from py.services.download_manager import DownloadManager
|
|
from py.services.download_queue_service import DownloadQueueService
|
|
from py.services.downloader import DownloadStreamControl
|
|
from py.services import download_manager
|
|
from py.services import aria2_transfer_state
|
|
from py.services.errors import DownloadRateLimitError, RateLimitError
|
|
from py.services.rate_limit_coordinator import RateLimitCoordinator
|
|
from py.services.service_registry import ServiceRegistry
|
|
from py.services.settings_manager import SettingsManager, get_settings_manager
|
|
from py.utils.metadata_manager import MetadataManager
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_download_manager():
|
|
"""Ensure each test operates on a fresh singleton."""
|
|
DownloadManager._instance = None
|
|
yield
|
|
DownloadManager._instance = None
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def isolate_settings(monkeypatch, tmp_path):
|
|
"""Point settings writes at a temporary directory to avoid touching real files."""
|
|
manager = get_settings_manager()
|
|
default_settings = manager._get_default_settings()
|
|
default_settings.update(
|
|
{
|
|
"default_lora_root": str(tmp_path),
|
|
"default_checkpoint_root": str(tmp_path / "checkpoints"),
|
|
"default_embedding_root": str(tmp_path / "embeddings"),
|
|
"download_path_templates": {
|
|
"lora": "{base_model}/{first_tag}",
|
|
"checkpoint": "{base_model}/{first_tag}",
|
|
"embedding": "{base_model}/{first_tag}",
|
|
},
|
|
"base_model_path_mappings": {"BaseModel": "MappedModel"},
|
|
}
|
|
)
|
|
monkeypatch.setattr(manager, "settings", default_settings)
|
|
monkeypatch.setattr(SettingsManager, "_save_settings", lambda self: None)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def isolate_aria2_state(monkeypatch, tmp_path):
|
|
state_path = tmp_path / "cache" / "aria2" / "downloads.json"
|
|
monkeypatch.setattr(
|
|
aria2_transfer_state,
|
|
"get_aria2_state_path",
|
|
lambda: str(state_path),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def queue_service(tmp_path, monkeypatch):
|
|
"""Return a tmp-backed DownloadQueueService and stub the singleton."""
|
|
service = DownloadQueueService(db_path=str(tmp_path / "queue.sqlite"))
|
|
|
|
async def fake_get_instance(_cls=None):
|
|
return service
|
|
|
|
monkeypatch.setattr(DownloadQueueService, "get_instance", fake_get_instance)
|
|
return service
|
|
|
|
|
|
@pytest.fixture
|
|
def reset_rate_limit_coordinator():
|
|
RateLimitCoordinator._instance = None
|
|
yield
|
|
RateLimitCoordinator._instance = None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_retries_urls(monkeypatch, tmp_path):
|
|
"""Test that download retries multiple URLs on failure."""
|
|
manager = DownloadManager()
|
|
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
initial_path = save_dir / "file.safetensors"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, _path):
|
|
return None
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
metadata = DummyMetadata(initial_path)
|
|
version_info = {"images": []}
|
|
download_urls = [
|
|
"https://first.example/file.safetensors",
|
|
"https://second.example/file.safetensors",
|
|
]
|
|
|
|
class DummyDownloader:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def download_file(
|
|
self, url, path, progress_callback=None, use_auth=None, **_kwargs
|
|
):
|
|
self.calls.append((url, path, use_auth))
|
|
if len(self.calls) == 1:
|
|
return False, "first failed"
|
|
# Create the target file to simulate a successful download
|
|
Path(path).write_text("content")
|
|
return True, "second success"
|
|
|
|
dummy_downloader = DummyDownloader()
|
|
monkeypatch.setattr(
|
|
download_manager, "get_downloader", AsyncMock(return_value=dummy_downloader)
|
|
)
|
|
|
|
class DummyScanner:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def add_model_to_cache(self, metadata_dict, relative_path):
|
|
self.calls.append((metadata_dict, relative_path))
|
|
|
|
dummy_scanner = DummyScanner()
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_get_checkpoint_scanner",
|
|
AsyncMock(return_value=dummy_scanner),
|
|
)
|
|
monkeypatch.setattr(
|
|
ServiceRegistry, "get_embedding_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=download_urls,
|
|
save_dir=str(save_dir),
|
|
metadata=metadata,
|
|
version_info=version_info,
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id=None,
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert [url for url, *_ in dummy_downloader.calls] == download_urls
|
|
assert dummy_scanner.calls # ensure cache updated
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_uses_aria2_backend_for_model_files(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
settings = get_settings_manager()
|
|
settings.settings["download_backend"] = "aria2"
|
|
settings.settings["civitai_api_key"] = "secret-key"
|
|
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
target_path = save_dir / "file.safetensors"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, _path):
|
|
return None
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
class DummyAria2Downloader:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def download_file(
|
|
self,
|
|
url,
|
|
save_path,
|
|
*,
|
|
download_id,
|
|
progress_callback=None,
|
|
headers=None,
|
|
):
|
|
self.calls.append(
|
|
{
|
|
"url": url,
|
|
"save_path": save_path,
|
|
"download_id": download_id,
|
|
"headers": headers,
|
|
}
|
|
)
|
|
Path(save_path).write_text("content")
|
|
return True, save_path
|
|
|
|
dummy_aria2 = DummyAria2Downloader()
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=dummy_aria2),
|
|
)
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_downloader",
|
|
AsyncMock(side_effect=AssertionError("python downloader should not be used")),
|
|
)
|
|
|
|
class DummyScanner:
|
|
async def add_model_to_cache(self, metadata_dict, relative_path):
|
|
return {"metadata": metadata_dict, "relative_path": relative_path}
|
|
|
|
dummy_scanner = DummyScanner()
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_get_checkpoint_scanner",
|
|
AsyncMock(return_value=dummy_scanner),
|
|
)
|
|
monkeypatch.setattr(
|
|
ServiceRegistry, "get_embedding_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=["https://civitai.com/api/download/models/1"],
|
|
save_dir=str(save_dir),
|
|
metadata=DummyMetadata(target_path),
|
|
version_info={"images": []},
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id="download-1",
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert dummy_aria2.calls == [
|
|
{
|
|
"url": "https://civitai.com/api/download/models/1",
|
|
"save_path": str(target_path),
|
|
"download_id": "download-1",
|
|
"headers": {"Authorization": "Bearer secret-key"},
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_allows_anonymous_civitai_with_aria2(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
settings = get_settings_manager()
|
|
settings.settings["download_backend"] = "aria2"
|
|
settings.settings["civitai_api_key"] = ""
|
|
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
target_path = save_dir / "file.safetensors"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, _path):
|
|
return None
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
class DummyAria2Downloader:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def download_file(
|
|
self,
|
|
url,
|
|
save_path,
|
|
*,
|
|
download_id,
|
|
progress_callback=None,
|
|
headers=None,
|
|
):
|
|
self.calls.append({"url": url, "headers": headers, "download_id": download_id})
|
|
Path(save_path).write_text("content")
|
|
return True, save_path
|
|
|
|
dummy_aria2 = DummyAria2Downloader()
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=dummy_aria2),
|
|
)
|
|
|
|
dummy_scanner = SimpleNamespace(add_model_to_cache=AsyncMock(return_value=None))
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=["https://civitai.com/api/download/models/1"],
|
|
save_dir=str(save_dir),
|
|
metadata=DummyMetadata(target_path),
|
|
version_info={"images": []},
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id="download-2",
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert dummy_aria2.calls == [
|
|
{
|
|
"url": "https://civitai.com/api/download/models/1",
|
|
"headers": None,
|
|
"download_id": "download-2",
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_adjusts_checkpoint_sub_type(monkeypatch, tmp_path):
|
|
"""Test that checkpoint sub_type is adjusted during download."""
|
|
manager = DownloadManager()
|
|
|
|
root_dir = tmp_path / "checkpoints"
|
|
root_dir.mkdir()
|
|
save_dir = root_dir
|
|
target_path = save_dir / "model.safetensors"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = path.as_posix()
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
self.preview_nsfw_level = 0
|
|
self.sub_type = "checkpoint"
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, updated_path):
|
|
self.file_path = Path(updated_path).as_posix()
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"file_path": self.file_path,
|
|
"sub_type": self.sub_type,
|
|
"sha256": self.sha256,
|
|
}
|
|
|
|
metadata = DummyMetadata(target_path)
|
|
version_info = {"images": []}
|
|
download_urls = ["https://example.invalid/model.safetensors"]
|
|
|
|
class DummyDownloader:
|
|
async def download_file(
|
|
self, _url, path, progress_callback=None, use_auth=None, **_kwargs
|
|
):
|
|
Path(path).write_text("content")
|
|
return True, "ok"
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_downloader",
|
|
AsyncMock(return_value=DummyDownloader()),
|
|
)
|
|
|
|
class DummyCheckpointScanner:
|
|
def __init__(self, root: Path):
|
|
self.root = root.as_posix()
|
|
self.add_calls = []
|
|
|
|
def _find_root_for_file(self, file_path: str):
|
|
return self.root if file_path.startswith(self.root) else None
|
|
|
|
def adjust_metadata(
|
|
self, metadata_obj, _file_path: str, root_path: Optional[str]
|
|
):
|
|
if root_path:
|
|
metadata_obj.sub_type = "diffusion_model"
|
|
return metadata_obj
|
|
|
|
def adjust_cached_entry(self, entry):
|
|
if entry.get("file_path", "").startswith(self.root):
|
|
entry["sub_type"] = "diffusion_model"
|
|
return entry
|
|
|
|
async def add_model_to_cache(self, metadata_dict, relative_path):
|
|
self.add_calls.append((metadata_dict, relative_path))
|
|
return True
|
|
|
|
dummy_scanner = DummyCheckpointScanner(root_dir)
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_get_checkpoint_scanner",
|
|
AsyncMock(return_value=dummy_scanner),
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=download_urls,
|
|
save_dir=str(save_dir),
|
|
metadata=metadata,
|
|
version_info=version_info,
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="checkpoint",
|
|
download_id=None,
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert metadata.sub_type == "diffusion_model"
|
|
saved_metadata = MetadataManager.save_metadata.await_args.args[1]
|
|
assert saved_metadata.sub_type == "diffusion_model"
|
|
assert dummy_scanner.add_calls
|
|
cached_entry, _ = dummy_scanner.add_calls[0]
|
|
assert cached_entry["sub_type"] == "diffusion_model"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_extracts_zip_single_model(monkeypatch, tmp_path):
|
|
"""Test extraction of single model from ZIP file."""
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
zip_path = save_dir / "bundle.zip"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, updated_path):
|
|
self.file_path = str(updated_path)
|
|
self.file_name = Path(updated_path).stem
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
metadata = DummyMetadata(zip_path)
|
|
# Simulate a zip entry with an API-reported AutoV3: it must be cleared
|
|
# for extracted models, which resolve their own header hash instead.
|
|
metadata.autov3 = "zipzipzip123"
|
|
version_info = {"images": []}
|
|
download_urls = ["https://example.invalid/model.zip"]
|
|
|
|
class DummyDownloader:
|
|
async def download_file(self, *_args, **_kwargs):
|
|
with zipfile.ZipFile(str(zip_path), "w") as archive:
|
|
archive.writestr("inner/model.safetensors", b"model")
|
|
archive.writestr("docs/readme.txt", b"ignore")
|
|
return True, "ok"
|
|
|
|
monkeypatch.setattr(
|
|
download_manager, "get_downloader", AsyncMock(return_value=DummyDownloader())
|
|
)
|
|
|
|
class ImmediateLoop:
|
|
async def run_in_executor(self, executor, func, *args):
|
|
return func(*args)
|
|
|
|
monkeypatch.setattr(download_manager.asyncio, "get_running_loop", lambda: ImmediateLoop())
|
|
|
|
dummy_scanner = SimpleNamespace(add_model_to_cache=AsyncMock(return_value=None))
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=download_urls,
|
|
save_dir=str(save_dir),
|
|
metadata=metadata,
|
|
version_info=version_info,
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id=None,
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert not zip_path.exists()
|
|
extracted = save_dir / "model.safetensors"
|
|
assert extracted.exists()
|
|
saved_call = MetadataManager.save_metadata.await_args
|
|
assert saved_call.args[0] == str(extracted).replace(os.sep, "/")
|
|
# SHA256 comes from metadata (API value), not recalculated
|
|
assert saved_call.args[1].sha256 == "sha256"
|
|
# The zip-level AutoV3 was cleared; the extracted file is not safetensors,
|
|
# so the entry is marked checked-unavailable rather than inheriting it.
|
|
assert saved_call.args[1].autov3 == ""
|
|
assert dummy_scanner.add_model_to_cache.await_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_extracts_zip_multiple_models(monkeypatch, tmp_path):
|
|
"""Test extraction of multiple models from ZIP file."""
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
zip_path = save_dir / "bundle.zip"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, updated_path):
|
|
self.file_path = str(updated_path)
|
|
self.file_name = Path(updated_path).stem
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
metadata = DummyMetadata(zip_path)
|
|
version_info = {"images": []}
|
|
download_urls = ["https://example.invalid/model.zip"]
|
|
|
|
class DummyDownloader:
|
|
async def download_file(self, *_args, **_kwargs):
|
|
with zipfile.ZipFile(str(zip_path), "w") as archive:
|
|
archive.writestr("first/model-one.safetensors", b"one")
|
|
archive.writestr("second/model-two.safetensors", b"two")
|
|
archive.writestr("readme.md", b"ignore")
|
|
return True, "ok"
|
|
|
|
monkeypatch.setattr(
|
|
download_manager, "get_downloader", AsyncMock(return_value=DummyDownloader())
|
|
)
|
|
|
|
class ImmediateLoop:
|
|
async def run_in_executor(self, executor, func, *args):
|
|
return func(*args)
|
|
|
|
monkeypatch.setattr(download_manager.asyncio, "get_running_loop", lambda: ImmediateLoop())
|
|
|
|
dummy_scanner = SimpleNamespace(add_model_to_cache=AsyncMock(return_value=None))
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=download_urls,
|
|
save_dir=str(save_dir),
|
|
metadata=metadata,
|
|
version_info=version_info,
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id=None,
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert not zip_path.exists()
|
|
extracted_one = save_dir / "model-one.safetensors"
|
|
extracted_two = save_dir / "model-two.safetensors"
|
|
assert extracted_one.exists()
|
|
assert extracted_two.exists()
|
|
|
|
assert MetadataManager.save_metadata.await_count == 2
|
|
assert dummy_scanner.add_model_to_cache.await_count == 2
|
|
|
|
metadata_calls = MetadataManager.save_metadata.await_args_list
|
|
assert metadata_calls[0].args[0] == str(extracted_one)
|
|
# SHA256 comes from metadata (API value), not recalculated
|
|
assert metadata_calls[0].args[1].sha256 == "sha256"
|
|
assert metadata_calls[1].args[0] == str(extracted_two)
|
|
assert metadata_calls[1].args[1].sha256 == "sha256"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_extracts_zip_pt_embedding(monkeypatch, tmp_path):
|
|
"""Test extraction of .pt embedding files from ZIP."""
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
zip_path = save_dir / "bundle.zip"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, updated_path):
|
|
self.file_path = str(updated_path)
|
|
self.file_name = Path(updated_path).stem
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
metadata = DummyMetadata(zip_path)
|
|
version_info = {"images": []}
|
|
download_urls = ["https://example.invalid/model.zip"]
|
|
|
|
class DummyDownloader:
|
|
async def download_file(self, *_args, **_kwargs):
|
|
with zipfile.ZipFile(str(zip_path), "w") as archive:
|
|
archive.writestr("inner/embedding.pt", b"embedding")
|
|
archive.writestr("docs/readme.txt", b"ignore")
|
|
return True, "ok"
|
|
|
|
monkeypatch.setattr(
|
|
download_manager, "get_downloader", AsyncMock(return_value=DummyDownloader())
|
|
)
|
|
|
|
class ImmediateLoop:
|
|
async def run_in_executor(self, executor, func, *args):
|
|
return func(*args)
|
|
|
|
monkeypatch.setattr(download_manager.asyncio, "get_running_loop", lambda: ImmediateLoop())
|
|
|
|
dummy_scanner = SimpleNamespace(add_model_to_cache=AsyncMock(return_value=None))
|
|
monkeypatch.setattr(
|
|
ServiceRegistry, "get_embedding_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=download_urls,
|
|
save_dir=str(save_dir),
|
|
metadata=metadata,
|
|
version_info=version_info,
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="embedding",
|
|
download_id=None,
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert not zip_path.exists()
|
|
extracted = save_dir / "embedding.pt"
|
|
assert extracted.exists()
|
|
saved_call = MetadataManager.save_metadata.await_args
|
|
assert saved_call.args[0] == str(extracted).replace(os.sep, "/")
|
|
# SHA256 comes from metadata (API value), not recalculated
|
|
assert saved_call.args[1].sha256 == "sha256"
|
|
assert dummy_scanner.add_model_to_cache.await_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_extract_model_files_from_archive_uses_executor(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
archive_path = tmp_path / "bundle.zip"
|
|
with zipfile.ZipFile(archive_path, "w") as archive:
|
|
archive.writestr("inner/model.safetensors", b"model")
|
|
|
|
captured = {}
|
|
|
|
class ImmediateLoop:
|
|
async def run_in_executor(self, executor, func, *args):
|
|
captured["executor"] = executor
|
|
return func(*args)
|
|
|
|
monkeypatch.setattr(
|
|
download_manager.asyncio,
|
|
"get_running_loop",
|
|
lambda: ImmediateLoop(),
|
|
)
|
|
|
|
extracted = await manager._extract_model_files_from_archive(
|
|
str(archive_path),
|
|
{".safetensors"},
|
|
)
|
|
|
|
assert captured["executor"] is manager._archive_executor
|
|
assert len(extracted) == 1
|
|
assert extracted[0].endswith("model.safetensors")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_download_updates_state():
|
|
"""Test that pause_download updates download state correctly."""
|
|
manager = DownloadManager()
|
|
|
|
download_id = "dl"
|
|
manager._download_tasks[download_id] = object()
|
|
pause_control = DownloadStreamControl()
|
|
manager._pause_events[download_id] = pause_control
|
|
manager._active_downloads[download_id] = {
|
|
"status": "downloading",
|
|
"bytes_per_second": 42.0,
|
|
}
|
|
|
|
result = await manager.pause_download(download_id)
|
|
|
|
assert result == {"success": True, "message": "Download paused successfully"}
|
|
assert download_id in manager._pause_events
|
|
assert manager._pause_events[download_id].is_set() is False
|
|
assert manager._active_downloads[download_id]["status"] == "paused"
|
|
assert manager._active_downloads[download_id]["bytes_per_second"] == 0.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_download_reverts_local_pause_when_aria2_pause_fails(monkeypatch):
|
|
manager = DownloadManager()
|
|
|
|
download_id = "dl"
|
|
manager._download_tasks[download_id] = object()
|
|
pause_control = DownloadStreamControl()
|
|
manager._pause_events[download_id] = pause_control
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "downloading",
|
|
"bytes_per_second": 42.0,
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
async def has_transfer(self, _download_id):
|
|
return True
|
|
|
|
async def pause_download(self, _download_id):
|
|
return {"success": False, "error": "rpc failed"}
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.pause_download(download_id)
|
|
|
|
assert result == {"success": False, "error": "rpc failed"}
|
|
assert pause_control.is_set() is True
|
|
assert manager._active_downloads[download_id]["status"] == "downloading"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_download_reverts_local_pause_when_aria2_probe_raises(monkeypatch):
|
|
manager = DownloadManager()
|
|
|
|
download_id = "dl"
|
|
manager._download_tasks[download_id] = object()
|
|
pause_control = DownloadStreamControl()
|
|
manager._pause_events[download_id] = pause_control
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "downloading",
|
|
"bytes_per_second": 42.0,
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
async def has_transfer(self, _download_id):
|
|
raise RuntimeError("rpc unavailable")
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.pause_download(download_id)
|
|
|
|
assert result == {"success": False, "error": "rpc unavailable"}
|
|
assert pause_control.is_set() is True
|
|
assert manager._active_downloads[download_id]["status"] == "downloading"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_download_returns_error_when_aria2_probe_raises(monkeypatch):
|
|
manager = DownloadManager()
|
|
|
|
download_id = "dl"
|
|
pause_control = DownloadStreamControl()
|
|
pause_control.pause()
|
|
manager._pause_events[download_id] = pause_control
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "paused",
|
|
"bytes_per_second": 0.0,
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
async def has_transfer(self, _download_id):
|
|
raise RuntimeError("rpc unavailable")
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.resume_download(download_id)
|
|
|
|
assert result == {"success": False, "error": "rpc unavailable"}
|
|
assert pause_control.is_paused() is True
|
|
assert manager._active_downloads[download_id]["status"] == "paused"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_download_does_not_spawn_restored_worker_when_aria2_resume_fails(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
|
|
download_id = "dl"
|
|
save_path = tmp_path / "file.safetensors"
|
|
pause_control = DownloadStreamControl()
|
|
pause_control.pause()
|
|
manager._pause_events[download_id] = pause_control
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "paused",
|
|
"bytes_per_second": 0.0,
|
|
}
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
download_id,
|
|
{
|
|
"download_id": download_id,
|
|
"transfer_backend": "aria2",
|
|
"status": "paused",
|
|
"save_path": str(save_path),
|
|
"file_path": str(save_path),
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"resume_context": {
|
|
"version_info": {"id": 34, "modelId": 12, "model": {"id": 12}},
|
|
"file_info": {
|
|
"name": "file.safetensors",
|
|
"downloadUrl": "https://example.com/file.safetensors",
|
|
},
|
|
"model_type": "lora",
|
|
"relative_path": "",
|
|
"save_dir": str(tmp_path),
|
|
"download_urls": ["https://example.com/file.safetensors"],
|
|
},
|
|
},
|
|
)
|
|
|
|
resume_restored = AsyncMock(return_value={"success": True})
|
|
monkeypatch.setattr(manager, "_resume_restored_aria2_download", resume_restored)
|
|
|
|
class DummyAria2Downloader:
|
|
async def has_transfer(self, _download_id):
|
|
return True
|
|
|
|
async def resume_download(self, _download_id):
|
|
return {"success": False, "error": "rpc unavailable"}
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.resume_download(download_id)
|
|
|
|
assert result == {"success": False, "error": "rpc unavailable"}
|
|
assert download_id not in manager._download_tasks
|
|
assert resume_restored.await_count == 0
|
|
assert pause_control.is_paused() is True
|
|
assert manager._active_downloads[download_id]["status"] == "paused"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_background_download_task_cleans_up_finished_restore_task():
|
|
manager = DownloadManager()
|
|
download_id = "download-1"
|
|
manager._pause_events[download_id] = DownloadStreamControl()
|
|
|
|
async def finished_restore():
|
|
return {"success": True}
|
|
|
|
task = manager._start_background_download_task(download_id, finished_restore())
|
|
await task
|
|
await asyncio.sleep(0)
|
|
|
|
assert download_id not in manager._download_tasks
|
|
assert download_id not in manager._pause_events
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_download_still_cancels_local_task_when_aria2_raises(monkeypatch):
|
|
manager = DownloadManager()
|
|
|
|
started = asyncio.Event()
|
|
|
|
async def blocked_task():
|
|
started.set()
|
|
await asyncio.sleep(60)
|
|
|
|
task = asyncio.create_task(blocked_task())
|
|
await started.wait()
|
|
|
|
download_id = "download-queued"
|
|
manager._download_tasks[download_id] = task
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "queued",
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
async def cancel_download(self, _download_id):
|
|
raise RuntimeError("rpc unavailable")
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.cancel_download(download_id)
|
|
|
|
assert result["success"] is True
|
|
assert task.cancelled() or task.done()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_download_preserves_tracking_when_aria2_returns_error(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
download_id = "download-queued"
|
|
save_path = tmp_path / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
(tmp_path / "file.safetensors.aria2").write_text("control")
|
|
|
|
pause_control = DownloadStreamControl()
|
|
manager._pause_events[download_id] = pause_control
|
|
manager._download_tasks[download_id] = object()
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "downloading",
|
|
"file_path": str(save_path),
|
|
}
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
download_id,
|
|
{
|
|
"download_id": download_id,
|
|
"transfer_backend": "aria2",
|
|
"status": "downloading",
|
|
"save_path": str(save_path),
|
|
"file_path": str(save_path),
|
|
},
|
|
)
|
|
|
|
cleanup_files = AsyncMock(return_value=None)
|
|
monkeypatch.setattr(manager, "_cleanup_cancelled_download_files", cleanup_files)
|
|
|
|
class DummyAria2Downloader:
|
|
async def cancel_download(self, _download_id):
|
|
return {"success": False, "error": "rpc unavailable"}
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.cancel_download(download_id)
|
|
|
|
assert result == {"success": False, "error": "rpc unavailable"}
|
|
assert download_id in manager._download_tasks
|
|
assert download_id in manager._pause_events
|
|
assert await manager._aria2_state_store.get(download_id) is not None
|
|
assert cleanup_files.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_download_rejects_completed_history_entry(tmp_path):
|
|
manager = DownloadManager()
|
|
download_id = "completed-download"
|
|
save_path = tmp_path / "file.safetensors"
|
|
metadata_path = tmp_path / "file.metadata.json"
|
|
preview_path = tmp_path / "file.jpeg"
|
|
save_path.write_text("complete")
|
|
metadata_path.write_text("{}")
|
|
preview_path.write_text("preview")
|
|
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "completed",
|
|
"file_path": str(save_path),
|
|
"preview_path": str(preview_path),
|
|
}
|
|
|
|
result = await manager.cancel_download(download_id)
|
|
|
|
assert result == {"success": False, "error": "Download task not found"}
|
|
assert save_path.exists()
|
|
assert metadata_path.exists()
|
|
assert preview_path.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_download_removes_preview_and_aria2_control_files(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
|
|
started = asyncio.Event()
|
|
|
|
async def blocked_task():
|
|
started.set()
|
|
await asyncio.sleep(60)
|
|
|
|
task = asyncio.create_task(blocked_task())
|
|
await started.wait()
|
|
|
|
save_path = tmp_path / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
aria2_path = tmp_path / "file.safetensors.aria2"
|
|
aria2_path.write_text("control")
|
|
preview_path = tmp_path / "file.jpeg"
|
|
preview_path.write_text("preview")
|
|
|
|
download_id = "download-queued"
|
|
manager._download_tasks[download_id] = task
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "queued",
|
|
"file_path": str(save_path),
|
|
"aria2_control_path": str(aria2_path),
|
|
"preview_path": str(preview_path),
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
async def cancel_download(self, _download_id):
|
|
return {"success": True, "message": "cancelled"}
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.cancel_download(download_id)
|
|
|
|
assert result["success"] is True
|
|
assert not save_path.exists()
|
|
assert not aria2_path.exists()
|
|
assert not preview_path.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_download_does_not_delete_untracked_same_basename_preview(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
|
|
started = asyncio.Event()
|
|
|
|
async def blocked_task():
|
|
started.set()
|
|
await asyncio.sleep(60)
|
|
|
|
task = asyncio.create_task(blocked_task())
|
|
await started.wait()
|
|
|
|
save_path = tmp_path / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
aria2_path = tmp_path / "file.safetensors.aria2"
|
|
aria2_path.write_text("control")
|
|
manual_preview_path = tmp_path / "file.jpg"
|
|
manual_preview_path.write_text("manual")
|
|
|
|
download_id = "download-queued"
|
|
manager._download_tasks[download_id] = task
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "queued",
|
|
"file_path": str(save_path),
|
|
"aria2_control_path": str(aria2_path),
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
async def cancel_download(self, _download_id):
|
|
return {"success": True, "message": "cancelled"}
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.cancel_download(download_id)
|
|
|
|
assert result["success"] is True
|
|
assert not save_path.exists()
|
|
assert not aria2_path.exists()
|
|
assert manual_preview_path.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_cancelled_download_files_retries_aria2_control_deletion(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
download_id = "download-1"
|
|
|
|
save_path = tmp_path / "file.safetensors"
|
|
aria2_path = tmp_path / "file.safetensors.aria2"
|
|
save_path.write_text("partial")
|
|
aria2_path.write_text("control")
|
|
|
|
original_unlink = os.unlink
|
|
attempts = {"count": 0}
|
|
|
|
def flaky_unlink(path):
|
|
if path == str(aria2_path) and attempts["count"] == 0:
|
|
attempts["count"] += 1
|
|
raise PermissionError("still locked")
|
|
return original_unlink(path)
|
|
|
|
monkeypatch.setattr(download_manager.os, "unlink", flaky_unlink)
|
|
monkeypatch.setattr("py.services.download_manager.asyncio.sleep", AsyncMock())
|
|
|
|
await manager._cleanup_cancelled_download_files(
|
|
download_id,
|
|
{
|
|
"file_path": str(save_path),
|
|
"aria2_control_path": str(aria2_path),
|
|
"transfer_backend": "aria2",
|
|
},
|
|
)
|
|
|
|
assert attempts["count"] == 1
|
|
assert not save_path.exists()
|
|
assert not aria2_path.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_waits_for_paused_pre_transfer_gate(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
target_path = save_dir / "file.safetensors"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, _path):
|
|
return None
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
pause_control = DownloadStreamControl()
|
|
pause_control.pause()
|
|
manager._pause_events["download-1"] = pause_control
|
|
manager._active_downloads["download-1"] = {
|
|
"status": "downloading",
|
|
"bytes_per_second": 42.0,
|
|
}
|
|
|
|
dummy_scanner = SimpleNamespace(add_model_to_cache=AsyncMock(return_value=None))
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
started = asyncio.Event()
|
|
allow_finish = asyncio.Event()
|
|
captured = {"calls": 0}
|
|
|
|
async def fake_download_model_file(
|
|
self,
|
|
download_url,
|
|
save_path,
|
|
*,
|
|
backend,
|
|
progress_callback,
|
|
use_auth,
|
|
download_id,
|
|
pause_control,
|
|
):
|
|
captured["calls"] += 1
|
|
started.set()
|
|
await allow_finish.wait()
|
|
Path(save_path).write_text("content")
|
|
return True, save_path
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_download_model_file",
|
|
fake_download_model_file,
|
|
)
|
|
|
|
task = asyncio.create_task(
|
|
manager._execute_download(
|
|
download_urls=["https://civitai.com/api/download/models/1"],
|
|
save_dir=str(save_dir),
|
|
metadata=DummyMetadata(target_path),
|
|
version_info={"images": []},
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id="download-1",
|
|
transfer_backend="aria2",
|
|
)
|
|
)
|
|
|
|
await asyncio.sleep(0)
|
|
assert started.is_set() is False
|
|
assert captured["calls"] == 0
|
|
assert manager._active_downloads["download-1"]["status"] == "paused"
|
|
|
|
pause_control.resume()
|
|
await asyncio.wait_for(started.wait(), timeout=1.0)
|
|
assert captured["calls"] == 1
|
|
assert manager._active_downloads["download-1"]["status"] == "downloading"
|
|
|
|
allow_finish.set()
|
|
result = await task
|
|
|
|
assert result == {"success": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_reuses_existing_aria2_partial_path(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
target_path = save_dir / "file.safetensors"
|
|
target_path.write_text("partial")
|
|
control_path = save_dir / "file.safetensors.aria2"
|
|
control_path.write_text("control")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"download-1",
|
|
{
|
|
"download_id": "download-1",
|
|
"transfer_backend": "aria2",
|
|
"save_path": str(target_path),
|
|
"file_path": str(target_path),
|
|
"status": "paused",
|
|
},
|
|
)
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return "renamed.safetensors"
|
|
|
|
def update_file_info(self, _path):
|
|
return None
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
manager._active_downloads["download-1"] = {"transfer_backend": "aria2"}
|
|
dummy_scanner = SimpleNamespace(add_model_to_cache=AsyncMock(return_value=None))
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
async def fake_download_model_file(
|
|
self,
|
|
download_url,
|
|
save_path,
|
|
*,
|
|
backend,
|
|
progress_callback,
|
|
use_auth,
|
|
download_id,
|
|
pause_control,
|
|
):
|
|
Path(save_path).write_text("content")
|
|
return True, save_path
|
|
|
|
monkeypatch.setattr(DownloadManager, "_download_model_file", fake_download_model_file)
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=["https://example.com/file.safetensors"],
|
|
save_dir=str(save_dir),
|
|
metadata=DummyMetadata(target_path),
|
|
version_info={"images": []},
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id="download-1",
|
|
transfer_backend="aria2",
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert manager._active_downloads["download-1"]["file_path"] == str(target_path)
|
|
assert not (save_dir / "renamed.safetensors").exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_rejects_conflicting_aria2_partial_path(tmp_path):
|
|
manager = DownloadManager()
|
|
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
target_path = save_dir / "file.safetensors"
|
|
target_path.write_text("partial")
|
|
(save_dir / "file.safetensors.aria2").write_text("control")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"other-download",
|
|
{
|
|
"download_id": "other-download",
|
|
"transfer_backend": "aria2",
|
|
"save_path": str(target_path),
|
|
"file_path": str(target_path),
|
|
"status": "paused",
|
|
},
|
|
)
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
raise AssertionError("should not rename")
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=["https://example.com/file.safetensors"],
|
|
save_dir=str(save_dir),
|
|
metadata=DummyMetadata(target_path),
|
|
version_info={"images": []},
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id="download-1",
|
|
transfer_backend="aria2",
|
|
)
|
|
|
|
assert result["success"] is False
|
|
assert "already using" in result["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_reassigns_same_aria2_partial_to_new_download_id(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
target_path = save_dir / "file.safetensors"
|
|
target_path.write_text("partial")
|
|
(save_dir / "file.safetensors.aria2").write_text("control")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"old-download",
|
|
{
|
|
"download_id": "old-download",
|
|
"transfer_backend": "aria2",
|
|
"save_path": str(target_path),
|
|
"file_path": str(target_path),
|
|
"status": "paused",
|
|
"model_id": 11,
|
|
"model_version_id": 22,
|
|
},
|
|
)
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
raise AssertionError("should not rename")
|
|
|
|
def update_file_info(self, _path):
|
|
return None
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
class DummyAria2Downloader:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def reassign_transfer(self, previous_download_id, new_download_id):
|
|
self.calls.append(("reassign_transfer", previous_download_id, new_download_id))
|
|
return None
|
|
|
|
dummy_aria2 = DummyAria2Downloader()
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=dummy_aria2),
|
|
)
|
|
|
|
manager._active_downloads["old-download"] = {
|
|
"transfer_backend": "aria2",
|
|
"model_id": 11,
|
|
"model_version_id": 22,
|
|
"status": "paused",
|
|
}
|
|
manager._active_downloads["new-download"] = {
|
|
"transfer_backend": "aria2",
|
|
"model_id": 11,
|
|
"model_version_id": 22,
|
|
"status": "queued",
|
|
}
|
|
|
|
resolved, path = await manager._resolve_download_target_path(
|
|
str(save_dir),
|
|
DummyMetadata(target_path),
|
|
transfer_backend="aria2",
|
|
download_id="new-download",
|
|
)
|
|
|
|
assert resolved is True
|
|
assert path == str(target_path)
|
|
assert "old-download" not in manager._active_downloads
|
|
assert manager._active_downloads["new-download"]["file_path"] == str(target_path)
|
|
assert dummy_aria2.calls == [("reassign_transfer", "old-download", "new-download")]
|
|
assert await manager._aria2_state_store.get("old-download") is None
|
|
persisted = await manager._aria2_state_store.get("new-download")
|
|
assert persisted is not None
|
|
assert persisted["save_path"] == str(target_path)
|
|
|
|
|
|
def test_is_same_aria2_download_request_requires_version_id_match():
|
|
manager = DownloadManager()
|
|
|
|
assert (
|
|
manager._is_same_aria2_download_request(
|
|
{"model_id": 1, "model_version_id": None},
|
|
{"model_id": 1, "model_version_id": 2},
|
|
)
|
|
is False
|
|
)
|
|
assert (
|
|
manager._is_same_aria2_download_request(
|
|
{"model_id": 1, "model_version_id": 3},
|
|
{"model_id": 1, "model_version_id": None},
|
|
)
|
|
is False
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_adopt_existing_aria2_download_cancels_old_running_task(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
save_path = tmp_path / "file.safetensors"
|
|
|
|
started = asyncio.Event()
|
|
cancelled = asyncio.Event()
|
|
call_order = []
|
|
|
|
async def old_download():
|
|
started.set()
|
|
try:
|
|
await asyncio.sleep(60)
|
|
except asyncio.CancelledError:
|
|
call_order.append("old-task-cancelled")
|
|
cancelled.set()
|
|
raise
|
|
|
|
old_task = asyncio.create_task(old_download())
|
|
await started.wait()
|
|
|
|
manager._download_tasks["old-download"] = old_task
|
|
old_pause_control = DownloadStreamControl()
|
|
old_pause_control.pause()
|
|
manager._pause_events["old-download"] = old_pause_control
|
|
manager._active_downloads["old-download"] = {
|
|
"transfer_backend": "aria2",
|
|
"model_id": 11,
|
|
"model_version_id": 22,
|
|
"status": "downloading",
|
|
}
|
|
manager._active_downloads["new-download"] = {
|
|
"transfer_backend": "aria2",
|
|
"model_id": 11,
|
|
"model_version_id": 22,
|
|
"status": "queued",
|
|
}
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"old-download",
|
|
{
|
|
"download_id": "old-download",
|
|
"transfer_backend": "aria2",
|
|
"save_path": str(save_path),
|
|
"file_path": str(save_path),
|
|
"status": "downloading",
|
|
"model_id": 11,
|
|
"model_version_id": 22,
|
|
},
|
|
)
|
|
|
|
class DummyAria2Downloader:
|
|
async def reassign_transfer(self, previous_download_id, new_download_id):
|
|
call_order.append("reassign-transfer")
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
await manager._adopt_existing_aria2_download(
|
|
"old-download",
|
|
"new-download",
|
|
{"model_id": 11, "model_version_id": 22},
|
|
str(save_path),
|
|
)
|
|
|
|
assert cancelled.is_set() is True
|
|
assert "old-download" not in manager._download_tasks
|
|
assert call_order == ["reassign-transfer", "old-task-cancelled"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_download_rejects_unknown_task():
|
|
"""Test that pause_download rejects unknown download tasks."""
|
|
manager = DownloadManager()
|
|
|
|
result = await manager.pause_download("missing")
|
|
|
|
assert result == {"success": False, "error": "Download task not found"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_download_sets_event_and_status():
|
|
"""Test that resume_download sets event and updates status."""
|
|
manager = DownloadManager()
|
|
|
|
download_id = "dl"
|
|
pause_control = DownloadStreamControl()
|
|
pause_control.pause()
|
|
pause_control.mark_progress()
|
|
manager._pause_events[download_id] = pause_control
|
|
manager._active_downloads[download_id] = {
|
|
"status": "paused",
|
|
"bytes_per_second": 0.0,
|
|
}
|
|
|
|
result = await manager.resume_download(download_id)
|
|
|
|
assert result == {"success": True, "message": "Download resumed successfully"}
|
|
assert manager._pause_events[download_id].is_set() is True
|
|
assert manager._active_downloads[download_id]["status"] == "downloading"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_download_requests_reconnect_for_stalled_stream():
|
|
"""Test that resume_download requests reconnect for stalled streams."""
|
|
manager = DownloadManager()
|
|
|
|
download_id = "dl"
|
|
pause_control = DownloadStreamControl(stall_timeout=40)
|
|
pause_control.pause()
|
|
pause_control.last_progress_timestamp = datetime.now().timestamp() - 120
|
|
manager._pause_events[download_id] = pause_control
|
|
manager._active_downloads[download_id] = {
|
|
"status": "paused",
|
|
"bytes_per_second": 0.0,
|
|
}
|
|
|
|
result = await manager.resume_download(download_id)
|
|
|
|
assert result == {"success": True, "message": "Download resumed successfully"}
|
|
assert pause_control.is_set() is True
|
|
assert pause_control.has_reconnect_request() is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_download_rejects_when_not_paused():
|
|
"""Test that resume_download rejects when download is not paused."""
|
|
|
|
manager = DownloadManager()
|
|
|
|
download_id = "dl"
|
|
pause_control = DownloadStreamControl()
|
|
manager._pause_events[download_id] = pause_control
|
|
|
|
result = await manager.resume_download(download_id)
|
|
|
|
assert result == {"success": False, "error": "Download is not paused"}
|
|
|
|
|
|
class _SharedTargetMetadata:
|
|
"""Metadata double whose generate_unique_filename mirrors BaseModelMetadata."""
|
|
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.file_name = path.stem
|
|
self.sha256 = "abcdef1234567890"
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, target_dir, base_name, extension, hash_provider=None):
|
|
original_filename = f"{base_name}{extension}"
|
|
if not os.path.exists(os.path.join(target_dir, original_filename)):
|
|
return original_filename
|
|
short_hash = (hash_provider() if hash_provider else "0000")[:4]
|
|
unique_filename = f"{base_name}-{short_hash}{extension}"
|
|
counter = 1
|
|
while os.path.exists(os.path.join(target_dir, unique_filename)):
|
|
unique_filename = f"{base_name}-{short_hash}-{counter}{extension}"
|
|
counter += 1
|
|
return unique_filename
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_downloads_with_same_target_path_do_not_destroy_each_other(
|
|
monkeypatch, tmp_path
|
|
):
|
|
"""Two downloads racing on the same target filename must not corrupt each other.
|
|
|
|
Regression scenario (#issue: same-name files of model versions 2665422 and
|
|
2925310 downloaded concurrently): one task fails mid-flight and its failure
|
|
cleanup deletes the target file the other task just finished downloading,
|
|
crashing the survivor in _build_metadata_entries with FileNotFoundError.
|
|
"""
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
target_path = save_dir / "GTO-MuraiJuria-RG4535.safetensors"
|
|
|
|
# Deterministic interleave control
|
|
both_resolved = asyncio.Event()
|
|
a_wrote_file = asyncio.Event()
|
|
build_started = asyncio.Event()
|
|
release_build = asyncio.Event()
|
|
resolve_count = {"n": 0}
|
|
|
|
original_resolve = manager._resolve_download_target_path
|
|
|
|
async def tracking_resolve(*args, **kwargs):
|
|
result = await original_resolve(*args, **kwargs)
|
|
resolve_count["n"] += 1
|
|
if resolve_count["n"] >= 2:
|
|
both_resolved.set()
|
|
return result
|
|
|
|
monkeypatch.setattr(manager, "_resolve_download_target_path", tracking_resolve)
|
|
|
|
original_build = manager._build_metadata_entries
|
|
|
|
async def paced_build(base_metadata, file_paths):
|
|
build_started.set()
|
|
await release_build.wait()
|
|
return await original_build(base_metadata, file_paths)
|
|
|
|
monkeypatch.setattr(manager, "_build_metadata_entries", paced_build)
|
|
|
|
downloader_paths = []
|
|
|
|
class DummyDownloader:
|
|
async def download_file(
|
|
self, url, path, progress_callback=None, use_auth=None, **_kwargs
|
|
):
|
|
downloader_paths.append(str(path))
|
|
if "2665422" in url:
|
|
# Task A: wait until both tasks resolved the same target path,
|
|
# then complete successfully.
|
|
try:
|
|
await asyncio.wait_for(both_resolved.wait(), timeout=0.5)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
Path(path).write_text("model-a-content")
|
|
a_wrote_file.set()
|
|
return True, "ok"
|
|
# Task B: fail as soon as A's file physically exists, mirroring a
|
|
# conflicting aria2/transfer error on the shared target path.
|
|
await a_wrote_file.wait()
|
|
return False, "simulated conflict failure"
|
|
|
|
monkeypatch.setattr(
|
|
download_manager, "get_downloader", AsyncMock(return_value=DummyDownloader())
|
|
)
|
|
|
|
class DummyScanner:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def add_model_to_cache(self, metadata_dict, relative_path):
|
|
self.calls.append((metadata_dict, relative_path))
|
|
|
|
dummy_scanner = DummyScanner()
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_get_lora_scanner", AsyncMock(return_value=dummy_scanner)
|
|
)
|
|
monkeypatch.setattr(MetadataManager, "save_metadata", AsyncMock(return_value=True))
|
|
|
|
version_info = {"images": []}
|
|
|
|
task_a = asyncio.create_task(
|
|
manager._execute_download(
|
|
download_urls=["https://civitai.example/api/download/models/2665422"],
|
|
save_dir=str(save_dir),
|
|
metadata=_SharedTargetMetadata(target_path),
|
|
version_info=version_info,
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id=None,
|
|
transfer_backend="python",
|
|
)
|
|
)
|
|
task_b = asyncio.create_task(
|
|
manager._execute_download(
|
|
download_urls=["https://civitai.example/api/download/models/2925310"],
|
|
save_dir=str(save_dir),
|
|
metadata=_SharedTargetMetadata(target_path),
|
|
version_info=version_info,
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id=None,
|
|
transfer_backend="python",
|
|
)
|
|
)
|
|
|
|
# Wait until task A reached the metadata-build seam (file written, download
|
|
# reported success), letting task B's failure cleanup run meanwhile.
|
|
await asyncio.wait_for(build_started.wait(), timeout=5.0)
|
|
release_build.set()
|
|
|
|
result_b = await asyncio.wait_for(task_b, timeout=10.0)
|
|
result_a = await asyncio.wait_for(task_a, timeout=10.0)
|
|
|
|
# Task A must succeed: its completed file must not be deleted by task B.
|
|
assert result_a.get("success") is True, (
|
|
f"Surviving download failed due to the same-path race: {result_a}"
|
|
)
|
|
assert target_path.exists(), (
|
|
"Task B's failure cleanup deleted task A's completed download"
|
|
)
|
|
assert target_path.read_text() == "model-a-content"
|
|
|
|
# Task B legitimately fails (simulated transfer error)...
|
|
assert result_b.get("success") is False
|
|
# ...but only after being serialized behind A, resolving to a unique
|
|
# filename instead of colliding on A's target path.
|
|
b_targets = [p for p in downloader_paths if "2925310" in str(p) or True]
|
|
assert b_targets, "Task B never attempted a download"
|
|
assert all(os.path.basename(p) != target_path.name for p in [downloader_paths[-1]]) or (
|
|
os.path.basename(downloader_paths[-1]) != target_path.name
|
|
), (
|
|
f"Task B reused task A's exact target path instead of a unique one: "
|
|
f"{downloader_paths}"
|
|
)
|
|
|
|
|
|
def test_reconcile_failed_aria2_partial_preserves_resumable_pair(tmp_path):
|
|
"""Payload and .aria2 control file are a resumable pair: keep both."""
|
|
save_path = tmp_path / "model.safetensors"
|
|
save_path.write_text("partial")
|
|
control_path = tmp_path / "model.safetensors.aria2"
|
|
control_path.write_text("control")
|
|
|
|
DownloadManager._reconcile_failed_aria2_partial(str(save_path))
|
|
|
|
assert save_path.exists()
|
|
assert control_path.exists()
|
|
|
|
|
|
def test_reconcile_failed_aria2_partial_preserves_likely_complete_file(tmp_path):
|
|
"""Payload without a control file means aria2 finished it: keep it."""
|
|
save_path = tmp_path / "model.safetensors"
|
|
save_path.write_text("complete")
|
|
|
|
DownloadManager._reconcile_failed_aria2_partial(str(save_path))
|
|
|
|
assert save_path.exists()
|
|
assert not (tmp_path / "model.safetensors.aria2").exists()
|
|
|
|
|
|
def test_reconcile_failed_aria2_partial_removes_orphaned_control_file(tmp_path):
|
|
"""A control file without its payload cannot resume: remove the orphan."""
|
|
save_path = tmp_path / "model.safetensors"
|
|
control_path = tmp_path / "model.safetensors.aria2"
|
|
control_path.write_text("control")
|
|
|
|
DownloadManager._reconcile_failed_aria2_partial(str(save_path))
|
|
|
|
assert not control_path.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_download_preserves_aria2_partial_pair_on_failure(
|
|
monkeypatch, tmp_path
|
|
):
|
|
"""A failed aria2 transfer keeps the payload and control file for resume."""
|
|
manager = DownloadManager()
|
|
settings = get_settings_manager()
|
|
settings.settings["download_backend"] = "aria2"
|
|
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
target_path = save_dir / "file.safetensors"
|
|
|
|
class DummyMetadata:
|
|
def __init__(self, path: Path):
|
|
self.file_path = str(path)
|
|
self.sha256 = "sha256"
|
|
self.file_name = path.stem
|
|
self.preview_url = None
|
|
self.autov3: Optional[str] = None
|
|
|
|
def generate_unique_filename(self, *_args, **_kwargs):
|
|
return os.path.basename(self.file_path)
|
|
|
|
def update_file_info(self, _path):
|
|
return None
|
|
|
|
def to_dict(self):
|
|
return {"file_path": self.file_path}
|
|
|
|
class FailingAria2Downloader:
|
|
async def download_file(
|
|
self,
|
|
url,
|
|
save_path,
|
|
*,
|
|
download_id,
|
|
progress_callback=None,
|
|
headers=None,
|
|
):
|
|
Path(save_path).write_text("partial")
|
|
Path(f"{save_path}.aria2").write_text("control")
|
|
return False, "No URI available."
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=FailingAria2Downloader()),
|
|
)
|
|
|
|
result = await manager._execute_download(
|
|
download_urls=["https://civitai.com/api/download/models/1"],
|
|
save_dir=str(save_dir),
|
|
metadata=DummyMetadata(target_path),
|
|
version_info={"images": []},
|
|
relative_path="",
|
|
progress_callback=None,
|
|
model_type="lora",
|
|
download_id="download-fail",
|
|
)
|
|
|
|
assert result["success"] is False
|
|
assert result["error"] == "No URI available."
|
|
assert target_path.exists()
|
|
assert (save_dir / "file.safetensors.aria2").exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_persisted_downloads_removes_orphaned_aria2_control_file(
|
|
monkeypatch, tmp_path
|
|
):
|
|
"""A persisted record whose .aria2 control file lost its payload is
|
|
dropped, and the orphaned control file is removed."""
|
|
manager = DownloadManager()
|
|
download_id = "download-orphan"
|
|
save_path = tmp_path / "file.safetensors"
|
|
control_path = tmp_path / "file.safetensors.aria2"
|
|
control_path.write_text("control")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
download_id,
|
|
{
|
|
"download_id": download_id,
|
|
"transfer_backend": "aria2",
|
|
"status": "failed",
|
|
"save_path": str(save_path),
|
|
"file_path": str(save_path),
|
|
},
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=SimpleNamespace()),
|
|
)
|
|
|
|
await manager._restore_persisted_downloads()
|
|
|
|
assert not control_path.exists()
|
|
assert await manager._aria2_state_store.get(download_id) is None
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Structured 429 rate-limit handling (download queue contract)
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
def _prepare_tracked_download(manager, download_id, status="downloading"):
|
|
manager._active_downloads[download_id] = {
|
|
"transfer_backend": "python",
|
|
"status": status,
|
|
"bytes_per_second": 0.0,
|
|
}
|
|
manager._pause_events[download_id] = DownloadStreamControl()
|
|
|
|
|
|
async def _run_download(manager, download_id, tmp_path, source=None):
|
|
return await manager._download_with_semaphore(
|
|
download_id,
|
|
1,
|
|
None,
|
|
str(tmp_path),
|
|
"",
|
|
None,
|
|
False,
|
|
source,
|
|
None,
|
|
False,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rate_limited_download_is_requeued_not_failed(
|
|
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
|
):
|
|
"""A 429 during the transfer re-queues the item with a structured result."""
|
|
manager = DownloadManager()
|
|
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
|
|
|
download_id = "dl-429"
|
|
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
|
await queue_service.update_status(download_id, "downloading")
|
|
_prepare_tracked_download(manager, download_id)
|
|
|
|
monkeypatch.setattr(
|
|
manager,
|
|
"_execute_original_download",
|
|
AsyncMock(
|
|
side_effect=DownloadRateLimitError(
|
|
"Download rate limited (429), retry after 120s",
|
|
retry_after=120,
|
|
host="civitai.com",
|
|
)
|
|
),
|
|
)
|
|
|
|
result = await _run_download(manager, download_id, tmp_path)
|
|
|
|
assert result["success"] is False
|
|
assert result["reason"] == "rate_limited"
|
|
assert result["retry_after"] == 120
|
|
assert "rate limited" in result["error"].lower()
|
|
|
|
# Back to "queued" — NOT moved to history as failed.
|
|
queue = await queue_service.get_queue()
|
|
assert len(queue) == 1
|
|
assert queue[0]["status"] == "queued"
|
|
history = await queue_service.get_history()
|
|
assert history["items"] == []
|
|
assert manager._active_downloads[download_id]["status"] == "queued"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_preflight_gate_blocks_when_host_in_cooldown(
|
|
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
|
):
|
|
"""Host in cooldown: immediate structured 429, no transfer, no slot held."""
|
|
manager = DownloadManager()
|
|
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
|
|
|
coordinator = await RateLimitCoordinator.get_instance()
|
|
coordinator.register_rate_limit("civitai.com", 300)
|
|
|
|
download_id = "dl-cooldown"
|
|
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
|
_prepare_tracked_download(manager, download_id, status="waiting")
|
|
|
|
execute = AsyncMock(
|
|
side_effect=AssertionError("download must not start during cooldown")
|
|
)
|
|
monkeypatch.setattr(manager, "_execute_original_download", execute)
|
|
|
|
result = await _run_download(manager, download_id, tmp_path)
|
|
|
|
assert execute.await_count == 0
|
|
assert result["success"] is False
|
|
assert result["reason"] == "rate_limited"
|
|
assert 290 <= result["retry_after"] <= 300
|
|
|
|
queue = await queue_service.get_queue()
|
|
assert len(queue) == 1
|
|
assert queue[0]["status"] == "queued"
|
|
history = await queue_service.get_history()
|
|
assert history["items"] == []
|
|
|
|
# The concurrency slot was never occupied by the gated download.
|
|
await asyncio.wait_for(manager._download_semaphore.acquire(), timeout=0.1)
|
|
manager._download_semaphore.release()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_preflight_gate_ignores_civarchive_cooldown_for_civitai_downloads(
|
|
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
|
):
|
|
"""A civarchive.com cooldown (armed e.g. by bulk metadata fetches) must
|
|
not block a plain CivitAI download that never touches that host."""
|
|
manager = DownloadManager()
|
|
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
|
|
|
coordinator = await RateLimitCoordinator.get_instance()
|
|
coordinator.register_rate_limit("civarchive.com", 1800)
|
|
|
|
download_id = "dl-civarchive-cooldown"
|
|
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
|
_prepare_tracked_download(manager, download_id, status="waiting")
|
|
|
|
execute = AsyncMock(return_value={"success": True})
|
|
monkeypatch.setattr(manager, "_execute_original_download", execute)
|
|
|
|
result = await _run_download(manager, download_id, tmp_path)
|
|
|
|
assert execute.await_count == 1
|
|
assert result["success"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_preflight_gate_blocks_civarchive_source_during_civarchive_cooldown(
|
|
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
|
):
|
|
"""Downloads explicitly sourced from CivArchive still gate on its cooldown."""
|
|
manager = DownloadManager()
|
|
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
|
|
|
coordinator = await RateLimitCoordinator.get_instance()
|
|
coordinator.register_rate_limit("civarchive.com", 1800)
|
|
|
|
download_id = "dl-civarchive-source"
|
|
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
|
_prepare_tracked_download(manager, download_id, status="waiting")
|
|
|
|
execute = AsyncMock(
|
|
side_effect=AssertionError("download must not start during cooldown")
|
|
)
|
|
monkeypatch.setattr(manager, "_execute_original_download", execute)
|
|
|
|
result = await _run_download(manager, download_id, tmp_path, source="civarchive")
|
|
|
|
assert execute.await_count == 0
|
|
assert result["success"] is False
|
|
assert result["reason"] == "rate_limited"
|
|
assert "civarchive.com" in result["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rate_limit_retry_after_falls_back_to_coordinator_backoff(
|
|
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
|
):
|
|
"""When the 429 carries no Retry-After, the coordinator backoff supplies it."""
|
|
manager = DownloadManager()
|
|
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
|
|
|
# Cooldown on a mirror host that is NOT in the pre-flight host list, so
|
|
# the failure must come from the transfer itself.
|
|
coordinator = await RateLimitCoordinator.get_instance()
|
|
backoff = coordinator.register_rate_limit("mirror.example.com", None)
|
|
assert backoff == 30.0
|
|
|
|
download_id = "dl-429-mirror"
|
|
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
|
_prepare_tracked_download(manager, download_id)
|
|
|
|
monkeypatch.setattr(
|
|
manager,
|
|
"_execute_original_download",
|
|
AsyncMock(
|
|
side_effect=DownloadRateLimitError(
|
|
"Download rate limited (429)",
|
|
retry_after=None,
|
|
host="mirror.example.com",
|
|
)
|
|
),
|
|
)
|
|
|
|
result = await _run_download(manager, download_id, tmp_path)
|
|
|
|
assert result["success"] is False
|
|
assert result["reason"] == "rate_limited"
|
|
assert 25 <= result["retry_after"] <= 30
|
|
queue = await queue_service.get_queue()
|
|
assert queue[0]["status"] == "queued"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_metadata_fetch_rate_limit_error_is_also_requeued(
|
|
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
|
):
|
|
"""A plain RateLimitError (metadata fetch path) gets the same treatment."""
|
|
manager = DownloadManager()
|
|
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
|
|
|
download_id = "dl-meta-429"
|
|
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
|
_prepare_tracked_download(manager, download_id)
|
|
|
|
monkeypatch.setattr(
|
|
manager,
|
|
"_execute_original_download",
|
|
AsyncMock(
|
|
side_effect=RateLimitError(
|
|
"Request rate limited", retry_after=45, provider="civitai_api"
|
|
)
|
|
),
|
|
)
|
|
|
|
result = await _run_download(manager, download_id, tmp_path)
|
|
|
|
assert result["success"] is False
|
|
assert result["reason"] == "rate_limited"
|
|
assert result["retry_after"] == 45
|
|
queue = await queue_service.get_queue()
|
|
assert queue[0]["status"] == "queued"
|
|
history = await queue_service.get_history()
|
|
assert history["items"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generic_failure_still_moves_to_history_as_failed(
|
|
monkeypatch, tmp_path, queue_service, reset_rate_limit_coordinator
|
|
):
|
|
"""Non-rate-limit failures keep the legacy move-to-history behavior."""
|
|
manager = DownloadManager()
|
|
monkeypatch.setattr(manager, "_cleanup_download_record", AsyncMock())
|
|
|
|
download_id = "dl-generic-fail"
|
|
await queue_service.add_to_queue(download_id=download_id, model_id=1)
|
|
_prepare_tracked_download(manager, download_id)
|
|
|
|
monkeypatch.setattr(
|
|
manager,
|
|
"_execute_original_download",
|
|
AsyncMock(side_effect=RuntimeError("disk full")),
|
|
)
|
|
|
|
result = await _run_download(manager, download_id, tmp_path)
|
|
|
|
assert result == {"success": False, "error": "disk full"}
|
|
queue = await queue_service.get_queue()
|
|
assert queue == []
|
|
history = await queue_service.get_history()
|
|
assert len(history["items"]) == 1
|
|
assert history["items"][0]["status"] == "failed"
|