mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-23 12:04:09 -03:00
Will Miao
74736f7560
Follow-up to the keyword-dump guard. The reported model's tag list is ["lora, character, ... face", "base model"], so skipping the dump left the "base model" label to be picked as the folder name. That label describes Civitai's listing rather than the model's content, which makes it as meaningless as a folder as the blob was. Add CIVITAI_META_TAGS and is_civitai_meta_tag(), and skip those labels in the automatic fallback. An explicit priority entry still matches them, so a user who does want a "base model" folder can configure one. The reported model now resolves to "Krea 2/no tags" instead of "Krea 2/base model". Also correct a comment that listed ".civitai.info" among the files sitting next to a model. LoRA Manager only reads that sidecar -- other tools write it -- and writes ".metadata.json" itself.
2317 lines
73 KiB
Python
2317 lines
73 KiB
Python
"""Core functionality tests for DownloadManager."""
|
|
|
|
import asyncio
|
|
import os
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from py.services.download_manager import DownloadManager
|
|
from py.services import download_manager
|
|
from py.services import aria2_transfer_state
|
|
from py.services.service_registry import ServiceRegistry
|
|
from py.services.settings_manager import SettingsManager, get_settings_manager
|
|
|
|
|
|
@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"},
|
|
"skip_previously_downloaded_model_versions": False,
|
|
"download_skip_base_models": [],
|
|
}
|
|
)
|
|
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(autouse=True)
|
|
def stub_metadata(monkeypatch):
|
|
class _StubMetadata:
|
|
def __init__(self, save_path: str):
|
|
self.file_path = save_path
|
|
self.sha256 = "sha256"
|
|
self.file_name = Path(save_path).stem
|
|
|
|
def _factory(save_path: str):
|
|
return _StubMetadata(save_path)
|
|
|
|
def _make_class():
|
|
@staticmethod
|
|
def from_civitai_info(_version_info, _file_info, save_path):
|
|
return _factory(save_path)
|
|
|
|
return type("StubMetadata", (), {"from_civitai_info": from_civitai_info})
|
|
|
|
stub_class = _make_class()
|
|
monkeypatch.setattr(download_manager, "LoraMetadata", stub_class)
|
|
monkeypatch.setattr(download_manager, "CheckpointMetadata", stub_class)
|
|
monkeypatch.setattr(download_manager, "EmbeddingMetadata", stub_class)
|
|
|
|
|
|
class DummyScanner:
|
|
def __init__(self, exists: bool = False, raw_data=None):
|
|
self.exists = exists
|
|
self.calls = []
|
|
self._cache = SimpleNamespace(raw_data=list(raw_data or []))
|
|
|
|
async def check_model_version_exists(self, version_id):
|
|
self.calls.append(version_id)
|
|
return self.exists
|
|
|
|
async def get_cached_data(self):
|
|
return self._cache
|
|
|
|
|
|
@pytest.fixture
|
|
def scanners(monkeypatch):
|
|
lora_scanner = DummyScanner()
|
|
checkpoint_scanner = DummyScanner()
|
|
embedding_scanner = DummyScanner()
|
|
|
|
monkeypatch.setattr(
|
|
ServiceRegistry, "get_lora_scanner", AsyncMock(return_value=lora_scanner)
|
|
)
|
|
monkeypatch.setattr(
|
|
ServiceRegistry,
|
|
"get_checkpoint_scanner",
|
|
AsyncMock(return_value=checkpoint_scanner),
|
|
)
|
|
monkeypatch.setattr(
|
|
ServiceRegistry,
|
|
"get_embedding_scanner",
|
|
AsyncMock(return_value=embedding_scanner),
|
|
)
|
|
|
|
return SimpleNamespace(
|
|
lora=lora_scanner,
|
|
checkpoint=checkpoint_scanner,
|
|
embedding=embedding_scanner,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def metadata_provider(monkeypatch):
|
|
class DummyProvider:
|
|
def __init__(self):
|
|
self.calls = []
|
|
self.payload = {
|
|
"id": 42,
|
|
"model": {"type": "LoRA", "tags": ["fantasy"]},
|
|
"baseModel": "BaseModel",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"type": "Model",
|
|
"primary": True,
|
|
"downloadUrl": "https://example.invalid/file.safetensors",
|
|
"name": "file.safetensors",
|
|
}
|
|
],
|
|
}
|
|
|
|
async def get_model_version(self, model_id, model_version_id):
|
|
self.calls.append((model_id, model_version_id))
|
|
return self.payload
|
|
|
|
provider = DummyProvider()
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_default_metadata_provider",
|
|
AsyncMock(return_value=provider),
|
|
)
|
|
return provider
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def noop_cleanup(monkeypatch):
|
|
async def _cleanup(self, task_id):
|
|
if task_id in self._active_downloads:
|
|
self._active_downloads[task_id]["cleaned"] = True
|
|
|
|
monkeypatch.setattr(DownloadManager, "_cleanup_download_record", _cleanup)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_requires_identifier():
|
|
"""Test that download fails when no identifier is provided."""
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai()
|
|
assert result == {
|
|
"success": False,
|
|
"error": "Either model_id or model_version_id must be provided",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_successful_download_uses_defaults(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""Test successful download with default settings."""
|
|
manager = DownloadManager()
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(
|
|
self,
|
|
*,
|
|
download_urls,
|
|
save_dir,
|
|
metadata,
|
|
version_info,
|
|
relative_path,
|
|
progress_callback,
|
|
model_type,
|
|
download_id,
|
|
transfer_backend=None,
|
|
):
|
|
captured.update(
|
|
{
|
|
"download_urls": download_urls,
|
|
"save_dir": Path(save_dir),
|
|
"relative_path": relative_path,
|
|
"progress_callback": progress_callback,
|
|
"model_type": model_type,
|
|
"download_id": download_id,
|
|
"metadata_path": metadata.file_path,
|
|
}
|
|
)
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=99,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert "download_id" in result
|
|
assert manager._download_tasks == {}
|
|
assert manager._active_downloads[result["download_id"]]["status"] == "completed"
|
|
|
|
assert captured["relative_path"] == "MappedModel/fantasy"
|
|
expected_dir = (
|
|
Path(get_settings_manager().get("default_lora_root"))
|
|
/ "MappedModel"
|
|
/ "fantasy"
|
|
)
|
|
assert captured["save_dir"] == expected_dir
|
|
assert captured["model_type"] == "lora"
|
|
assert captured["download_urls"] == ["https://example.invalid/file.safetensors"]
|
|
|
|
|
|
def test_calculate_relative_path_ignores_keyword_dump_tag():
|
|
"""The #1119 download flow: the real tag list must not become a folder.
|
|
|
|
The model's only two tags are the keyword dump and Civitai's "base model"
|
|
label, so nothing usable is left and the template falls back to "no tags".
|
|
"""
|
|
keyword_dump = (
|
|
"lora, character, rosie, irish, redhead, auburn, freckles, green eyes, "
|
|
"curly hair, woman, female, photorealistic, realistic, krea2, dark beast, "
|
|
"kreativity, nsfw, nude, portrait, face"
|
|
)
|
|
manager = DownloadManager()
|
|
|
|
relative_path = manager._calculate_relative_path(
|
|
{
|
|
"baseModel": "BaseModel",
|
|
"creator": {"username": "mad_macs"},
|
|
"name": "v1.2",
|
|
"model": {"name": "Rosie", "tags": [keyword_dump, "base model"]},
|
|
},
|
|
"lora",
|
|
)
|
|
|
|
assert relative_path == "MappedModel/no tags"
|
|
assert keyword_dump not in relative_path
|
|
assert len(relative_path) < 50
|
|
|
|
|
|
def test_calculate_relative_path_sanitizes_tag_segment():
|
|
"""A tag with path separators must not create nested folders."""
|
|
manager = DownloadManager()
|
|
|
|
relative_path = manager._calculate_relative_path(
|
|
{
|
|
"baseModel": "BaseModel",
|
|
"creator": {"username": "author"},
|
|
"name": "v1.2",
|
|
"model": {"name": "Rosie", "tags": ["a/b:c"]},
|
|
},
|
|
"lora",
|
|
)
|
|
|
|
assert relative_path == "MappedModel/a_b_c"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_accepts_enhancement_lora_primary_file(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""A version whose only file has type 'Enhancement LoRA' (Anima/AIR
|
|
image-editing LoRAs) must download — previously failed with
|
|
"No suitable file found in metadata" because the type was missing from
|
|
the primary-file weights allowlist."""
|
|
manager = DownloadManager()
|
|
metadata_provider.payload = {
|
|
"id": 3219121,
|
|
"model": {"type": "LORA", "tags": ["style"]},
|
|
"baseModel": "Anima",
|
|
"creator": {"username": "Deskup"},
|
|
"files": [
|
|
{
|
|
"id": 3100968,
|
|
"type": "Enhancement LoRA",
|
|
"primary": True,
|
|
"name": "deskup-anima-edit-general.safetensors",
|
|
"sizeKB": 358501.13,
|
|
"downloadUrl": "https://example.invalid/deskup-anima-edit-general.safetensors",
|
|
}
|
|
],
|
|
}
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(self, **kwargs):
|
|
captured.update(
|
|
{
|
|
"download_urls": kwargs["download_urls"],
|
|
"model_type": kwargs["model_type"],
|
|
}
|
|
)
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_id=2850692,
|
|
model_version_id=3219121,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert captured["model_type"] == "lora"
|
|
assert captured["download_urls"] == [
|
|
"https://example.invalid/deskup-anima-edit-general.safetensors"
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_falls_back_to_civitai_primary_flag_regardless_of_type(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""If no weights-type file exists, trust CivitAI's `primary` flag on any
|
|
file — mirrors CivitAI's getPrimaryFile() which never excludes a file by
|
|
type."""
|
|
manager = DownloadManager()
|
|
metadata_provider.payload = {
|
|
"id": 77,
|
|
"model": {"type": "LORA", "tags": ["concept"]},
|
|
"baseModel": "Anima",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"id": 100,
|
|
"type": "Other",
|
|
"primary": True,
|
|
"name": "custom-type-lora.safetensors",
|
|
"downloadUrl": "https://example.invalid/custom-type-lora.safetensors",
|
|
}
|
|
],
|
|
}
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(self, **kwargs):
|
|
captured["download_urls"] = kwargs["download_urls"]
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=77,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert captured["download_urls"] == [
|
|
"https://example.invalid/custom-type-lora.safetensors"
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_prefers_weights_file_over_non_weights_primary(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""A Config/Archive-type primary must never replace an existing weights
|
|
file — the weights file wins even without the primary flag."""
|
|
manager = DownloadManager()
|
|
metadata_provider.payload = {
|
|
"id": 78,
|
|
"model": {"type": "LORA", "tags": ["concept"]},
|
|
"baseModel": "BaseModel",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"id": 201,
|
|
"type": "Config",
|
|
"primary": True,
|
|
"name": "config.json",
|
|
"downloadUrl": "https://example.invalid/config.json",
|
|
},
|
|
{
|
|
"id": 202,
|
|
"type": "Model",
|
|
"primary": False,
|
|
"name": "weights.safetensors",
|
|
"downloadUrl": "https://example.invalid/weights.safetensors",
|
|
},
|
|
],
|
|
}
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(self, **kwargs):
|
|
captured["download_urls"] = kwargs["download_urls"]
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=78,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert captured["download_urls"] == [
|
|
"https://example.invalid/weights.safetensors"
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_keeps_save_dir_when_use_save_dir_as_root(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""use_default_paths with use_save_dir_as_root resolves the template under
|
|
the provided save_dir instead of switching to the default root."""
|
|
manager = DownloadManager()
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(
|
|
self,
|
|
*,
|
|
download_urls,
|
|
save_dir,
|
|
metadata,
|
|
version_info,
|
|
relative_path,
|
|
progress_callback,
|
|
model_type,
|
|
download_id,
|
|
transfer_backend=None,
|
|
):
|
|
captured.update(
|
|
{
|
|
"save_dir": Path(save_dir),
|
|
"relative_path": relative_path,
|
|
"model_type": model_type,
|
|
}
|
|
)
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
custom_root = tmp_path / "custom_root"
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=99,
|
|
save_dir=str(custom_root),
|
|
use_default_paths=True,
|
|
use_save_dir_as_root=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert captured["relative_path"] == "MappedModel/fantasy"
|
|
assert captured["save_dir"] == custom_root / "MappedModel" / "fantasy"
|
|
assert captured["model_type"] == "lora"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_successful_download_schedules_auto_example_images(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
scheduled = []
|
|
|
|
async def fake_execute_download(
|
|
self,
|
|
*,
|
|
download_urls,
|
|
save_dir,
|
|
metadata,
|
|
version_info,
|
|
relative_path,
|
|
progress_callback,
|
|
model_type,
|
|
download_id,
|
|
transfer_backend=None,
|
|
):
|
|
return {"success": True}
|
|
|
|
async def fake_schedule(self, *, metadata, model_type):
|
|
scheduled.append({"metadata": metadata, "model_type": model_type})
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_schedule_auto_example_images_download",
|
|
fake_schedule,
|
|
raising=False,
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=99,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert len(scheduled) == 1
|
|
assert scheduled[0]["model_type"] == "lora"
|
|
assert scheduled[0]["metadata"].sha256 == "sha256"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_example_images_download_uses_settings_payload(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
settings = get_settings_manager()
|
|
settings.settings["auto_download_example_images"] = True
|
|
settings.settings["example_images_path"] = str(tmp_path / "examples")
|
|
settings.settings["optimize_example_images"] = False
|
|
|
|
calls = []
|
|
|
|
class DummyExampleImagesManager:
|
|
async def start_force_download(self, payload):
|
|
calls.append(payload)
|
|
return {"success": True}
|
|
|
|
from py.utils import example_images_download_manager
|
|
|
|
monkeypatch.setattr(
|
|
ServiceRegistry,
|
|
"get_websocket_manager",
|
|
AsyncMock(return_value=object()),
|
|
)
|
|
monkeypatch.setattr(
|
|
example_images_download_manager,
|
|
"get_default_download_manager",
|
|
lambda _ws_manager: DummyExampleImagesManager(),
|
|
)
|
|
|
|
metadata = SimpleNamespace(sha256="ABCDEF", file_path="model.safetensors")
|
|
await manager._schedule_auto_example_images_download(
|
|
metadata=metadata,
|
|
model_type="lora",
|
|
)
|
|
|
|
for _ in range(10):
|
|
if calls:
|
|
break
|
|
await asyncio.sleep(0)
|
|
|
|
assert calls == [
|
|
{
|
|
"model_hashes": ["abcdef"],
|
|
"optimize": False,
|
|
"model_types": ["lora"],
|
|
"delay": 0,
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_example_images_download_skips_without_configuration(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
settings = get_settings_manager()
|
|
settings.settings["auto_download_example_images"] = True
|
|
settings.settings["example_images_path"] = ""
|
|
|
|
get_ws_manager = AsyncMock(return_value=object())
|
|
monkeypatch.setattr(ServiceRegistry, "get_websocket_manager", get_ws_manager)
|
|
|
|
await manager._schedule_auto_example_images_download(
|
|
metadata=SimpleNamespace(sha256="abcdef", file_path="model.safetensors"),
|
|
model_type="lora",
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
get_ws_manager.assert_not_called()
|
|
|
|
settings.settings["example_images_path"] = str(tmp_path / "examples")
|
|
await manager._schedule_auto_example_images_download(
|
|
metadata=SimpleNamespace(sha256="", file_path="model.safetensors"),
|
|
model_type="lora",
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
get_ws_manager.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_uses_active_mirrors(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""Test that active mirrors are used when available."""
|
|
manager = DownloadManager()
|
|
|
|
metadata_with_mirrors = {
|
|
"id": 42,
|
|
"model": {"type": "LoRA", "tags": ["fantasy"]},
|
|
"baseModel": "BaseModel",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"type": "Model",
|
|
"primary": True,
|
|
"downloadUrl": "https://example.invalid/file.safetensors",
|
|
"mirrors": [
|
|
{
|
|
"url": "https://mirror.example/file.safetensors",
|
|
"deletedAt": None,
|
|
},
|
|
{
|
|
"url": "https://mirror.example/old.safetensors",
|
|
"deletedAt": "2024-01-01",
|
|
},
|
|
],
|
|
"name": "file.safetensors",
|
|
}
|
|
],
|
|
}
|
|
|
|
metadata_provider.get_model_version = AsyncMock(return_value=metadata_with_mirrors)
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(
|
|
self,
|
|
*,
|
|
download_urls,
|
|
save_dir,
|
|
metadata,
|
|
version_info,
|
|
relative_path,
|
|
progress_callback,
|
|
model_type,
|
|
download_id,
|
|
transfer_backend=None,
|
|
):
|
|
captured["download_urls"] = download_urls
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=99,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert captured["download_urls"] == ["https://mirror.example/file.safetensors"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_resume_cancel_delegate_to_aria2_backend(monkeypatch):
|
|
manager = DownloadManager()
|
|
|
|
task = asyncio.create_task(asyncio.sleep(60))
|
|
manager._download_tasks["download-1"] = task
|
|
manager._pause_events["download-1"] = download_manager.DownloadStreamControl()
|
|
manager._active_downloads["download-1"] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "downloading",
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def pause_download(self, download_id):
|
|
self.calls.append(("pause", download_id))
|
|
return {"success": True, "message": "paused"}
|
|
|
|
async def resume_download(self, download_id):
|
|
self.calls.append(("resume", download_id))
|
|
return {"success": True, "message": "resumed"}
|
|
|
|
async def cancel_download(self, download_id):
|
|
self.calls.append(("cancel", download_id))
|
|
return {"success": True, "message": "cancelled"}
|
|
|
|
async def has_transfer(self, download_id):
|
|
self.calls.append(("has_transfer", download_id))
|
|
return True
|
|
|
|
dummy_aria2 = DummyAria2Downloader()
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=dummy_aria2),
|
|
)
|
|
|
|
pause_result = await manager.pause_download("download-1")
|
|
assert pause_result["success"] is True
|
|
assert manager._active_downloads["download-1"]["status"] == "paused"
|
|
|
|
resume_result = await manager.resume_download("download-1")
|
|
assert resume_result["success"] is True
|
|
assert manager._active_downloads["download-1"]["status"] == "downloading"
|
|
|
|
cancel_result = await manager.cancel_download("download-1")
|
|
assert cancel_result["success"] is True
|
|
assert task.cancelled() or task.done()
|
|
assert dummy_aria2.calls == [
|
|
("has_transfer", "download-1"),
|
|
("pause", "download-1"),
|
|
("has_transfer", "download-1"),
|
|
("resume", "download-1"),
|
|
("cancel", "download-1"),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_allows_queued_aria2_task_without_transfer(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()
|
|
|
|
manager._download_tasks["download-queued"] = task
|
|
manager._active_downloads["download-queued"] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "queued",
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
async def cancel_download(self, download_id):
|
|
return {"success": False, "error": "Download task not found"}
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
result = await manager.cancel_download("download-queued")
|
|
|
|
assert result["success"] is True
|
|
assert task.cancelled() or task.done()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_resume_queued_aria2_task_without_transfer(monkeypatch):
|
|
manager = DownloadManager()
|
|
|
|
task = asyncio.create_task(asyncio.sleep(60))
|
|
manager._download_tasks["download-queued"] = task
|
|
manager._pause_events["download-queued"] = download_manager.DownloadStreamControl()
|
|
manager._active_downloads["download-queued"] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "waiting",
|
|
"bytes_per_second": 12.0,
|
|
}
|
|
|
|
class DummyAria2Downloader:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def has_transfer(self, download_id):
|
|
self.calls.append(("has_transfer", download_id))
|
|
return False
|
|
|
|
async def pause_download(self, download_id):
|
|
self.calls.append(("pause", download_id))
|
|
return {"success": True, "message": "paused"}
|
|
|
|
async def resume_download(self, download_id):
|
|
self.calls.append(("resume", download_id))
|
|
return {"success": True, "message": "resumed"}
|
|
|
|
dummy_aria2 = DummyAria2Downloader()
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=dummy_aria2),
|
|
)
|
|
|
|
pause_result = await manager.pause_download("download-queued")
|
|
assert pause_result == {"success": True, "message": "Download paused successfully"}
|
|
assert manager._active_downloads["download-queued"]["status"] == "paused"
|
|
assert manager._pause_events["download-queued"].is_paused() is True
|
|
|
|
resume_result = await manager.resume_download("download-queued")
|
|
assert resume_result == {"success": True, "message": "Download resumed successfully"}
|
|
assert manager._active_downloads["download-queued"]["status"] == "downloading"
|
|
assert manager._pause_events["download-queued"].is_set() is True
|
|
assert dummy_aria2.calls == [
|
|
("has_transfer", "download-queued"),
|
|
("has_transfer", "download-queued"),
|
|
]
|
|
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_download_restores_persisted_aria2_task(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
save_path = save_dir / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
(save_dir / "file.safetensors.aria2").write_text("control")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"download-1",
|
|
{
|
|
"download_id": "download-1",
|
|
"transfer_backend": "aria2",
|
|
"status": "paused",
|
|
"save_dir": str(save_dir),
|
|
"relative_path": "",
|
|
"use_default_paths": False,
|
|
"save_path": str(save_path),
|
|
"file_path": str(save_path),
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
},
|
|
)
|
|
|
|
created = {}
|
|
|
|
async def fake_download_with_semaphore(
|
|
self,
|
|
task_id,
|
|
model_id,
|
|
model_version_id,
|
|
save_dir,
|
|
relative_path,
|
|
progress_callback=None,
|
|
use_default_paths=False,
|
|
source=None,
|
|
file_params=None,
|
|
use_save_dir_as_root=False,
|
|
):
|
|
created.update(
|
|
{
|
|
"task_id": task_id,
|
|
"model_id": model_id,
|
|
"model_version_id": model_version_id,
|
|
"save_dir": save_dir,
|
|
}
|
|
)
|
|
return {"success": True}
|
|
|
|
class DummyAria2Downloader:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def get_status_by_gid(self, gid):
|
|
return None
|
|
|
|
async def has_transfer(self, download_id):
|
|
self.calls.append(("has_transfer", download_id))
|
|
return False
|
|
|
|
async def resume_download(self, download_id):
|
|
self.calls.append(("resume", download_id))
|
|
return {"success": True, "message": "resumed"}
|
|
|
|
async def restore_transfer(self, download_id, gid, save_path):
|
|
self.calls.append(("restore_transfer", download_id, gid, save_path))
|
|
|
|
dummy_aria2 = DummyAria2Downloader()
|
|
monkeypatch.setattr(
|
|
download_manager, "_download_with_semaphore", None, raising=False
|
|
)
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_download_with_semaphore",
|
|
fake_download_with_semaphore,
|
|
)
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=dummy_aria2),
|
|
)
|
|
|
|
result = await manager.resume_download("download-1")
|
|
await asyncio.sleep(0)
|
|
|
|
assert result == {"success": True, "message": "Download resumed successfully"}
|
|
assert created["task_id"] == "download-1"
|
|
assert created["model_version_id"] == 34
|
|
assert manager._active_downloads["download-1"]["status"] == "downloading"
|
|
assert manager._pause_events["download-1"].is_set() is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_active_downloads_restores_persisted_aria2_entries(monkeypatch, tmp_path):
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
save_path = save_dir / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
(save_dir / "file.safetensors.aria2").write_text("control")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"download-1",
|
|
{
|
|
"download_id": "download-1",
|
|
"transfer_backend": "aria2",
|
|
"status": "paused",
|
|
"save_path": str(save_path),
|
|
"file_path": str(save_path),
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
},
|
|
)
|
|
|
|
class DummyAria2Downloader:
|
|
async def get_status_by_gid(self, gid):
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
downloads = await manager.get_active_downloads()
|
|
|
|
assert downloads["downloads"] == [
|
|
{
|
|
"download_id": "download-1",
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"progress": 0,
|
|
"status": "paused",
|
|
"error": None,
|
|
"bytes_downloaded": 0,
|
|
"total_bytes": None,
|
|
"bytes_per_second": 0.0,
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_active_downloads_restores_orphaned_aria2_partial_as_paused(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
save_path = save_dir / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
(save_dir / "file.safetensors.aria2").write_text("control")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"download-1",
|
|
{
|
|
"download_id": "download-1",
|
|
"transfer_backend": "aria2",
|
|
"status": "downloading",
|
|
"save_path": str(save_path),
|
|
"file_path": str(save_path),
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"gid": "missing-gid",
|
|
},
|
|
)
|
|
|
|
class DummyAria2Downloader:
|
|
async def get_status_by_gid(self, gid):
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
downloads = await manager.get_active_downloads()
|
|
persisted = await manager._aria2_state_store.get("download-1")
|
|
assert persisted is not None
|
|
|
|
assert downloads["downloads"] == [
|
|
{
|
|
"download_id": "download-1",
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"progress": 0,
|
|
"status": "paused",
|
|
"error": None,
|
|
"bytes_downloaded": 0,
|
|
"total_bytes": None,
|
|
"bytes_per_second": 0.0,
|
|
}
|
|
]
|
|
assert manager._pause_events["download-1"].is_paused() is True
|
|
assert persisted["status"] == "paused"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_active_downloads_restarts_from_resume_context_for_active_restored_aria2(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
save_path = save_dir / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"download-1",
|
|
{
|
|
"download_id": "download-1",
|
|
"transfer_backend": "aria2",
|
|
"status": "downloading",
|
|
"save_path": str(save_path),
|
|
"file_path": str(save_path),
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"gid": "gid-1",
|
|
"resume_context": {
|
|
"version_info": {
|
|
"id": 34,
|
|
"modelId": 12,
|
|
"model": {"id": 12, "type": "LoRA", "tags": ["fantasy"]},
|
|
"images": [],
|
|
},
|
|
"file_info": {
|
|
"name": "file.safetensors",
|
|
"type": "Model",
|
|
"primary": True,
|
|
"downloadUrl": "https://example.com/file.safetensors",
|
|
},
|
|
"model_type": "lora",
|
|
"relative_path": "",
|
|
"save_dir": str(save_dir),
|
|
"download_urls": ["https://example.com/file.safetensors"],
|
|
},
|
|
},
|
|
)
|
|
|
|
restarted = {}
|
|
|
|
class DummyAria2Downloader:
|
|
async def get_status_by_gid(self, gid):
|
|
return {"gid": gid, "status": "active"}
|
|
|
|
async def restore_transfer(self, download_id, gid, restored_path):
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
async def fake_resume_restored_aria2_download(self, download_id, record):
|
|
restarted.update(
|
|
{
|
|
"download_id": download_id,
|
|
"model_id": record.get("model_id"),
|
|
"model_version_id": record.get("model_version_id"),
|
|
"save_dir": record.get("save_dir"),
|
|
"resume_context": record.get("resume_context"),
|
|
}
|
|
)
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_resume_restored_aria2_download",
|
|
fake_resume_restored_aria2_download,
|
|
)
|
|
execute_original = AsyncMock(side_effect=AssertionError("should not refetch metadata"))
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_execute_original_download",
|
|
execute_original,
|
|
)
|
|
|
|
downloads = await manager.get_active_downloads()
|
|
assert downloads["downloads"][0]["status"] == "downloading"
|
|
restarted_task = manager._download_tasks["download-1"]
|
|
await restarted_task
|
|
|
|
assert restarted["download_id"] == "download-1"
|
|
assert restarted["model_id"] == 12
|
|
assert restarted["model_version_id"] == 34
|
|
assert restarted["save_dir"] is None
|
|
assert restarted["resume_context"]["model_type"] == "lora"
|
|
assert execute_original.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_active_downloads_restores_persisted_aria2_without_initial_save_path(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
save_dir = tmp_path / "downloads"
|
|
save_dir.mkdir()
|
|
save_path = save_dir / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
(save_dir / "file.safetensors.aria2").write_text("control")
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"download-1",
|
|
{
|
|
"download_id": "download-1",
|
|
"transfer_backend": "aria2",
|
|
"status": "paused",
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"resume_context": {
|
|
"version_info": {
|
|
"id": 34,
|
|
"modelId": 12,
|
|
"model": {"id": 12, "type": "LoRA"},
|
|
"images": [],
|
|
},
|
|
"file_info": {
|
|
"name": "file.safetensors",
|
|
"type": "Model",
|
|
"primary": True,
|
|
"downloadUrl": "https://example.com/file.safetensors",
|
|
},
|
|
"model_type": "lora",
|
|
"relative_path": "",
|
|
"save_dir": str(save_dir),
|
|
"download_urls": ["https://example.com/file.safetensors"],
|
|
},
|
|
},
|
|
)
|
|
|
|
class DummyAria2Downloader:
|
|
async def get_status_by_gid(self, gid):
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
downloads = await manager.get_active_downloads()
|
|
persisted = await manager._aria2_state_store.get("download-1")
|
|
assert persisted is not None
|
|
|
|
assert downloads["downloads"] == [
|
|
{
|
|
"download_id": "download-1",
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"progress": 0,
|
|
"status": "paused",
|
|
"error": None,
|
|
"bytes_downloaded": 0,
|
|
"total_bytes": None,
|
|
"bytes_per_second": 0.0,
|
|
}
|
|
]
|
|
assert manager._active_downloads["download-1"]["file_path"] == str(save_path)
|
|
assert persisted["save_path"] == str(save_path)
|
|
assert persisted["file_path"] == str(save_path)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_restored_aria2_download_updates_terminal_status_and_cleanup(monkeypatch):
|
|
manager = DownloadManager()
|
|
manager._active_downloads["download-1"] = {
|
|
"transfer_backend": "aria2",
|
|
"status": "paused",
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"bytes_per_second": 10.0,
|
|
}
|
|
|
|
persist_state = AsyncMock()
|
|
cleanup_record = AsyncMock(return_value=None)
|
|
execute_download = AsyncMock(return_value={"success": True})
|
|
record_history = AsyncMock(return_value=None)
|
|
sync_version = AsyncMock(return_value=None)
|
|
|
|
monkeypatch.setattr(manager, "_persist_aria2_state", persist_state)
|
|
monkeypatch.setattr(manager, "_cleanup_download_record", cleanup_record)
|
|
monkeypatch.setattr(manager, "_execute_download", execute_download)
|
|
monkeypatch.setattr(manager, "_record_downloaded_version_history", record_history)
|
|
monkeypatch.setattr(manager, "_sync_downloaded_version", sync_version)
|
|
|
|
scheduled_tasks = []
|
|
original_create_task = asyncio.create_task
|
|
|
|
def tracking_create_task(coro):
|
|
task = original_create_task(coro)
|
|
scheduled_tasks.append(task)
|
|
return task
|
|
|
|
monkeypatch.setattr(download_manager.asyncio, "create_task", tracking_create_task)
|
|
|
|
result = await manager._resume_restored_aria2_download(
|
|
"download-1",
|
|
{
|
|
"download_id": "download-1",
|
|
"save_path": "/tmp/file.safetensors",
|
|
"file_path": "/tmp/file.safetensors",
|
|
"model_id": 12,
|
|
"model_version_id": 34,
|
|
"resume_context": {
|
|
"version_info": {
|
|
"id": 34,
|
|
"modelId": 12,
|
|
"model": {"id": 12},
|
|
"images": [],
|
|
},
|
|
"file_info": {
|
|
"name": "file.safetensors",
|
|
"downloadUrl": "https://example.com/file.safetensors",
|
|
},
|
|
"model_type": "lora",
|
|
"relative_path": "",
|
|
"save_dir": "/tmp",
|
|
"download_urls": ["https://example.com/file.safetensors"],
|
|
},
|
|
},
|
|
)
|
|
|
|
assert result == {"success": True}
|
|
assert manager._active_downloads["download-1"]["status"] == "completed"
|
|
assert manager._active_downloads["download-1"]["bytes_per_second"] == 0.0
|
|
assert persist_state.await_count == 2
|
|
assert len(scheduled_tasks) == 1
|
|
await asyncio.gather(*scheduled_tasks)
|
|
cleanup_record.assert_awaited_once_with("download-1")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_uses_captured_backend_when_settings_change(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
settings = get_settings_manager()
|
|
settings.settings["download_backend"] = "aria2"
|
|
|
|
semaphore = asyncio.Semaphore(0)
|
|
manager._download_semaphore = semaphore
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_original_download(
|
|
self,
|
|
model_id,
|
|
model_version_id,
|
|
save_dir,
|
|
relative_path,
|
|
progress_callback,
|
|
use_default_paths,
|
|
download_id=None,
|
|
transfer_backend="python",
|
|
source=None,
|
|
file_params=None,
|
|
use_save_dir_as_root=False,
|
|
):
|
|
captured["transfer_backend"] = transfer_backend
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager,
|
|
"_execute_original_download",
|
|
fake_execute_original_download,
|
|
)
|
|
|
|
download_task = asyncio.create_task(
|
|
manager.download_from_civitai(
|
|
model_version_id=99,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
)
|
|
|
|
await asyncio.sleep(0)
|
|
assert len(manager._active_downloads) == 1
|
|
download_id = next(iter(manager._active_downloads))
|
|
assert manager._active_downloads[download_id]["transfer_backend"] == "aria2"
|
|
|
|
settings.settings["download_backend"] = "python"
|
|
semaphore.release()
|
|
|
|
result = await download_task
|
|
|
|
assert result["success"] is True
|
|
assert captured["transfer_backend"] == "aria2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_aborts_when_version_exists(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""Test that download aborts when version already exists."""
|
|
scanners.lora.exists = True
|
|
|
|
manager = DownloadManager()
|
|
|
|
execute_mock = AsyncMock(return_value={"success": True})
|
|
monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock)
|
|
|
|
result = await manager.download_from_civitai(model_version_id=101, save_dir="/tmp")
|
|
|
|
assert result["success"] is False
|
|
assert result["error"] == "Model version already exists in lora library"
|
|
assert "download_id" in result
|
|
assert execute_mock.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_handles_metadata_errors(monkeypatch, scanners):
|
|
"""Test that download handles metadata fetch failures gracefully."""
|
|
async def failing_provider(*_args, **_kwargs):
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_default_metadata_provider",
|
|
AsyncMock(
|
|
return_value=SimpleNamespace(get_model_version=AsyncMock(return_value=None))
|
|
),
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
|
|
result = await manager.download_from_civitai(model_version_id=5, save_dir="/tmp")
|
|
|
|
assert result["success"] is False
|
|
assert result["error"] == "Failed to fetch model metadata"
|
|
assert "download_id" in result
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_rejects_unsupported_model_type(monkeypatch, scanners):
|
|
"""Test that unsupported model types are rejected."""
|
|
class Provider:
|
|
async def get_model_version(self, *_args, **_kwargs):
|
|
return {
|
|
"model": {"type": "Unsupported", "tags": []},
|
|
"files": [],
|
|
}
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_default_metadata_provider",
|
|
AsyncMock(return_value=Provider()),
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
|
|
result = await manager.download_from_civitai(model_version_id=5, save_dir="/tmp")
|
|
|
|
assert result["success"] is False
|
|
assert result["error"].startswith("Model type")
|
|
|
|
|
|
def test_embedding_relative_path_replaces_spaces():
|
|
"""Test that embedding paths replace spaces with underscores."""
|
|
manager = DownloadManager()
|
|
|
|
version_info = {
|
|
"baseModel": "Base Model",
|
|
"model": {"tags": ["tag with space"]},
|
|
"creator": {"username": "Author Name"},
|
|
}
|
|
|
|
relative_path = manager._calculate_relative_path(version_info, "embedding")
|
|
|
|
assert relative_path == "Base_Model/tag_with_space"
|
|
|
|
|
|
def test_relative_path_supports_model_and_version_placeholders():
|
|
"""Test that relative path supports {model_name} and {version_name} placeholders."""
|
|
manager = DownloadManager()
|
|
settings_manager = get_settings_manager()
|
|
settings_manager.settings["download_path_templates"]["lora"] = (
|
|
"{model_name}/{version_name}"
|
|
)
|
|
|
|
version_info = {
|
|
"baseModel": "BaseModel",
|
|
"name": "Version One",
|
|
"model": {"name": "Fancy Model", "tags": []},
|
|
}
|
|
|
|
relative_path = manager._calculate_relative_path(version_info, "lora")
|
|
|
|
assert relative_path == "Fancy Model/Version One"
|
|
|
|
|
|
def test_relative_path_sanitizes_model_and_version_placeholders():
|
|
"""Test that relative path sanitizes special characters in placeholders."""
|
|
manager = DownloadManager()
|
|
settings_manager = get_settings_manager()
|
|
settings_manager.settings["download_path_templates"]["lora"] = (
|
|
"{model_name}/{version_name}"
|
|
)
|
|
|
|
version_info = {
|
|
"baseModel": "BaseModel",
|
|
"name": "Version:One?",
|
|
"model": {"name": "Fancy:Model*", "tags": []},
|
|
}
|
|
|
|
relative_path = manager._calculate_relative_path(version_info, "lora")
|
|
|
|
assert relative_path == "Fancy_Model/Version_One"
|
|
|
|
|
|
def test_relative_path_empty_first_tag_fallback():
|
|
"""Test that empty first_tag falls back to 'no tags'."""
|
|
manager = DownloadManager()
|
|
settings_manager = get_settings_manager()
|
|
settings_manager.settings["download_path_templates"]["lora"] = (
|
|
"{base_model}/{first_tag}"
|
|
)
|
|
|
|
version_info = {
|
|
"baseModel": "SDXL",
|
|
"model": {"name": "Test Model", "tags": []},
|
|
"creator": {"username": "Author"},
|
|
}
|
|
|
|
relative_path = manager._calculate_relative_path(version_info, "lora")
|
|
|
|
assert relative_path == "SDXL/no tags"
|
|
|
|
|
|
def test_relative_path_empty_base_model_and_first_tag():
|
|
"""Test that empty base_model + empty first_tag does NOT produce a leading slash."""
|
|
manager = DownloadManager()
|
|
settings_manager = get_settings_manager()
|
|
settings_manager.settings["download_path_templates"]["lora"] = (
|
|
"{base_model}/{first_tag}"
|
|
)
|
|
|
|
version_info = {
|
|
"baseModel": "",
|
|
"model": {"name": "Test Model", "tags": []},
|
|
"creator": {"username": "Author"},
|
|
}
|
|
|
|
relative_path = manager._calculate_relative_path(version_info, "lora")
|
|
|
|
assert not relative_path.startswith("/")
|
|
assert relative_path == "no tags"
|
|
|
|
|
|
def test_relative_path_sanitizes_double_slashes():
|
|
"""Test that empty placeholder substitutions don't produce double slashes."""
|
|
manager = DownloadManager()
|
|
settings_manager = get_settings_manager()
|
|
settings_manager.settings["download_path_templates"]["lora"] = (
|
|
"{base_model}/{first_tag}/{author}"
|
|
)
|
|
|
|
version_info = {
|
|
"baseModel": "SDXL",
|
|
"model": {"name": "Test Model", "tags": []},
|
|
"creator": {"username": "Author"},
|
|
}
|
|
|
|
relative_path = manager._calculate_relative_path(version_info, "lora")
|
|
|
|
assert "//" not in relative_path
|
|
assert relative_path == "SDXL/no tags/Author"
|
|
|
|
|
|
def test_download_containment_accepts_symlink_save_dir(tmp_path):
|
|
"""Verify the download path containment check (download_manager.py:1395-1397)
|
|
accepts save directories reached through user-created symlinks inside the
|
|
library root — reproducing the symlink scenario from issue #1028."""
|
|
# Library root with a symlink subdirectory pointing to an external drive
|
|
lora_root = tmp_path / "loras"
|
|
lora_root.mkdir()
|
|
|
|
external_drive = tmp_path / "external" / "models"
|
|
external_drive.mkdir(parents=True)
|
|
|
|
symlink = lora_root / "Krea 2"
|
|
symlink.symlink_to(str(external_drive))
|
|
|
|
# Simulate a download: base_save_dir = library root,
|
|
# relative_path = "Krea 2/concept/NewModel"
|
|
base_save_dir = str(lora_root)
|
|
save_dir = os.path.join(base_save_dir, "Krea 2", "concept", "NewModel")
|
|
|
|
# Replicate the exact containment check from download_manager.py
|
|
resolved_dir = os.path.abspath(os.path.normpath(save_dir))
|
|
base_dir = os.path.abspath(os.path.normpath(base_save_dir))
|
|
|
|
# Must NOT be rejected — symlinks are legitimate business paths
|
|
assert resolved_dir.startswith(base_dir + os.sep)
|
|
|
|
|
|
def test_download_containment_rejects_dot_dot_traversal(tmp_path):
|
|
"""Verify the download path containment check still blocks ``..`` traversal
|
|
after the realpath → abspath change."""
|
|
lora_root = tmp_path / "loras"
|
|
lora_root.mkdir()
|
|
|
|
base_save_dir = str(lora_root)
|
|
save_dir = os.path.join(base_save_dir, "..", "..", "etc", "passwd")
|
|
|
|
resolved_dir = os.path.abspath(os.path.normpath(save_dir))
|
|
base_dir = os.path.abspath(os.path.normpath(base_save_dir))
|
|
|
|
# Must be rejected — dot-dot escapes the library root
|
|
assert not resolved_dir.startswith(base_dir + os.sep)
|
|
assert resolved_dir != base_dir
|
|
|
|
|
|
def test_distribute_preview_to_entries_moves_and_copies(tmp_path):
|
|
"""Test that preview distribution moves file to first entry and copies to others."""
|
|
manager = DownloadManager()
|
|
preview_file = tmp_path / "bundle.webp"
|
|
preview_file.write_bytes(b"image-data")
|
|
|
|
entries = [
|
|
SimpleNamespace(file_path=str(tmp_path / "model-one.safetensors")),
|
|
SimpleNamespace(file_path=str(tmp_path / "model-two.safetensors")),
|
|
]
|
|
|
|
targets = manager._distribute_preview_to_entries(str(preview_file), entries)
|
|
|
|
assert targets == [
|
|
str(tmp_path / "model-one.webp"),
|
|
str(tmp_path / "model-two.webp"),
|
|
]
|
|
assert not preview_file.exists()
|
|
assert Path(targets[0]).read_bytes() == b"image-data"
|
|
assert Path(targets[1]).read_bytes() == b"image-data"
|
|
|
|
|
|
def test_distribute_preview_to_entries_keeps_existing_file(tmp_path):
|
|
"""Test that existing preview files are not overwritten."""
|
|
manager = DownloadManager()
|
|
existing_preview = tmp_path / "model-one.webp"
|
|
existing_preview.write_bytes(b"preview")
|
|
|
|
entries = [
|
|
SimpleNamespace(file_path=str(tmp_path / "model-one.safetensors")),
|
|
SimpleNamespace(file_path=str(tmp_path / "model-two.safetensors")),
|
|
]
|
|
|
|
targets = manager._distribute_preview_to_entries(str(existing_preview), entries)
|
|
|
|
assert targets[0] == str(existing_preview)
|
|
assert Path(targets[1]).read_bytes() == b"preview"
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_skips_excluded_base_model(monkeypatch, scanners, metadata_provider):
|
|
manager = DownloadManager()
|
|
get_settings_manager().settings["download_skip_base_models"] = ["SDXL 1.0"]
|
|
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value={
|
|
"id": 99,
|
|
"model": {"type": "LoRA", "tags": ["fantasy"]},
|
|
"baseModel": "SDXL 1.0",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"type": "Model",
|
|
"primary": True,
|
|
"downloadUrl": "https://example.invalid/file.safetensors",
|
|
"name": "file.safetensors",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
execute_download = AsyncMock()
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=99,
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert result["skipped"] is True
|
|
assert result["status"] == "skipped"
|
|
assert result["reason"] == "base_model_excluded"
|
|
assert result["base_model"] == "SDXL 1.0"
|
|
assert result["file_name"] == "file.safetensors"
|
|
assert "file.safetensors" in result["message"]
|
|
execute_download.assert_not_called()
|
|
assert manager._active_downloads[result["download_id"]]["status"] == "skipped"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_skips_previously_downloaded_version(monkeypatch, scanners, metadata_provider):
|
|
manager = DownloadManager()
|
|
get_settings_manager().settings["skip_previously_downloaded_model_versions"] = True
|
|
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value={
|
|
"id": 42,
|
|
"model": {"type": "LoRA", "tags": ["fantasy"]},
|
|
"baseModel": "SDXL 1.0",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"type": "Model",
|
|
"primary": True,
|
|
"downloadUrl": "https://example.invalid/file.safetensors",
|
|
"name": "file.safetensors",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
history_service = AsyncMock()
|
|
history_service.has_been_downloaded = AsyncMock(return_value=True)
|
|
monkeypatch.setattr(
|
|
ServiceRegistry,
|
|
"get_downloaded_version_history_service",
|
|
AsyncMock(return_value=history_service),
|
|
)
|
|
|
|
execute_download = AsyncMock()
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=99,
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert result["skipped"] is True
|
|
assert result["status"] == "skipped"
|
|
assert result["reason"] == "previously_downloaded_version"
|
|
assert result["model_version_id"] == 99
|
|
assert result["file_name"] == "file.safetensors"
|
|
history_service.has_been_downloaded.assert_awaited_once_with("lora", 99)
|
|
execute_download.assert_not_called()
|
|
assert manager._active_downloads[result["download_id"]]["status"] == "skipped"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_proceeds_when_history_skip_disabled(monkeypatch, scanners, metadata_provider):
|
|
manager = DownloadManager()
|
|
get_settings_manager().settings["skip_previously_downloaded_model_versions"] = False
|
|
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value={
|
|
"id": 42,
|
|
"model": {"type": "LoRA", "tags": ["fantasy"]},
|
|
"baseModel": "SDXL 1.0",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"type": "Model",
|
|
"primary": True,
|
|
"downloadUrl": "https://example.invalid/file.safetensors",
|
|
"name": "file.safetensors",
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
history_service = AsyncMock()
|
|
history_service.has_been_downloaded = AsyncMock(return_value=True)
|
|
monkeypatch.setattr(
|
|
ServiceRegistry,
|
|
"get_downloaded_version_history_service",
|
|
AsyncMock(return_value=history_service),
|
|
)
|
|
|
|
execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=99,
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert result.get("skipped") is not True
|
|
history_service.has_been_downloaded.assert_not_called()
|
|
execute_download.assert_awaited_once()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Multi-file downloads within a single model version (#1058)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _multi_file_payload(include_hashes: bool = True):
|
|
"""Version payload with two weight files under the same version."""
|
|
files = [
|
|
{
|
|
"id": 1001,
|
|
"type": "Model",
|
|
"primary": True,
|
|
"downloadUrl": "https://example.invalid/file-a.safetensors",
|
|
"name": "file-a.safetensors",
|
|
},
|
|
{
|
|
"id": 1002,
|
|
"type": "Model",
|
|
"primary": False,
|
|
"downloadUrl": "https://example.invalid/file-b.safetensors",
|
|
"name": "file-b.safetensors",
|
|
},
|
|
]
|
|
if include_hashes:
|
|
files[0]["hashes"] = {"SHA256": "AAA111"}
|
|
files[1]["hashes"] = {"SHA256": "BBB222"}
|
|
return {
|
|
"id": 42,
|
|
"modelId": 7,
|
|
"model": {"type": "LoRA", "tags": ["fantasy"]},
|
|
"baseModel": "BaseModel",
|
|
"creator": {"username": "Author"},
|
|
"files": files,
|
|
}
|
|
|
|
|
|
def _local_entry(file_name: str, sha256: str, version_id: int = 42):
|
|
"""A library cache entry for one already-downloaded file of a version."""
|
|
return {
|
|
"file_name": file_name,
|
|
"file_path": f"/tmp/{file_name}.safetensors",
|
|
"sha256": sha256,
|
|
"civitai": {"id": version_id, "modelId": 7},
|
|
}
|
|
|
|
|
|
def _stub_history_service(monkeypatch, has_been_downloaded: bool = False):
|
|
history_service = AsyncMock()
|
|
history_service.has_been_downloaded = AsyncMock(
|
|
return_value=has_been_downloaded
|
|
)
|
|
history_service.mark_downloaded = AsyncMock()
|
|
monkeypatch.setattr(
|
|
ServiceRegistry,
|
|
"get_downloaded_version_history_service",
|
|
AsyncMock(return_value=history_service),
|
|
)
|
|
return history_service
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_allows_other_file_of_in_library_version(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""A different file of an in-library version must still download (#1058)."""
|
|
scanners.lora._cache.raw_data.append(_local_entry("file-a", "aaa111"))
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value=_multi_file_payload()
|
|
)
|
|
_stub_history_service(monkeypatch)
|
|
|
|
execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=42,
|
|
save_dir="/tmp",
|
|
file_params={"id": 1002, "type": "Model", "name": "file-b.safetensors"},
|
|
)
|
|
|
|
assert result["success"] is True
|
|
execute_download.assert_awaited_once()
|
|
# Version-level gates must not run for an explicit file selection
|
|
assert scanners.lora.calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_blocks_same_file_of_in_library_version(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""Re-downloading the SAME file of an in-library version stays blocked."""
|
|
scanners.lora._cache.raw_data.append(_local_entry("file-a", "aaa111"))
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value=_multi_file_payload()
|
|
)
|
|
_stub_history_service(monkeypatch)
|
|
|
|
execute_download = AsyncMock(return_value={"success": True})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=42,
|
|
save_dir="/tmp",
|
|
file_params={"id": 1001, "type": "Model", "name": "file-a.safetensors"},
|
|
)
|
|
|
|
assert result["success"] is False
|
|
assert "file-a.safetensors" in result["error"]
|
|
assert "already exists in lora library" in result["error"]
|
|
execute_download.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_unresolvable_file_params_falls_back_to_version_gate(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""file_params that match no file must not bypass version-level gates."""
|
|
scanners.lora.exists = True
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value=_multi_file_payload()
|
|
)
|
|
_stub_history_service(monkeypatch)
|
|
|
|
execute_download = AsyncMock(return_value={"success": True})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=42,
|
|
save_dir="/tmp",
|
|
file_params={"id": 9999, "type": "Model"},
|
|
)
|
|
|
|
assert result["success"] is False
|
|
assert result["error"] == "Model version already exists in lora library"
|
|
execute_download.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_empty_file_params_treated_as_no_selection(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""An empty file_params dict is normalized to None (version-level gates)."""
|
|
scanners.lora.exists = True
|
|
|
|
execute_download = AsyncMock(return_value={"success": True})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=42,
|
|
save_dir="/tmp",
|
|
file_params={},
|
|
)
|
|
|
|
assert result["success"] is False
|
|
assert result["error"] == "Model version already exists in lora library"
|
|
# Early version-level gate fired (file_params normalized to None)
|
|
assert scanners.lora.calls == [42]
|
|
execute_download.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_explicit_file_bypasses_history_skip(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""Explicit file selection bypasses the previously-downloaded skip."""
|
|
get_settings_manager().settings[
|
|
"skip_previously_downloaded_model_versions"
|
|
] = True
|
|
scanners.lora._cache.raw_data.append(_local_entry("file-a", "aaa111"))
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value=_multi_file_payload()
|
|
)
|
|
history_service = _stub_history_service(monkeypatch, has_been_downloaded=True)
|
|
|
|
execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=42,
|
|
save_dir="/tmp",
|
|
file_params={"id": 1002, "type": "Model", "name": "file-b.safetensors"},
|
|
)
|
|
|
|
assert result["success"] is True
|
|
execute_download.assert_awaited_once()
|
|
# History gate is bypassed before it even queries the service
|
|
history_service.has_been_downloaded.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_file_match_by_name_when_local_hash_missing(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""Legacy local metadata without sha256 falls back to name matching."""
|
|
scanners.lora._cache.raw_data.append(_local_entry("file-a", ""))
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value=_multi_file_payload()
|
|
)
|
|
_stub_history_service(monkeypatch)
|
|
|
|
execute_download = AsyncMock(return_value={"success": True})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=42,
|
|
save_dir="/tmp",
|
|
file_params={"id": 1001, "type": "Model"},
|
|
)
|
|
|
|
assert result["success"] is False
|
|
assert "already exists in lora library" in result["error"]
|
|
execute_download.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_no_false_positive_when_both_hashes_empty(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""Two empty hashes must never compare equal; name decides instead."""
|
|
# Local entry for a DIFFERENT file of the same version, no hash stored
|
|
scanners.lora._cache.raw_data.append(_local_entry("file-b", ""))
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value=_multi_file_payload(include_hashes=False)
|
|
)
|
|
_stub_history_service(monkeypatch)
|
|
|
|
execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=42,
|
|
save_dir="/tmp",
|
|
file_params={"id": 1001, "type": "Model"},
|
|
)
|
|
|
|
assert result["success"] is True
|
|
execute_download.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_model_id_only_with_file_params_resolves_file(
|
|
monkeypatch, scanners, metadata_provider
|
|
):
|
|
"""model_id-only requests with file_params resolve the file post-fetch."""
|
|
scanners.lora.exists = True # version-level index says "in library"
|
|
scanners.lora._cache.raw_data.append(_local_entry("file-a", "aaa111"))
|
|
metadata_provider.get_model_version = AsyncMock(
|
|
return_value=_multi_file_payload()
|
|
)
|
|
_stub_history_service(monkeypatch)
|
|
|
|
execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", execute_download, raising=False
|
|
)
|
|
|
|
manager = DownloadManager()
|
|
result = await manager.download_from_civitai(
|
|
model_id=7,
|
|
save_dir="/tmp",
|
|
file_params={"id": 1002, "type": "Model"},
|
|
)
|
|
|
|
assert result["success"] is True
|
|
execute_download.assert_awaited_once()
|
|
|
|
|
|
def test_resolve_target_file_by_id():
|
|
files = _multi_file_payload()["files"]
|
|
resolved = DownloadManager._resolve_target_file(files, {"id": 1002})
|
|
assert resolved is files[1]
|
|
|
|
|
|
def test_resolve_target_file_by_primary_flag():
|
|
files = _multi_file_payload()["files"]
|
|
resolved = DownloadManager._resolve_target_file(files, {"isPrimary": True})
|
|
assert resolved is files[0]
|
|
|
|
|
|
def test_resolve_target_file_returns_none_for_no_match():
|
|
files = _multi_file_payload()["files"]
|
|
assert DownloadManager._resolve_target_file(files, {"id": 9999}) is None
|
|
assert DownloadManager._resolve_target_file(files, None) is None
|
|
assert DownloadManager._resolve_target_file(files, {}) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_drops_unrestorable_persisted_records(monkeypatch, tmp_path):
|
|
"""Records without any resolvable target path can never be restored;
|
|
the restore sweep must delete them instead of skipping them forever."""
|
|
manager = DownloadManager()
|
|
|
|
await manager._aria2_state_store.upsert(
|
|
"download-orphan",
|
|
{
|
|
"download_id": "download-orphan",
|
|
"transfer_backend": "aria2",
|
|
"status": "failed",
|
|
# no save_path / file_path / resume_context
|
|
},
|
|
)
|
|
|
|
class DummyAria2Downloader:
|
|
async def get_status_by_gid(self, gid):
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
downloads = await manager.get_active_downloads()
|
|
|
|
assert downloads["downloads"] == []
|
|
assert await manager._aria2_state_store.get("download-orphan") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_discard_cleared_downloads_stops_tracking_and_preserves_files(
|
|
monkeypatch, tmp_path
|
|
):
|
|
manager = DownloadManager()
|
|
|
|
save_path = tmp_path / "file.safetensors"
|
|
save_path.write_text("partial")
|
|
control_path = tmp_path / "file.safetensors.aria2"
|
|
control_path.write_text("control")
|
|
|
|
async def _pending():
|
|
await asyncio.sleep(3600)
|
|
|
|
task = asyncio.create_task(_pending())
|
|
manager._download_tasks["download-1"] = task
|
|
manager._pause_events["download-1"] = download_manager.DownloadStreamControl()
|
|
manager._active_downloads["download-1"] = {
|
|
"status": "downloading",
|
|
"transfer_backend": "aria2",
|
|
"file_path": str(save_path),
|
|
}
|
|
await manager._aria2_state_store.upsert(
|
|
"download-1",
|
|
{
|
|
"download_id": "download-1",
|
|
"transfer_backend": "aria2",
|
|
"status": "downloading",
|
|
"save_path": str(save_path),
|
|
"gid": "gid-1",
|
|
},
|
|
)
|
|
|
|
cancelled = []
|
|
|
|
class DummyAria2Downloader:
|
|
async def has_transfer(self, download_id):
|
|
return True
|
|
|
|
async def cancel_download(self, download_id):
|
|
cancelled.append(download_id)
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_aria2_downloader",
|
|
AsyncMock(return_value=DummyAria2Downloader()),
|
|
)
|
|
|
|
discarded = await manager.discard_cleared_downloads(["download-1", "unknown-id"])
|
|
|
|
assert discarded == 1
|
|
assert cancelled == ["download-1"]
|
|
assert task.cancelled()
|
|
assert "download-1" not in manager._download_tasks
|
|
assert "download-1" not in manager._active_downloads
|
|
assert "download-1" not in manager._pause_events
|
|
assert await manager._aria2_state_store.get("download-1") is None
|
|
# Partial files are preserved for a future resume from disk.
|
|
assert save_path.exists()
|
|
assert control_path.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_uses_raw_file_name_from_mini_endpoint(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""#1100: when the REST name is rewritten ("{model}_{version}"), the raw
|
|
stored filename from the mini endpoint wins for the on-disk name."""
|
|
manager = DownloadManager()
|
|
get_settings_manager().settings["default_unet_root"] = str(tmp_path / "unet")
|
|
metadata_provider.payload = {
|
|
"id": 3284136,
|
|
"model": {"type": "Checkpoint", "tags": ["realistic"]},
|
|
"baseModel": "ZImageTurbo",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"id": 3168412,
|
|
"type": "Model",
|
|
"primary": True,
|
|
"name": "cyberrealisticZImage_v80.safetensors",
|
|
"downloadUrl": "https://civitai.com/api/download/models/3284136?fileId=3168412",
|
|
}
|
|
],
|
|
}
|
|
metadata_provider.get_version_file_mini = AsyncMock(
|
|
return_value={"fileName": "CyberRealistic_zit_v8.0_bf16.safetensors"}
|
|
)
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(self, **kwargs):
|
|
captured["download_urls"] = kwargs["download_urls"]
|
|
captured["file_path"] = kwargs["metadata"].file_path
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=3284136,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True, result
|
|
metadata_provider.get_version_file_mini.assert_awaited_once_with(3284136, 3168412)
|
|
assert captured["file_path"].endswith("CyberRealistic_zit_v8.0_bf16.safetensors")
|
|
# The file's own pinned downloadUrl is untouched.
|
|
assert captured["download_urls"] == [
|
|
"https://civitai.com/api/download/models/3284136?fileId=3168412"
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_falls_back_to_rest_name_when_mini_fails(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""A failed/absent mini lookup must keep the previous behavior."""
|
|
manager = DownloadManager()
|
|
metadata_provider.payload = {
|
|
"id": 42,
|
|
"model": {"type": "Checkpoint", "tags": ["fantasy"]},
|
|
"baseModel": "BaseModel",
|
|
"creator": {"username": "Author"},
|
|
"files": [
|
|
{
|
|
"id": 1001,
|
|
"type": "Model",
|
|
"primary": True,
|
|
"name": "rewritten_v10.safetensors",
|
|
"downloadUrl": "https://example.invalid/file.safetensors",
|
|
}
|
|
],
|
|
}
|
|
metadata_provider.get_version_file_mini = AsyncMock(return_value=None)
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(self, **kwargs):
|
|
captured["file_path"] = kwargs["metadata"].file_path
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=42,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source=None,
|
|
)
|
|
|
|
assert result["success"] is True
|
|
assert captured["file_path"].endswith("rewritten_v10.safetensors")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_download_skips_mini_lookup_for_civarchive_source(
|
|
monkeypatch, scanners, metadata_provider, tmp_path
|
|
):
|
|
"""CivArchive already serves raw stored names — no mini call."""
|
|
manager = DownloadManager()
|
|
mini_mock = AsyncMock(return_value={"fileName": "should_not_be_used.safetensors"})
|
|
metadata_provider.get_version_file_mini = mini_mock
|
|
|
|
monkeypatch.setattr(
|
|
download_manager,
|
|
"get_metadata_provider",
|
|
AsyncMock(return_value=metadata_provider),
|
|
)
|
|
|
|
captured = {}
|
|
|
|
async def fake_execute_download(self, **kwargs):
|
|
captured["file_path"] = kwargs["metadata"].file_path
|
|
return {"success": True}
|
|
|
|
monkeypatch.setattr(
|
|
DownloadManager, "_execute_download", fake_execute_download, raising=False
|
|
)
|
|
|
|
result = await manager.download_from_civitai(
|
|
model_version_id=99,
|
|
save_dir=str(tmp_path),
|
|
use_default_paths=True,
|
|
progress_callback=None,
|
|
source="civarchive",
|
|
)
|
|
|
|
assert result["success"] is True
|
|
mini_mock.assert_not_called()
|
|
assert captured["file_path"].endswith("file.safetensors")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_raw_file_name_edge_cases():
|
|
"""_fetch_raw_file_name never raises and strips path components."""
|
|
manager = DownloadManager()
|
|
|
|
provider = SimpleNamespace()
|
|
|
|
# Missing version id / file id short-circuit before any provider call.
|
|
provider.get_version_file_mini = AsyncMock()
|
|
assert await manager._fetch_raw_file_name(provider, None, 1) is None
|
|
assert await manager._fetch_raw_file_name(provider, 1, None) is None
|
|
provider.get_version_file_mini.assert_not_called()
|
|
|
|
# Provider without the method (older mocks / non-CivitAI providers).
|
|
assert await manager._fetch_raw_file_name(object(), 1, 2) is None
|
|
|
|
# Non-dict payload, empty fileName.
|
|
provider.get_version_file_mini = AsyncMock(return_value="oops")
|
|
assert await manager._fetch_raw_file_name(provider, 1, 2) is None
|
|
provider.get_version_file_mini = AsyncMock(return_value={"fileName": " "})
|
|
assert await manager._fetch_raw_file_name(provider, 1, 2) is None
|
|
|
|
# Path components are stripped defensively.
|
|
provider.get_version_file_mini = AsyncMock(
|
|
return_value={"fileName": "../evil/model.safetensors"}
|
|
)
|
|
assert await manager._fetch_raw_file_name(provider, 1, 2) == "model.safetensors"
|
|
|
|
# Provider exceptions degrade to None.
|
|
provider.get_version_file_mini = AsyncMock(side_effect=RuntimeError("boom"))
|
|
assert await manager._fetch_raw_file_name(provider, 1, 2) is None
|