mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-12 12:32:11 -03:00
Will Miao
c16de68640
The new informational Windows job surfaced 91 failures: four real bugs (fixed in
the previous two commits) and a long tail of POSIX assumptions in the tests
themselves. None of them assert behaviour that differs on Windows, so they are
fixed rather than skipped wherever the behaviour is the same:
* Path spelling: expectations built with `str(tmp_path)` are compared through an
identity helper (`normalize_root_identity`) or the platform's separator, since
the code hands out forward-slashed business paths while `str()` is native.
* Root labels: `_root_display_labels` prefixes the drive on Windows ("C: loras"),
so label assertions compare the final path component or derive the expected
label from the same helper.
* `~` expansion reads USERPROFILE (ntpath), not HOME; timer granularity rounds
0.1s down; a case-only rename counts as "already named" via `os.path.normcase`.
* Genuinely POSIX-only strategies are skipped with a reason instead of failing:
the headless clipboard fallback, dbus file-manager selection, xdg-open, and
the case-fold index test that simulates a case-sensitive filesystem.
781 lines
28 KiB
Python
781 lines
28 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
|
|
import os
|
|
|
|
|
|
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]:
|
|
# Compare the basename: the sidecar path is platform-spelled.
|
|
if os.path.basename(metadata_path) == "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([
|
|
# The names differ beyond case: on Windows a case-only difference counts
|
|
# as "already named that" (os.path.normcase) and would be skipped.
|
|
_filename_template_model("/library/alpha.safetensors", "Alpha Two"),
|
|
_filename_template_model("/library/beta.safetensors", "Beta Two"),
|
|
]))
|
|
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) |