mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-03 16:45:33 -03:00
Will Miao
896ce5eddb
Applying a filename template to a large library re-did O(library) work for every renamed file: a full natsort resort plus whole-table SQLite rewrite and download-history resync after each rename, and a full scan plus resort of the entire recipe collection per renamed LoRA. On a 20k-model library with 300k recipes on a HDD this pushed "Apply to Library" into multi-day runs. - ModelScanner.defer_cache_persist(): bulk loops update the in-memory entry and indexes only; resort + persist + download-history sync run once at context exit, forced even on cancellation/error since files are already renamed on disk. Single-rename callers keep immediate per-call behavior. - RecipeScanner.build_lora_hash_index(): one-shot hash -> recipes index so per-file lookups are O(1); update_lora_filename_by_hash gains hash_index / defer_maintenance params, with a single finalize_bulk_filename_updates() resort at the end of a bulk session. - ModelLifecycleService.bulk_rename_session() / BulkRenameContext wire the deferred path through rename_model (hash index built lazily on first recipe-touching rename). - Blocking os.rename sequence offloaded via asyncio.to_thread so one file's HDD I/O no longer stalls the event loop (no cross-file parallelism). - Skip logic, per-batch WebSocket progress, cancellation, and result counters unchanged.
777 lines
27 KiB
Python
777 lines
27 KiB
Python
import asyncio
|
|
import logging
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from typing import Any, AsyncIterator, Dict, List, Optional
|
|
|
|
import pytest
|
|
|
|
from py.services.download_coordinator import DownloadCoordinator
|
|
from py.services.metadata_sync_service import MetadataSyncService
|
|
from py.services.model_file_service import AutoOrganizeResult, ModelFileService
|
|
from py.services.use_cases import (
|
|
AutoOrganizeInProgressError,
|
|
AutoOrganizeUseCase,
|
|
BulkMetadataRefreshUseCase,
|
|
DownloadExampleImagesConfigurationError,
|
|
DownloadExampleImagesInProgressError,
|
|
DownloadExampleImagesUseCase,
|
|
DownloadModelEarlyAccessError,
|
|
DownloadModelUseCase,
|
|
DownloadModelValidationError,
|
|
FilenameTemplateUseCase,
|
|
ImportExampleImagesUseCase,
|
|
ImportExampleImagesValidationError,
|
|
)
|
|
from py.utils.example_images_download_manager import (
|
|
DownloadConfigurationError,
|
|
DownloadInProgressError,
|
|
ExampleImagesDownloadError,
|
|
)
|
|
from py.utils.example_images_processor import (
|
|
ExampleImagesImportError,
|
|
ExampleImagesProcessor,
|
|
ExampleImagesValidationError,
|
|
)
|
|
from py.utils.metadata_manager import MetadataManager
|
|
from tests.conftest import MockCache, MockModelService, MockScanner
|
|
|
|
|
|
class StubLockProvider:
|
|
def __init__(self) -> None:
|
|
self._lock = asyncio.Lock()
|
|
self.running = False
|
|
|
|
def is_auto_organize_running(self) -> bool:
|
|
return self.running
|
|
|
|
async def get_auto_organize_lock(self) -> asyncio.Lock:
|
|
return self._lock
|
|
|
|
|
|
class StubFileService(ModelFileService):
|
|
def __init__(self) -> None:
|
|
super().__init__(scanner=None, model_type="lora")
|
|
self.calls: List[Dict[str, Any]] = []
|
|
|
|
async def auto_organize_models(
|
|
self,
|
|
file_paths: Optional[List[str]] = None,
|
|
progress_callback=None,
|
|
exclusion_patterns=None,
|
|
) -> AutoOrganizeResult:
|
|
result = AutoOrganizeResult()
|
|
result.total = len(file_paths or [])
|
|
self.calls.append({
|
|
"file_paths": file_paths,
|
|
"progress_callback": progress_callback,
|
|
"exclusion_patterns": exclusion_patterns,
|
|
})
|
|
return result
|
|
|
|
|
|
class StubMetadataSync(MetadataSyncService):
|
|
def __init__(self) -> None:
|
|
super().__init__(
|
|
metadata_manager=object(),
|
|
preview_service=object(),
|
|
settings=StubSettings(), # pyright: ignore[reportArgumentType]
|
|
default_metadata_provider_factory=lambda: asyncio.sleep(0, result=None), # pyright: ignore[reportArgumentType]
|
|
metadata_provider_selector=lambda _name=None: asyncio.sleep(0, result=None), # pyright: ignore[reportArgumentType]
|
|
)
|
|
self.calls: List[Dict[str, Any]] = []
|
|
|
|
async def fetch_and_update_model(self, **kwargs: Any):
|
|
self.calls.append(kwargs)
|
|
model_data = kwargs["model_data"]
|
|
model_data["model_name"] = model_data.get("model_name", "model") + "-updated"
|
|
return True, None
|
|
|
|
|
|
@dataclass
|
|
class StubSettings:
|
|
enable_metadata_archive_db: bool = False
|
|
|
|
def get(self, key: str, default: Any = None) -> Any:
|
|
if key == "enable_metadata_archive_db":
|
|
return self.enable_metadata_archive_db
|
|
return default
|
|
|
|
|
|
class ProgressCollector:
|
|
def __init__(self) -> None:
|
|
self.events: List[Dict[str, Any]] = []
|
|
|
|
async def on_progress(self, payload: Dict[str, Any]) -> None:
|
|
self.events.append(payload)
|
|
|
|
|
|
class StubDownloadCoordinator(DownloadCoordinator):
|
|
def __init__(self, *, error: Optional[str] = None) -> None:
|
|
super().__init__(
|
|
ws_manager=SimpleNamespace(generate_download_id=lambda: "abc123"),
|
|
download_manager_factory=lambda: asyncio.sleep(0, result=None),
|
|
)
|
|
self.error = error
|
|
self.payloads: List[Dict[str, Any]] = []
|
|
|
|
async def schedule_download(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
self.payloads.append(payload)
|
|
if self.error == "validation":
|
|
raise ValueError("Missing required parameter: Please provide either 'model_id' or 'model_version_id'")
|
|
if self.error == "401":
|
|
raise RuntimeError("401 Unauthorized")
|
|
return {"success": True, "download_id": "abc123"}
|
|
|
|
|
|
class StubExampleImagesDownloadManager:
|
|
def __init__(self) -> None:
|
|
self.payloads: List[Dict[str, Any]] = []
|
|
self.error: Optional[str] = None
|
|
self.progress_snapshot = {"status": "running"}
|
|
|
|
async def start_download(self, payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
self.payloads.append(payload)
|
|
if self.error == "in_progress":
|
|
raise DownloadInProgressError(self.progress_snapshot)
|
|
if self.error == "configuration":
|
|
raise DownloadConfigurationError("path missing")
|
|
if self.error == "generic":
|
|
raise ExampleImagesDownloadError("boom")
|
|
return {"success": True, "message": "ok"}
|
|
|
|
|
|
class StubExampleImagesProcessor(ExampleImagesProcessor):
|
|
def __init__(self) -> None:
|
|
self.calls: List[Dict[str, Any]] = []
|
|
self.error: Optional[str] = None
|
|
self.response: Dict[str, Any] = {"success": True}
|
|
self.resolved_hash: Optional[str] = None
|
|
self.resolve_calls: List[str] = []
|
|
|
|
async def import_images(self, model_hash: str, files: List[str]) -> Dict[str, Any]: # pyright: ignore[reportIncompatibleMethodOverride]
|
|
self.calls.append({"model_hash": model_hash, "files": files})
|
|
if self.error == "validation":
|
|
raise ExampleImagesValidationError("missing")
|
|
if self.error == "generic":
|
|
raise ExampleImagesImportError("boom")
|
|
return self.response
|
|
|
|
async def resolve_hash_for_file_path(self, file_path: str) -> str:
|
|
self.resolve_calls.append(file_path)
|
|
return self.resolved_hash or ""
|
|
|
|
|
|
async def test_auto_organize_use_case_executes_with_lock() -> None:
|
|
file_service = StubFileService()
|
|
lock_provider = StubLockProvider()
|
|
use_case = AutoOrganizeUseCase(file_service=file_service, lock_provider=lock_provider)
|
|
|
|
result = await use_case.execute(file_paths=["model1"], progress_callback=None)
|
|
|
|
assert isinstance(result, AutoOrganizeResult)
|
|
assert file_service.calls[0]["file_paths"] == ["model1"]
|
|
assert file_service.calls[0]["exclusion_patterns"] is None
|
|
|
|
|
|
async def test_auto_organize_use_case_rejects_when_running() -> None:
|
|
file_service = StubFileService()
|
|
lock_provider = StubLockProvider()
|
|
lock_provider.running = True
|
|
use_case = AutoOrganizeUseCase(file_service=file_service, lock_provider=lock_provider)
|
|
|
|
with pytest.raises(AutoOrganizeInProgressError):
|
|
await use_case.execute(file_paths=None, progress_callback=None)
|
|
|
|
|
|
async def test_bulk_metadata_refresh_emits_progress_and_updates_cache(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
scanner = MockScanner()
|
|
scanner._cache.raw_data = [
|
|
{
|
|
"file_path": "model1.safetensors",
|
|
"sha256": "hash",
|
|
"from_civitai": True,
|
|
"model_name": "Demo",
|
|
}
|
|
]
|
|
service = MockModelService(scanner)
|
|
metadata_sync = StubMetadataSync()
|
|
settings = StubSettings()
|
|
progress = ProgressCollector()
|
|
|
|
hydration_calls: list[str] = []
|
|
|
|
async def fake_hydrate(model_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
hydration_calls.append(model_data.get("file_path", ""))
|
|
model_data.clear()
|
|
model_data.update(
|
|
{
|
|
"file_path": "model1.safetensors",
|
|
"sha256": "hash",
|
|
"from_civitai": True,
|
|
"model_name": "Demo",
|
|
"extra": "value",
|
|
"civitai": {"images": [{"url": "existing.png", "type": "image"}]},
|
|
}
|
|
)
|
|
return model_data
|
|
|
|
monkeypatch.setattr(MetadataManager, "hydrate_model_data", staticmethod(fake_hydrate))
|
|
|
|
use_case = BulkMetadataRefreshUseCase(
|
|
service=service,
|
|
metadata_sync=metadata_sync,
|
|
settings_service=settings,
|
|
logger=logging.getLogger("test"),
|
|
)
|
|
|
|
result = await use_case.execute_with_error_handling(progress_callback=progress)
|
|
|
|
assert result["success"] is True
|
|
assert progress.events[0]["status"] == "started"
|
|
assert progress.events[-1]["status"] == "completed"
|
|
assert metadata_sync.calls
|
|
assert metadata_sync.calls[0]["model_data"]["extra"] == "value"
|
|
assert scanner._cache.raw_data[0]["extra"] == "value"
|
|
assert hydration_calls == ["model1.safetensors"]
|
|
assert scanner._cache.resort_calls == 1
|
|
|
|
|
|
async def test_bulk_metadata_refresh_reports_errors() -> None:
|
|
class FailingScanner(MockScanner):
|
|
async def get_cached_data(self, force_refresh: bool = False):
|
|
raise RuntimeError("boom")
|
|
|
|
service = MockModelService(FailingScanner())
|
|
metadata_sync = StubMetadataSync()
|
|
settings = StubSettings()
|
|
progress = ProgressCollector()
|
|
|
|
use_case = BulkMetadataRefreshUseCase(
|
|
service=service,
|
|
metadata_sync=metadata_sync,
|
|
settings_service=settings,
|
|
logger=logging.getLogger("test"),
|
|
)
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await use_case.execute_with_error_handling(progress_callback=progress)
|
|
|
|
assert progress.events
|
|
assert progress.events[-1]["status"] == "error"
|
|
assert progress.events[-1]["error"] == "boom"
|
|
|
|
|
|
async def test_bulk_metadata_refresh_skips_confirmed_not_found_models(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Models marked as from_civitai=False and civitai_deleted=True should be skipped."""
|
|
scanner = MockScanner()
|
|
scanner._cache.raw_data = [
|
|
{
|
|
"file_path": "model1.safetensors",
|
|
"sha256": "hash1",
|
|
"from_civitai": False,
|
|
"civitai_deleted": True,
|
|
"model_name": "NotOnCivitAI",
|
|
},
|
|
{
|
|
"file_path": "model2.safetensors",
|
|
"sha256": "hash2",
|
|
"from_civitai": True,
|
|
"model_name": "OnCivitAI",
|
|
},
|
|
]
|
|
service = MockModelService(scanner)
|
|
metadata_sync = StubMetadataSync()
|
|
settings = StubSettings(enable_metadata_archive_db=False)
|
|
progress = ProgressCollector()
|
|
|
|
async def fake_hydrate(model_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
# Preserve the original data (simulating no metadata file on disk)
|
|
return model_data
|
|
|
|
monkeypatch.setattr(MetadataManager, "hydrate_model_data", staticmethod(fake_hydrate))
|
|
|
|
use_case = BulkMetadataRefreshUseCase(
|
|
service=service,
|
|
metadata_sync=metadata_sync,
|
|
settings_service=settings,
|
|
logger=logging.getLogger("test"),
|
|
)
|
|
|
|
result = await use_case.execute_with_error_handling(progress_callback=progress)
|
|
|
|
assert result["success"] is True
|
|
# Only model2 should be processed (model1 is skipped)
|
|
assert result["processed"] == 1
|
|
assert result["updated"] == 1
|
|
assert len(metadata_sync.calls) == 1
|
|
assert metadata_sync.calls[0]["file_path"] == "model2.safetensors"
|
|
|
|
|
|
async def test_bulk_metadata_refresh_skips_when_archive_checked(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Models with db_checked=True should be skipped even if archive DB is enabled."""
|
|
scanner = MockScanner()
|
|
scanner._cache.raw_data = [
|
|
{
|
|
"file_path": "model1.safetensors",
|
|
"sha256": "hash1",
|
|
"from_civitai": False,
|
|
"civitai_deleted": True,
|
|
"db_checked": True,
|
|
"model_name": "ArchiveChecked",
|
|
},
|
|
{
|
|
"file_path": "model2.safetensors",
|
|
"sha256": "hash2",
|
|
"from_civitai": False,
|
|
"civitai_deleted": True,
|
|
"db_checked": False,
|
|
"model_name": "ArchiveNotChecked",
|
|
},
|
|
]
|
|
service = MockModelService(scanner)
|
|
metadata_sync = StubMetadataSync()
|
|
settings = StubSettings(enable_metadata_archive_db=True)
|
|
progress = ProgressCollector()
|
|
|
|
async def fake_hydrate(model_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
return model_data
|
|
|
|
monkeypatch.setattr(MetadataManager, "hydrate_model_data", staticmethod(fake_hydrate))
|
|
|
|
use_case = BulkMetadataRefreshUseCase(
|
|
service=service,
|
|
metadata_sync=metadata_sync,
|
|
settings_service=settings,
|
|
logger=logging.getLogger("test"),
|
|
)
|
|
|
|
result = await use_case.execute_with_error_handling(progress_callback=progress)
|
|
|
|
assert result["success"] is True
|
|
# Only model2 should be processed (model1 has db_checked=True)
|
|
assert result["processed"] == 1
|
|
assert result["updated"] == 1
|
|
assert len(metadata_sync.calls) == 1
|
|
assert metadata_sync.calls[0]["file_path"] == "model2.safetensors"
|
|
|
|
|
|
async def test_bulk_metadata_refresh_processes_never_fetched_models(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Models that have never been fetched (from_civitai=None) should be processed."""
|
|
scanner = MockScanner()
|
|
scanner._cache.raw_data = [
|
|
{
|
|
"file_path": "model1.safetensors",
|
|
"sha256": "hash1",
|
|
"from_civitai": None,
|
|
"model_name": "NeverFetched",
|
|
},
|
|
{
|
|
"file_path": "model2.safetensors",
|
|
"sha256": "hash2",
|
|
"model_name": "NoFromCivitaiField",
|
|
},
|
|
]
|
|
service = MockModelService(scanner)
|
|
metadata_sync = StubMetadataSync()
|
|
settings = StubSettings(enable_metadata_archive_db=False)
|
|
progress = ProgressCollector()
|
|
|
|
async def fake_hydrate(model_data: Dict[str, Any]) -> Dict[str, Any]:
|
|
return model_data
|
|
|
|
monkeypatch.setattr(MetadataManager, "hydrate_model_data", staticmethod(fake_hydrate))
|
|
|
|
use_case = BulkMetadataRefreshUseCase(
|
|
service=service,
|
|
metadata_sync=metadata_sync,
|
|
settings_service=settings,
|
|
logger=logging.getLogger("test"),
|
|
)
|
|
|
|
result = await use_case.execute_with_error_handling(progress_callback=progress)
|
|
|
|
assert result["success"] is True
|
|
# Both models should be processed
|
|
assert result["processed"] == 2
|
|
assert result["updated"] == 2
|
|
assert len(metadata_sync.calls) == 2
|
|
|
|
|
|
async def test_download_model_use_case_raises_validation_error() -> None:
|
|
coordinator = StubDownloadCoordinator(error="validation")
|
|
use_case = DownloadModelUseCase(download_coordinator=coordinator)
|
|
|
|
with pytest.raises(DownloadModelValidationError):
|
|
await use_case.execute({})
|
|
|
|
|
|
async def test_download_model_use_case_raises_early_access() -> None:
|
|
coordinator = StubDownloadCoordinator(error="401")
|
|
use_case = DownloadModelUseCase(download_coordinator=coordinator)
|
|
|
|
with pytest.raises(DownloadModelEarlyAccessError):
|
|
await use_case.execute({"model_id": 1})
|
|
|
|
|
|
async def test_download_model_use_case_returns_result() -> None:
|
|
coordinator = StubDownloadCoordinator()
|
|
use_case = DownloadModelUseCase(download_coordinator=coordinator)
|
|
|
|
result = await use_case.execute({"model_id": 1})
|
|
|
|
assert result["success"] is True
|
|
assert result["download_id"] == "abc123"
|
|
|
|
|
|
async def test_download_example_images_use_case_triggers_manager() -> None:
|
|
manager = StubExampleImagesDownloadManager()
|
|
use_case = DownloadExampleImagesUseCase(download_manager=manager)
|
|
|
|
payload = {"optimize": True}
|
|
result = await use_case.execute(payload)
|
|
|
|
assert manager.payloads == [payload]
|
|
assert result == {"success": True, "message": "ok"}
|
|
|
|
|
|
async def test_download_example_images_use_case_maps_in_progress() -> None:
|
|
manager = StubExampleImagesDownloadManager()
|
|
manager.error = "in_progress"
|
|
use_case = DownloadExampleImagesUseCase(download_manager=manager)
|
|
|
|
with pytest.raises(DownloadExampleImagesInProgressError) as exc:
|
|
await use_case.execute({})
|
|
|
|
assert exc.value.progress == manager.progress_snapshot
|
|
|
|
|
|
async def test_download_example_images_use_case_maps_configuration() -> None:
|
|
manager = StubExampleImagesDownloadManager()
|
|
manager.error = "configuration"
|
|
use_case = DownloadExampleImagesUseCase(download_manager=manager)
|
|
|
|
with pytest.raises(DownloadExampleImagesConfigurationError):
|
|
await use_case.execute({})
|
|
|
|
|
|
async def test_download_example_images_use_case_propagates_generic_error() -> None:
|
|
manager = StubExampleImagesDownloadManager()
|
|
manager.error = "generic"
|
|
use_case = DownloadExampleImagesUseCase(download_manager=manager)
|
|
|
|
with pytest.raises(ExampleImagesDownloadError):
|
|
await use_case.execute({})
|
|
|
|
|
|
class DummyJsonRequest:
|
|
def __init__(self, payload: Dict[str, Any]) -> None:
|
|
self._payload = payload
|
|
self.content_type = "application/json"
|
|
|
|
async def json(self) -> Dict[str, Any]:
|
|
return self._payload
|
|
|
|
|
|
async def test_import_example_images_use_case_delegates() -> None:
|
|
processor = StubExampleImagesProcessor()
|
|
use_case = ImportExampleImagesUseCase(processor=processor)
|
|
|
|
request = DummyJsonRequest({"model_hash": "abc", "file_paths": ["/tmp/file"]})
|
|
result = await use_case.execute(request) # pyright: ignore[reportArgumentType]
|
|
|
|
assert processor.calls == [{"model_hash": "abc", "files": ["/tmp/file"]}]
|
|
assert result == {"success": True}
|
|
|
|
|
|
async def test_import_example_images_use_case_maps_validation_error() -> None:
|
|
processor = StubExampleImagesProcessor()
|
|
processor.error = "validation"
|
|
use_case = ImportExampleImagesUseCase(processor=processor)
|
|
request = DummyJsonRequest({"model_hash": None, "file_paths": []})
|
|
|
|
with pytest.raises(ImportExampleImagesValidationError):
|
|
await use_case.execute(request) # pyright: ignore[reportArgumentType]
|
|
|
|
|
|
async def test_import_example_images_use_case_propagates_generic_error() -> None:
|
|
processor = StubExampleImagesProcessor()
|
|
processor.error = "generic"
|
|
use_case = ImportExampleImagesUseCase(processor=processor)
|
|
request = DummyJsonRequest({"model_hash": "abc", "file_paths": ["/tmp/file"]})
|
|
|
|
with pytest.raises(ExampleImagesImportError):
|
|
await use_case.execute(request) # pyright: ignore[reportArgumentType]
|
|
|
|
|
|
async def test_import_example_images_use_case_resolves_hash_from_model_path() -> None:
|
|
"""Models with a deferred hash (checkpoints, Other) send model_path
|
|
instead of model_hash; the use case must resolve the hash on demand."""
|
|
processor = StubExampleImagesProcessor()
|
|
processor.resolved_hash = "f" * 64
|
|
use_case = ImportExampleImagesUseCase(processor=processor)
|
|
|
|
request = DummyJsonRequest(
|
|
{"model_path": "/models/vae/x.safetensors", "file_paths": ["/tmp/file"]}
|
|
)
|
|
result = await use_case.execute(request) # pyright: ignore[reportArgumentType]
|
|
|
|
assert processor.resolve_calls == ["/models/vae/x.safetensors"]
|
|
assert processor.calls == [{"model_hash": "f" * 64, "files": ["/tmp/file"]}]
|
|
assert result == {"success": True}
|
|
|
|
|
|
async def test_import_example_images_use_case_rejects_unresolvable_model_path() -> None:
|
|
processor = StubExampleImagesProcessor()
|
|
use_case = ImportExampleImagesUseCase(processor=processor)
|
|
request = DummyJsonRequest(
|
|
{"model_path": "/models/unknown.safetensors", "file_paths": []}
|
|
)
|
|
|
|
with pytest.raises(ImportExampleImagesValidationError):
|
|
await use_case.execute(request) # pyright: ignore[reportArgumentType]
|
|
|
|
|
|
class StubLifecycleService:
|
|
def __init__(self, scanner: Optional[MockScanner] = None) -> None:
|
|
self.renames: List[Dict[str, str]] = []
|
|
self.error: Optional[Exception] = None
|
|
self.cancel_on_rename = False
|
|
self._scanner = scanner
|
|
|
|
@asynccontextmanager
|
|
async def bulk_rename_session(self) -> AsyncIterator[None]:
|
|
yield None
|
|
|
|
async def rename_model(
|
|
self, *, file_path: str, new_file_name: str, bulk_context: Any = None
|
|
) -> Dict[str, Any]:
|
|
if self.error is not None:
|
|
raise self.error
|
|
self.renames.append({"file_path": file_path, "new_file_name": new_file_name})
|
|
if self.cancel_on_rename and self._scanner is not None:
|
|
self._scanner.cancel_task()
|
|
return {"success": True, "new_file_path": file_path}
|
|
|
|
|
|
def _filename_template_model(
|
|
file_path: str,
|
|
model_name: str,
|
|
sha256: str = "abcdef0123456789",
|
|
) -> Dict[str, Any]:
|
|
return {
|
|
"file_path": file_path,
|
|
"file_name": file_path.rsplit("/", 1)[-1].rsplit(".", 1)[0],
|
|
"model_name": model_name,
|
|
"sha256": sha256,
|
|
"civitai": {"id": 1},
|
|
}
|
|
|
|
|
|
def _set_filename_template(template: str, model_type: str = "lora") -> None:
|
|
from py.services.settings_manager import get_settings_manager
|
|
|
|
manager = get_settings_manager()
|
|
templates = dict(manager.settings.get("download_filename_templates") or {})
|
|
templates[model_type] = template
|
|
manager.settings["download_filename_templates"] = templates
|
|
|
|
|
|
def _make_filename_template_use_case(
|
|
scanner: MockScanner,
|
|
lifecycle: StubLifecycleService,
|
|
lock_provider: Optional[StubLockProvider] = None,
|
|
metadata_loader: Optional[Any] = None,
|
|
) -> FilenameTemplateUseCase:
|
|
kwargs: Dict[str, Any] = {}
|
|
if metadata_loader is not None:
|
|
kwargs["metadata_loader"] = metadata_loader
|
|
return FilenameTemplateUseCase(
|
|
scanner=scanner,
|
|
lifecycle_service=lifecycle, # pyright: ignore[reportArgumentType]
|
|
lock_provider=lock_provider or StubLockProvider(),
|
|
model_type="lora",
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
async def test_filename_template_use_case_renames_models() -> None:
|
|
_set_filename_template("{model_name}-{hash_short}")
|
|
scanner = MockScanner(cache=MockCache([
|
|
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
|
_filename_template_model("/library/beta.safetensors", "Beta"),
|
|
]))
|
|
lifecycle = StubLifecycleService()
|
|
progress = ProgressCollector()
|
|
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
|
|
|
result = await use_case.execute(progress_callback=progress)
|
|
|
|
assert result.status == "success"
|
|
assert result.operation_type == "filename_template"
|
|
assert result.total == 2
|
|
assert result.success_count == 2
|
|
assert result.failure_count == 0
|
|
assert lifecycle.renames == [
|
|
{"file_path": "/library/alpha.safetensors", "new_file_name": "Alpha-abcdef0123"},
|
|
{"file_path": "/library/beta.safetensors", "new_file_name": "Beta-abcdef0123"},
|
|
]
|
|
statuses = [event["status"] for event in progress.events]
|
|
assert statuses[0] == "started"
|
|
assert statuses[-1] == "completed"
|
|
assert all(event["type"] == "filename_template_progress" for event in progress.events)
|
|
|
|
|
|
async def test_filename_template_use_case_skips_unchanged_names() -> None:
|
|
_set_filename_template("{model_name}-{hash_short}")
|
|
scanner = MockScanner(cache=MockCache([
|
|
_filename_template_model("/library/Alpha-abcdef0123.safetensors", "Alpha"),
|
|
]))
|
|
lifecycle = StubLifecycleService()
|
|
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
|
|
|
result = await use_case.execute(progress_callback=None)
|
|
|
|
assert result.success_count == 0
|
|
assert result.skipped_count == 1
|
|
assert lifecycle.renames == []
|
|
|
|
|
|
async def test_filename_template_use_case_reverts_to_recorded_original_when_template_empty() -> None:
|
|
_set_filename_template("")
|
|
scanner = MockScanner(cache=MockCache([
|
|
_filename_template_model("/library/alpha-renamed.safetensors", "Alpha"),
|
|
_filename_template_model("/library/beta.safetensors", "Beta"),
|
|
]))
|
|
lifecycle = StubLifecycleService()
|
|
|
|
async def metadata_loader(metadata_path: str) -> Dict[str, Any]:
|
|
if metadata_path == "/library/alpha-renamed.metadata.json":
|
|
return {"original_file_name": "alpha-original"}
|
|
return {}
|
|
|
|
use_case = _make_filename_template_use_case(
|
|
scanner, lifecycle, metadata_loader=metadata_loader
|
|
)
|
|
|
|
result = await use_case.execute(progress_callback=None)
|
|
|
|
assert result.success_count == 1
|
|
assert result.skipped_count == 1
|
|
assert lifecycle.renames == [
|
|
{
|
|
"file_path": "/library/alpha-renamed.safetensors",
|
|
"new_file_name": "alpha-original",
|
|
}
|
|
]
|
|
|
|
|
|
async def test_filename_template_use_case_skips_revert_without_recorded_original() -> None:
|
|
_set_filename_template("")
|
|
scanner = MockScanner(cache=MockCache([
|
|
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
|
]))
|
|
lifecycle = StubLifecycleService()
|
|
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
|
|
|
result = await use_case.execute(progress_callback=None)
|
|
|
|
assert result.skipped_count == 1
|
|
assert lifecycle.renames == []
|
|
|
|
|
|
async def test_filename_template_use_case_skips_revert_matching_current_name() -> None:
|
|
_set_filename_template("")
|
|
scanner = MockScanner(cache=MockCache([
|
|
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
|
]))
|
|
lifecycle = StubLifecycleService()
|
|
|
|
async def metadata_loader(metadata_path: str) -> Dict[str, Any]:
|
|
return {"original_file_name": "alpha"}
|
|
|
|
use_case = _make_filename_template_use_case(
|
|
scanner, lifecycle, metadata_loader=metadata_loader
|
|
)
|
|
|
|
result = await use_case.execute(progress_callback=None)
|
|
|
|
assert result.success_count == 0
|
|
assert result.skipped_count == 1
|
|
assert lifecycle.renames == []
|
|
|
|
|
|
async def test_filename_template_use_case_counts_conflicts_as_failures() -> None:
|
|
_set_filename_template("{model_name}")
|
|
scanner = MockScanner(cache=MockCache([
|
|
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
|
_filename_template_model("/library/beta.safetensors", "Beta"),
|
|
]))
|
|
lifecycle = StubLifecycleService()
|
|
lifecycle.error = ValueError("A file with this name already exists")
|
|
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
|
|
|
result = await use_case.execute(progress_callback=None)
|
|
|
|
assert result.status == "success"
|
|
assert result.failure_count == 2
|
|
assert result.success_count == 0
|
|
assert len(result.results) == 2
|
|
|
|
|
|
async def test_filename_template_use_case_honours_cancellation() -> None:
|
|
_set_filename_template("{model_name}-{hash_short}")
|
|
scanner = MockScanner(cache=MockCache([
|
|
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
|
_filename_template_model("/library/beta.safetensors", "Beta"),
|
|
]))
|
|
lifecycle = StubLifecycleService(scanner=scanner)
|
|
lifecycle.cancel_on_rename = True
|
|
progress = ProgressCollector()
|
|
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
|
|
|
result = await use_case.execute(progress_callback=progress)
|
|
|
|
assert result.status == "cancelled"
|
|
assert len(lifecycle.renames) == 1
|
|
assert progress.events[-1]["status"] == "cancelled"
|
|
|
|
|
|
async def test_filename_template_use_case_filters_file_paths() -> None:
|
|
_set_filename_template("{model_name}-{hash_short}")
|
|
scanner = MockScanner(cache=MockCache([
|
|
_filename_template_model("/library/alpha.safetensors", "Alpha"),
|
|
_filename_template_model("/library/beta.safetensors", "Beta"),
|
|
]))
|
|
lifecycle = StubLifecycleService()
|
|
use_case = _make_filename_template_use_case(scanner, lifecycle)
|
|
|
|
result = await use_case.execute(
|
|
file_paths=["/library/beta.safetensors"], progress_callback=None
|
|
)
|
|
|
|
assert result.total == 1
|
|
assert lifecycle.renames == [
|
|
{"file_path": "/library/beta.safetensors", "new_file_name": "Beta-abcdef0123"}
|
|
]
|
|
|
|
|
|
async def test_filename_template_use_case_rejects_when_lock_held() -> None:
|
|
_set_filename_template("{model_name}")
|
|
scanner = MockScanner(cache=MockCache())
|
|
lifecycle = StubLifecycleService()
|
|
lock_provider = StubLockProvider()
|
|
lock_provider.running = True
|
|
use_case = _make_filename_template_use_case(scanner, lifecycle, lock_provider)
|
|
|
|
with pytest.raises(AutoOrganizeInProgressError):
|
|
await use_case.execute(progress_callback=None) |