alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-05-06 16:36:45 -03:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into codex/github-mention-fixnetwork-add-connectivityguard-to-short

Browse Source
This commit is contained in:
pixelpaws
2026-04-20 15:54:30 +08:00
committed by GitHub
parent 9443175abc 34f03d6495
commit fc19a145ff
28 changed files with 4469 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files

183
tests/config/test_config_save_paths.py
View File

@@ -46,6 +46,7 @@ def test_save_paths_renames_default_library(monkeypatch: pytest.MonkeyPatch, tmp
self.delete_calls = []
self.upsert_calls = []
self._renamed = False
self.active_library = "default"
def get_libraries(self):
if self._renamed:
@@ -62,6 +63,11 @@ def test_save_paths_renames_default_library(monkeypatch: pytest.MonkeyPatch, tmp
def rename_library(self, old_name: str, new_name: str):
self.rename_calls.append((old_name, new_name))
self._renamed = True
if self.active_library == old_name:
self.active_library = new_name
def get_active_library_name(self):
return self.active_library
def delete_library(self, name: str): # pragma: no cover - defensive guard
self.delete_calls.append(name)
@@ -104,6 +110,7 @@ def test_save_paths_logs_warning_when_upsert_fails(
class RaisingSettingsService:
def __init__(self):
self.upsert_attempts = []
self.active_library = "comfyui"
def get_libraries(self):
return {
@@ -116,6 +123,9 @@ def test_save_paths_logs_warning_when_upsert_fails(
def rename_library(self, *_):
raise AssertionError("rename_library should not be invoked")
def get_active_library_name(self):
return self.active_library
def upsert_library(self, name: str, **payload):
self.upsert_attempts.append((name, payload))
raise RuntimeError("boom")
@@ -135,6 +145,8 @@ def test_save_paths_repairs_empty_default_roots(monkeypatch: pytest.MonkeyPatch,
folder_paths = _setup_config_environment(monkeypatch, tmp_path)
class FakeSettingsService:
active_library = "comfyui"
def get_libraries(self):
return {
"comfyui": {
@@ -148,6 +160,9 @@ def test_save_paths_repairs_empty_default_roots(monkeypatch: pytest.MonkeyPatch,
def rename_library(self, *_):
raise AssertionError("rename_library should not be invoked")
def get_active_library_name(self):
return self.active_library
def upsert_library(self, name: str, **payload):
self.name = name
self.payload = payload
@@ -167,6 +182,8 @@ def test_save_paths_repairs_stale_default_roots(monkeypatch: pytest.MonkeyPatch,
folder_paths = _setup_config_environment(monkeypatch, tmp_path)
class FakeSettingsService:
active_library = "comfyui"
def get_libraries(self):
return {
"comfyui": {
@@ -180,6 +197,9 @@ def test_save_paths_repairs_stale_default_roots(monkeypatch: pytest.MonkeyPatch,
def rename_library(self, *_):
raise AssertionError("rename_library should not be invoked")
def get_active_library_name(self):
return self.active_library
def upsert_library(self, name: str, **payload):
self.name = name
self.payload = payload
@@ -199,6 +219,8 @@ def test_save_paths_keeps_valid_default_roots(monkeypatch: pytest.MonkeyPatch, t
folder_paths = _setup_config_environment(monkeypatch, tmp_path)
class FakeSettingsService:
active_library = "comfyui"
def get_libraries(self):
return {
"comfyui": {
@@ -212,6 +234,9 @@ def test_save_paths_keeps_valid_default_roots(monkeypatch: pytest.MonkeyPatch, t
def rename_library(self, *_):
raise AssertionError("rename_library should not be invoked")
def get_active_library_name(self):
return self.active_library
def upsert_library(self, name: str, **payload):
self.name = name
self.payload = payload
@@ -258,6 +283,7 @@ def test_save_paths_removes_template_default_library(monkeypatch, tmp_path):
self.rename_calls = []
self.delete_calls = []
self.upsert_calls = []
self.active_library = "default"
def get_libraries(self):
return self.libraries
@@ -265,6 +291,8 @@ def test_save_paths_removes_template_default_library(monkeypatch, tmp_path):
def rename_library(self, old_name: str, new_name: str):
self.rename_calls.append((old_name, new_name))
self.libraries[new_name] = self.libraries.pop(old_name)
if self.active_library == old_name:
self.active_library = new_name
def delete_library(self, name: str):
self.delete_calls.append(name)
@@ -273,6 +301,11 @@ def test_save_paths_removes_template_default_library(monkeypatch, tmp_path):
def upsert_library(self, name: str, **payload):
self.upsert_calls.append((name, payload))
self.libraries[name] = {**payload}
if payload.get("activate"):
self.active_library = name
def get_active_library_name(self):
return self.active_library
fake_settings = FakeSettingsService()
monkeypatch.setattr(settings_manager_module, "settings", fake_settings)
@@ -313,6 +346,156 @@ def test_save_paths_removes_template_default_library(monkeypatch, tmp_path):
assert payload["activate"] is True
def test_save_paths_keeps_default_roots_in_extra_paths(monkeypatch: pytest.MonkeyPatch, tmp_path):
folder_paths = _setup_config_environment(monkeypatch, tmp_path)
extra_lora_dir = tmp_path / "extra_loras"
extra_checkpoint_dir = tmp_path / "extra_checkpoints"
extra_embedding_dir = tmp_path / "extra_embeddings"
for directory in (extra_lora_dir, extra_checkpoint_dir, extra_embedding_dir):
directory.mkdir()
class FakeSettingsService:
active_library = "comfyui"
def get_libraries(self):
return {
"comfyui": {
"folder_paths": {key: list(value) for key, value in folder_paths.items()},
"extra_folder_paths": {
"loras": [str(extra_lora_dir)],
"checkpoints": [str(extra_checkpoint_dir)],
"embeddings": [str(extra_embedding_dir)],
},
"default_lora_root": str(extra_lora_dir),
"default_checkpoint_root": str(extra_checkpoint_dir),
"default_embedding_root": str(extra_embedding_dir),
}
}
def rename_library(self, *_):
raise AssertionError("rename_library should not be invoked")
def get_active_library_name(self):
return self.active_library
def upsert_library(self, name: str, **payload):
self.name = name
self.payload = payload
fake_settings = FakeSettingsService()
monkeypatch.setattr(settings_manager_module, "settings", fake_settings)
config_module.Config()
assert fake_settings.name == "comfyui"
assert fake_settings.payload["extra_folder_paths"]["loras"] == [str(extra_lora_dir).replace("\\", "/")]
assert fake_settings.payload["extra_folder_paths"]["checkpoints"] == [
str(extra_checkpoint_dir).replace("\\", "/")
]
assert fake_settings.payload["extra_folder_paths"]["embeddings"] == [
str(extra_embedding_dir).replace("\\", "/")
]
assert fake_settings.payload["default_lora_root"] == str(extra_lora_dir).replace("\\", "/")
assert fake_settings.payload["default_checkpoint_root"] == str(extra_checkpoint_dir).replace("\\", "/")
assert fake_settings.payload["default_embedding_root"] == str(extra_embedding_dir).replace("\\", "/")
assert fake_settings.payload["activate"] is True
def test_save_paths_repairs_empty_default_roots_to_extra_paths_when_primary_missing(
monkeypatch: pytest.MonkeyPatch, tmp_path
):
_setup_config_environment(monkeypatch, tmp_path)
extra_lora_dir = tmp_path / "extra_loras"
extra_lora_dir.mkdir()
monkeypatch.setattr(
config_module.folder_paths,
"get_folder_paths",
lambda kind: [] if kind == "loras" else [],
)
class FakeSettingsService:
active_library = "comfyui"
def get_libraries(self):
return {
"comfyui": {
"folder_paths": {
"loras": [],
"checkpoints": [],
"unet": [],
"embeddings": [],
},
"extra_folder_paths": {
"loras": [str(extra_lora_dir)],
},
"default_lora_root": "",
}
}
def rename_library(self, *_):
raise AssertionError("rename_library should not be invoked")
def get_active_library_name(self):
return self.active_library
def upsert_library(self, name: str, **payload):
self.name = name
self.payload = payload
fake_settings = FakeSettingsService()
monkeypatch.setattr(settings_manager_module, "settings", fake_settings)
config_module.Config()
assert fake_settings.name == "comfyui"
assert fake_settings.payload["default_lora_root"] == str(extra_lora_dir).replace("\\", "/")
def test_save_paths_does_not_activate_comfyui_library_when_another_library_is_active(
monkeypatch: pytest.MonkeyPatch, tmp_path
):
folder_paths = _setup_config_environment(monkeypatch, tmp_path)
class FakeSettingsService:
def __init__(self):
self.active_library = "studio"
self.upsert_calls = []
def get_libraries(self):
return {
"studio": {
"folder_paths": {"loras": ["/studio/loras"]},
},
"comfyui": {
"folder_paths": {key: list(value) for key, value in folder_paths.items()},
"default_lora_root": folder_paths["loras"][0],
"default_checkpoint_root": folder_paths["checkpoints"][0],
"default_embedding_root": folder_paths["embeddings"][0],
},
}
def rename_library(self, *_):
raise AssertionError("rename_library should not be invoked")
def get_active_library_name(self):
return self.active_library
def upsert_library(self, name: str, **payload):
self.upsert_calls.append((name, payload))
fake_settings = FakeSettingsService()
monkeypatch.setattr(settings_manager_module, "settings", fake_settings)
config_module.Config()
assert len(fake_settings.upsert_calls) == 1
name, payload = fake_settings.upsert_calls[0]
assert name == "comfyui"
assert payload["activate"] is False
def test_apply_library_settings_merges_extra_paths(monkeypatch, tmp_path):
"""Test that apply_library_settings correctly merges folder_paths with extra_folder_paths."""
loras_dir = tmp_path / "loras"

35
tests/frontend/managers/settingsManager.library.test.js
View File

@@ -305,4 +305,39 @@ describe('SettingsManager library controls', () => {
'success',
);
});
it('loads download backend settings and toggles the aria2 path field', () => {
const manager = createManager();
document.body.innerHTML = `
<select id="downloadBackend">
<option value="python">Python</option>
<option value="aria2">aria2</option>
</select>
<div id="aria2PathSetting" style="display: none;"></div>
<input id="aria2cPath" />
`;
state.global.settings = {
download_backend: 'aria2',
aria2c_path: '/usr/bin/aria2c',
};
const saveSpy = vi.spyOn(manager, 'saveSelectSetting').mockResolvedValue();
manager.loadDownloadBackendSettings();
const backendSelect = document.getElementById('downloadBackend');
const aria2PathSetting = document.getElementById('aria2PathSetting');
const aria2cPath = document.getElementById('aria2cPath');
expect(backendSelect.value).toBe('aria2');
expect(aria2cPath.value).toBe('/usr/bin/aria2c');
expect(aria2PathSetting.style.display).toBe('block');
backendSelect.value = 'python';
backendSelect.onchange();
expect(aria2PathSetting.style.display).toBe('none');
expect(saveSpy).toHaveBeenCalledWith('downloadBackend', 'download_backend');
});
});

354
tests/services/test_aria2_downloader.py Normal file
View File

@@ -0,0 +1,354 @@
from __future__ import annotations
import asyncio
from pathlib import Path
from unittest.mock import AsyncMock
import pytest
from py.services.aria2_downloader import Aria2Downloader, Aria2Error
from py.services.aria2_transfer_state import Aria2TransferStateStore
from py.services import aria2_transfer_state
@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.mark.asyncio
async def test_download_file_polls_until_complete(tmp_path, monkeypatch):
downloader = Aria2Downloader()
downloader._rpc_url = "http://127.0.0.1/jsonrpc"
downloader._rpc_secret = "secret"
save_path = tmp_path / "downloads" / "model.safetensors"
progress_events = []
rpc_calls = []
statuses = iter(
[
{
"gid": "gid-1",
"status": "active",
"completedLength": "5",
"totalLength": "10",
"downloadSpeed": "25",
},
{
"gid": "gid-1",
"status": "complete",
"completedLength": "10",
"totalLength": "10",
"downloadSpeed": "0",
"files": [{"path": str(save_path)}],
},
]
)
async def fake_rpc_call(method, params):
rpc_calls.append((method, params))
if method == "aria2.addUri":
return "gid-1"
if method == "aria2.tellStatus":
return next(statuses)
raise AssertionError(f"Unexpected RPC method: {method}")
monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
monkeypatch.setattr(
downloader,
"_resolve_authenticated_redirect_url",
AsyncMock(
return_value="https://signed.example.com/model.safetensors?token=abc"
),
)
monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
async def progress_callback(progress, snapshot=None):
progress_events.append(snapshot.percent_complete if snapshot else progress)
success, result = await downloader.download_file(
"https://civitai.com/api/download/models/123",
str(save_path),
download_id="download-1",
progress_callback=progress_callback,
headers={"Authorization": "Bearer token"},
)
assert success is True
assert result == str(save_path)
assert progress_events == [50.0, 100.0]
assert downloader._transfers == {}
assert rpc_calls[0][0] == "aria2.addUri"
assert rpc_calls[0][1][0] == [
"https://signed.example.com/model.safetensors?token=abc"
]
assert rpc_calls[0][1][1]["out"] == "model.safetensors"
assert "header" not in rpc_calls[0][1][1]
@pytest.mark.asyncio
async def test_transfer_state_store_shares_lock_and_preserves_concurrent_updates(tmp_path):
state_path = tmp_path / "cache" / "aria2" / "downloads.json"
store_a = Aria2TransferStateStore(str(state_path))
store_b = Aria2TransferStateStore(str(state_path))
assert store_a._lock is store_b._lock
await asyncio.gather(
store_a.upsert("download-1", {"status": "downloading", "gid": "gid-1"}),
store_b.upsert("download-2", {"status": "paused", "gid": "gid-2"}),
)
assert await store_a.get("download-1") == {"status": "downloading", "gid": "gid-1"}
assert await store_b.get("download-2") == {"status": "paused", "gid": "gid-2"}
@pytest.mark.asyncio
async def test_download_file_keeps_auth_headers_when_civitai_does_not_redirect(
tmp_path, monkeypatch
):
downloader = Aria2Downloader()
downloader._rpc_url = "http://127.0.0.1/jsonrpc"
downloader._rpc_secret = "secret"
save_path = tmp_path / "downloads" / "model.safetensors"
rpc_calls = []
statuses = iter(
[
{
"gid": "gid-1",
"status": "complete",
"completedLength": "10",
"totalLength": "10",
"downloadSpeed": "0",
"files": [{"path": str(save_path)}],
},
]
)
async def fake_rpc_call(method, params):
rpc_calls.append((method, params))
if method == "aria2.addUri":
return "gid-1"
if method == "aria2.tellStatus":
return next(statuses)
raise AssertionError(f"Unexpected RPC method: {method}")
monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
monkeypatch.setattr(
downloader,
"_resolve_authenticated_redirect_url",
AsyncMock(return_value="https://civitai.com/api/download/models/123"),
)
monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
success, result = await downloader.download_file(
"https://civitai.com/api/download/models/123",
str(save_path),
download_id="download-1",
headers={"Authorization": "Bearer token"},
)
assert success is True
assert result == str(save_path)
assert rpc_calls[0][1][0] == ["https://civitai.com/api/download/models/123"]
assert rpc_calls[0][1][1]["header"] == ["Authorization: Bearer token"]
@pytest.mark.asyncio
async def test_pause_resume_cancel_forward_to_rpc(monkeypatch):
downloader = Aria2Downloader()
downloader._transfers["download-1"] = type(
"Transfer", (), {"gid": "gid-1", "save_path": "/tmp/model.safetensors"}
)()
calls = []
async def fake_rpc_call(method, params):
calls.append((method, params))
return "gid-1"
monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
pause_result = await downloader.pause_download("download-1")
resume_result = await downloader.resume_download("download-1")
cancel_result = await downloader.cancel_download("download-1")
assert pause_result["success"] is True
assert resume_result["success"] is True
assert cancel_result["success"] is True
assert calls == [
("aria2.forcePause", ["gid-1"]),
("aria2.unpause", ["gid-1"]),
("aria2.forceRemove", ["gid-1"]),
]
@pytest.mark.asyncio
async def test_download_file_reuses_existing_transfer_without_add_uri(
tmp_path, monkeypatch
):
downloader = Aria2Downloader()
downloader._rpc_url = "http://127.0.0.1/jsonrpc"
downloader._rpc_secret = "secret"
save_path = tmp_path / "downloads" / "model.safetensors"
downloader._transfers["download-1"] = type(
"Transfer", (), {"gid": "gid-1", "save_path": str(save_path)}
)()
rpc_calls = []
statuses = iter(
[
{
"gid": "gid-1",
"status": "active",
"completedLength": "5",
"totalLength": "10",
"downloadSpeed": "25",
},
{
"gid": "gid-1",
"status": "complete",
"completedLength": "10",
"totalLength": "10",
"downloadSpeed": "0",
"files": [{"path": str(save_path)}],
},
]
)
async def fake_rpc_call(method, params):
rpc_calls.append((method, params))
if method == "aria2.tellStatus":
return next(statuses)
raise AssertionError(f"Unexpected RPC method: {method}")
monkeypatch.setattr(downloader, "_ensure_process", AsyncMock())
monkeypatch.setattr(downloader, "_rpc_call", fake_rpc_call)
monkeypatch.setattr("py.services.aria2_downloader.asyncio.sleep", AsyncMock())
success, result = await downloader.download_file(
"https://example.com/model.safetensors",
str(save_path),
download_id="download-1",
)
assert success is True
assert result == str(save_path)
assert [call[0] for call in rpc_calls] == ["aria2.tellStatus", "aria2.tellStatus"]
def test_build_progress_snapshot_normalizes_numeric_fields():
downloader = Aria2Downloader()
snapshot = downloader._build_progress_snapshot(
{
"completedLength": "75",
"totalLength": "100",
"downloadSpeed": "512",
}
)
assert snapshot.percent_complete == 75.0
assert snapshot.bytes_downloaded == 75
assert snapshot.total_bytes == 100
assert snapshot.bytes_per_second == 512.0
def test_resolve_executable_raises_when_binary_missing(monkeypatch):
downloader = Aria2Downloader()
settings = type("Settings", (), {"get": lambda self, key, default=None: ""})()
monkeypatch.setattr("py.services.aria2_downloader.get_settings_manager", lambda: settings)
monkeypatch.setattr("py.services.aria2_downloader.shutil.which", lambda _: None)
with pytest.raises(Aria2Error):
downloader._resolve_executable()
@pytest.mark.asyncio
async def test_rpc_call_surfaces_json_error_on_non_200(monkeypatch):
downloader = Aria2Downloader()
downloader._rpc_url = "http://127.0.0.1:6800/jsonrpc"
downloader._rpc_secret = "secret"
class FakeResponse:
status = 400
async def text(self):
return (
'{"jsonrpc":"2.0","id":"x","error":{"code":1,"message":"Unauthorized"}}'
)
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
class FakeSession:
def post(self, _url, json=None):
return FakeResponse()
monkeypatch.setattr(downloader, "_get_rpc_session", AsyncMock(return_value=FakeSession()))
with pytest.raises(Aria2Error) as exc_info:
await downloader._rpc_call("aria2.addUri", [["https://example.com/file"]])
assert "Unauthorized" in str(exc_info.value)
assert "aria2.addUri" in str(exc_info.value)
@pytest.mark.asyncio
async def test_resolve_authenticated_redirect_url_returns_location(monkeypatch):
downloader = Aria2Downloader()
class FakeResponse:
status = 307
headers = {"Location": "https://signed.example.com/file.safetensors"}
async def text(self):
return ""
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
class FakeSession:
def get(self, _url, headers=None, allow_redirects=False, proxy=None):
return FakeResponse()
class FakeDownloader:
default_headers = {"User-Agent": "ComfyUI-LoRA-Manager/1.0"}
proxy_url = None
@property
def session(self):
async def _session():
return FakeSession()
return _session()
fake_downloader = FakeDownloader()
monkeypatch.setattr(
"py.services.aria2_downloader.get_downloader",
AsyncMock(return_value=fake_downloader),
)
result = await downloader._resolve_authenticated_redirect_url(
"https://civitai.com/api/download/models/123",
{"Authorization": "Bearer token"},
)
assert result == "https://signed.example.com/file.safetensors"

62
tests/services/test_connectivity_guard.py
View File

@@ -39,6 +39,26 @@ async def test_connectivity_guard_enters_cooldown_after_threshold():
assert guard.cooldown_remaining_seconds() > 0
async def test_connectivity_guard_scopes_cooldown_to_destination():
guard = await ConnectivityGuard.get_instance()
destination_a = "civitai.com"
destination_b = "api.github.com"
guard.register_network_failure(
OSError(errno.ENETUNREACH, "unreachable"),
destination_a,
)
guard.register_network_failure(asyncio.TimeoutError("timeout"), destination_a)
guard.register_network_failure(ConnectionRefusedError("refused"), destination_a)
assert guard.should_block_request(destination_a) is True
assert guard.should_block_request(destination_b) is False
guard.register_success(destination_a)
assert guard.should_block_request(destination_a) is False
async def test_connectivity_guard_recovers_after_success():
guard = await ConnectivityGuard.get_instance()
guard.online = False
@@ -55,21 +75,51 @@ async def test_connectivity_guard_recovers_after_success():
async def test_downloader_short_circuits_all_request_helpers_during_cooldown():
guard = await ConnectivityGuard.get_instance()
guard.cooldown_until = datetime.now() + timedelta(seconds=30)
guard.online = False
guard.failure_count = 3
destination = "example.invalid"
guard.register_network_failure(
OSError(errno.ENETUNREACH, "unreachable"),
destination,
)
guard.register_network_failure(asyncio.TimeoutError("timeout"), destination)
guard.register_network_failure(
ConnectionRefusedError("refused"),
destination,
)
downloader = Downloader()
ok, payload = await downloader.make_request("GET", "https://example.invalid")
ok, payload = await downloader.make_request("GET", f"https://{destination}")
assert ok is False
assert payload == OFFLINE_COOLDOWN_ERROR
ok, payload, headers = await downloader.download_to_memory("https://example.invalid")
ok, payload, headers = await downloader.download_to_memory(f"https://{destination}")
assert ok is False
assert payload == OFFLINE_FRIENDLY_MESSAGE
assert headers is None
ok, payload = await downloader.get_response_headers("https://example.invalid")
ok, payload = await downloader.get_response_headers(f"https://{destination}")
assert ok is False
assert payload == OFFLINE_COOLDOWN_ERROR
async def test_downloader_only_short_circuits_requests_for_same_destination():
guard = await ConnectivityGuard.get_instance()
guard.register_network_failure(
OSError(errno.ENETUNREACH, "unreachable"),
"example.invalid",
)
guard.register_network_failure(asyncio.TimeoutError("timeout"), "example.invalid")
guard.register_network_failure(
ConnectionRefusedError("refused"),
"example.invalid",
)
downloader = Downloader()
ok, payload = await downloader.make_request("GET", "https://example.invalid")
assert ok is False
assert payload == OFFLINE_COOLDOWN_ERROR
assert (
guard.should_block_request(downloader._guard_destination("https://example.com"))
is False
)

651
tests/services/test_download_manager_basic.py
View File

@@ -10,6 +10,7 @@ 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
@@ -46,6 +47,16 @@ def isolate_settings(monkeypatch, tmp_path):
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:
@@ -179,6 +190,7 @@ async def test_successful_download_uses_defaults(
progress_callback,
model_type,
download_id,
transfer_backend=None,
):
captured.update(
{
@@ -268,6 +280,7 @@ async def test_download_uses_active_mirrors(
progress_callback,
model_type,
download_id,
transfer_backend=None,
):
captured["download_urls"] = download_urls
return {"success": True}
@@ -288,6 +301,644 @@ async def test_download_uses_active_mirrors(
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,
):
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 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 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,
):
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

1072
tests/services/test_download_manager_error.py
View File

File diff suppressed because it is too large Load Diff

42
tests/services/test_settings_manager.py
View File

@@ -147,6 +147,11 @@ def test_environment_variable_overrides_settings(tmp_path, monkeypatch):
assert mgr.get("civitai_api_key") == "secret"
def test_default_download_backend_is_python(manager):
assert manager.get("download_backend") == "python"
assert manager.get("aria2c_path") == ""
def _create_manager_with_settings(
tmp_path, monkeypatch, initial_settings, *, save_spy=None
):
@@ -327,6 +332,43 @@ def test_auto_set_default_roots_keeps_valid_values(manager):
assert manager.get("default_embedding_root") == "/embeddings"
def test_auto_set_default_roots_keeps_valid_extra_values(manager):
manager.settings["default_lora_root"] = "/extra-loras"
manager.settings["default_checkpoint_root"] = "/extra-checkpoints"
manager.settings["default_embedding_root"] = "/extra-embeddings"
manager.settings["default_unet_root"] = "/extra-unet"
manager.settings["folder_paths"] = {
"loras": ["/loras"],
"checkpoints": ["/checkpoints"],
"unet": ["/unet"],
"embeddings": ["/embeddings"],
}
manager.settings["extra_folder_paths"] = {
"loras": ["/extra-loras"],
"checkpoints": ["/extra-checkpoints"],
"unet": ["/extra-unet"],
"embeddings": ["/extra-embeddings"],
}
manager._auto_set_default_roots()
assert manager.get("default_lora_root") == "/extra-loras"
assert manager.get("default_checkpoint_root") == "/extra-checkpoints"
assert manager.get("default_unet_root") == "/extra-unet"
assert manager.get("default_embedding_root") == "/extra-embeddings"
def test_auto_set_default_roots_falls_back_to_extra_when_primary_missing(manager):
manager.settings["default_lora_root"] = ""
manager.settings["folder_paths"] = {"loras": []}
manager.settings["extra_folder_paths"] = {"loras": ["/extra-loras"]}
manager._auto_set_default_roots()
assert manager.get("default_lora_root") == "/extra-loras"
def test_delete_setting(manager):
manager.set("example", 1)
manager.delete("example")