mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-24 06:32:12 -03:00
fix(example-images): pin downloads to start library
This commit is contained in:
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -376,3 +377,80 @@ async def test_legacy_progress_file_migrates(monkeypatch: pytest.MonkeyPatch, tm
|
||||
assert new_progress.exists()
|
||||
contents = json.loads(new_progress.read_text(encoding="utf-8"))
|
||||
assert model_hash in contents.get("processed_models", [])
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("tmp_path")
|
||||
async def test_download_remains_in_initial_library(monkeypatch: pytest.MonkeyPatch, tmp_path):
|
||||
ws_manager = RecordingWebSocketManager()
|
||||
manager = download_module.DownloadManager(ws_manager=ws_manager)
|
||||
|
||||
monkeypatch.setitem(settings.settings, "example_images_path", str(tmp_path))
|
||||
monkeypatch.setitem(settings.settings, "libraries", {"LibraryA": {}, "LibraryB": {}})
|
||||
monkeypatch.setitem(settings.settings, "active_library", "LibraryA")
|
||||
|
||||
state = {"active": "LibraryA"}
|
||||
|
||||
def fake_get_active_library_name(self):
|
||||
return state["active"]
|
||||
|
||||
monkeypatch.setattr(type(settings), "get_active_library_name", fake_get_active_library_name)
|
||||
|
||||
model_hash = "f" * 64
|
||||
model_path = tmp_path / "example-model.safetensors"
|
||||
model_path.write_text("data", encoding="utf-8")
|
||||
|
||||
model = {
|
||||
"sha256": model_hash,
|
||||
"model_name": "Library Switch Model",
|
||||
"file_path": str(model_path),
|
||||
"file_name": "example-model.safetensors",
|
||||
}
|
||||
|
||||
_patch_scanner(monkeypatch, StubScanner([model]))
|
||||
|
||||
async def fake_process_local_examples(
|
||||
_file_path,
|
||||
_file_name,
|
||||
_model_name,
|
||||
model_dir,
|
||||
_optimize,
|
||||
):
|
||||
Path(model_dir).mkdir(parents=True, exist_ok=True)
|
||||
(Path(model_dir) / "local.txt").write_text("data", encoding="utf-8")
|
||||
state["active"] = "LibraryB"
|
||||
return True
|
||||
|
||||
async def fake_update_metadata(*_args, **_kwargs):
|
||||
return True
|
||||
|
||||
async def fake_get_downloader():
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_module.ExampleImagesProcessor,
|
||||
"process_local_examples",
|
||||
staticmethod(fake_process_local_examples),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
download_module.MetadataUpdater,
|
||||
"update_metadata_from_local_examples",
|
||||
staticmethod(fake_update_metadata),
|
||||
)
|
||||
monkeypatch.setattr(download_module, "get_downloader", fake_get_downloader)
|
||||
|
||||
result = await manager.start_download({"model_types": ["lora"], "delay": 0})
|
||||
assert result["success"] is True
|
||||
|
||||
if manager._download_task is not None:
|
||||
await asyncio.wait_for(manager._download_task, timeout=1)
|
||||
|
||||
library_a_root = tmp_path / "LibraryA"
|
||||
library_b_root = tmp_path / "LibraryB"
|
||||
|
||||
progress_file = library_a_root / ".download_progress.json"
|
||||
model_dir = library_a_root / model_hash
|
||||
|
||||
assert progress_file.exists()
|
||||
assert (model_dir / "local.txt").exists()
|
||||
assert not (library_b_root / ".download_progress.json").exists()
|
||||
assert not (library_b_root / model_hash).exists()
|
||||
|
||||
Reference in New Issue
Block a user