fix(example-images): pin downloads to start library

This commit is contained in:
pixelpaws
2025-10-05 09:10:25 +08:00
parent 2e79970e6e
commit 5b564cd8a3
2 changed files with 140 additions and 16 deletions

View File

@@ -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()