import asyncio import logging from dataclasses import dataclass from types import SimpleNamespace from typing import Any, 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} 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 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] 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 async def rename_model(self, *, file_path: str, new_file_name: str) -> 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)