"""DownloadManager support for the "other" model type (VAE, upscaler, ...). Covers the Phase-2 scatter points from docs/plans/other-models-page.md §9.1: type map acceptance, existence gates consulting the other scanner (never falling through to the lora scanner), per-sub_type default roots, resume metadata and the archive extension set. """ from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock import pytest from py.services import aria2_transfer_state from py.services import download_manager from py.services.download_manager import DownloadManager from py.services.service_registry import ServiceRegistry from py.services.settings_manager import SettingsManager, get_settings_manager @pytest.fixture(autouse=True) def reset_download_manager(): """Ensure each test operates on a fresh singleton.""" DownloadManager._instance = None yield DownloadManager._instance = None @pytest.fixture(autouse=True) def isolate_settings(monkeypatch, tmp_path): """Point settings writes at a temporary directory to avoid touching real files.""" manager = get_settings_manager() default_settings = manager._get_default_settings() default_settings.update( { "default_lora_root": str(tmp_path / "loras"), "default_checkpoint_root": str(tmp_path / "checkpoints"), "default_embedding_root": str(tmp_path / "embeddings"), "default_other_roots": { "vae": str(tmp_path / "vae"), "upscaler": str(tmp_path / "upscale_models"), "text_encoder": str(tmp_path / "text_encoders"), "clip_vision": str(tmp_path / "clip_vision"), }, "download_path_templates": { "lora": "{base_model}/{first_tag}", "checkpoint": "{base_model}/{first_tag}", "embedding": "{base_model}/{first_tag}", "other": "", }, "skip_previously_downloaded_model_versions": False, "download_skip_base_models": [], } ) monkeypatch.setattr(manager, "settings", default_settings) monkeypatch.setattr(SettingsManager, "_save_settings", lambda self: None) @pytest.fixture(autouse=True) def isolate_aria2_state(monkeypatch, tmp_path): state_path = tmp_path / "cache" / "aria2" / "downloads.json" monkeypatch.setattr( aria2_transfer_state, "get_aria2_state_path", lambda: str(state_path), ) @pytest.fixture(autouse=True) def stub_metadata(monkeypatch): class _StubMetadata: def __init__(self, save_path: str): self.file_path = save_path self.sha256 = "sha256" self.file_name = Path(save_path).stem def _make_class(name): @staticmethod def from_civitai_info(_version_info, _file_info, save_path): metadata = _StubMetadata(save_path) metadata.metadata_class = name return metadata return type(name, (), {"from_civitai_info": from_civitai_info}) monkeypatch.setattr(download_manager, "LoraMetadata", _make_class("LoraMetadata")) monkeypatch.setattr( download_manager, "CheckpointMetadata", _make_class("CheckpointMetadata") ) monkeypatch.setattr( download_manager, "EmbeddingMetadata", _make_class("EmbeddingMetadata") ) monkeypatch.setattr( download_manager, "OtherModelMetadata", _make_class("OtherModelMetadata") ) class DummyScanner: def __init__(self, exists: bool = False, raw_data=None): self.exists = exists self.calls = [] self._cache = SimpleNamespace(raw_data=list(raw_data or [])) async def check_model_version_exists(self, version_id): self.calls.append(version_id) return self.exists async def get_cached_data(self): return self._cache @pytest.fixture def scanners(monkeypatch): lora_scanner = DummyScanner() checkpoint_scanner = DummyScanner() embedding_scanner = DummyScanner() other_scanner = DummyScanner() monkeypatch.setattr( ServiceRegistry, "get_lora_scanner", AsyncMock(return_value=lora_scanner) ) monkeypatch.setattr( ServiceRegistry, "get_checkpoint_scanner", AsyncMock(return_value=checkpoint_scanner), ) monkeypatch.setattr( ServiceRegistry, "get_embedding_scanner", AsyncMock(return_value=embedding_scanner), ) monkeypatch.setattr( ServiceRegistry, "get_other_scanner", AsyncMock(return_value=other_scanner), ) return SimpleNamespace( lora=lora_scanner, checkpoint=checkpoint_scanner, embedding=embedding_scanner, other=other_scanner, ) def _other_payload(civitai_type: str, *, files=None) -> dict: return { "id": 42, "model": {"type": civitai_type, "tags": ["utility"]}, "baseModel": "SDXL 1.0", "creator": {"username": "Author"}, "files": files or [ { "type": "Model", "primary": True, "downloadUrl": "https://example.invalid/file.safetensors", "name": "file.safetensors", } ], } @pytest.fixture def metadata_provider(monkeypatch): class DummyProvider: def __init__(self): self.calls = [] self.payload = _other_payload("VAE") async def get_model_version(self, model_id, model_version_id): self.calls.append((model_id, model_version_id)) return self.payload provider = DummyProvider() monkeypatch.setattr( download_manager, "get_default_metadata_provider", AsyncMock(return_value=provider), ) return provider def _capture_execute(monkeypatch, captured): async def fake_execute_download(self, **kwargs): captured.update(kwargs) return {"success": True} monkeypatch.setattr( DownloadManager, "_execute_download", fake_execute_download, raising=False ) @pytest.mark.asyncio @pytest.mark.parametrize( "civitai_type", ["VAE", "Upscaler", "TextEncoder", "CLIP", "CLIPVision", "Controlnet", "Other"], ) async def test_download_accepts_other_model_types( monkeypatch, scanners, metadata_provider, tmp_path, civitai_type ): """All VALID_OTHER_CIVITAI_TYPES route to model_type 'other'.""" metadata_provider.payload = _other_payload(civitai_type) captured = {} _capture_execute(monkeypatch, captured) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=99, save_dir=str(tmp_path) ) assert result["success"] is True assert captured["model_type"] == "other" assert captured["metadata"].metadata_class == "OtherModelMetadata" @pytest.mark.asyncio async def test_download_rejects_unknown_model_type( monkeypatch, scanners, metadata_provider, tmp_path ): metadata_provider.payload = _other_payload("Workflow") manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=99, save_dir=str(tmp_path) ) assert result["success"] is False assert result["error"].startswith("Model type") @pytest.mark.asyncio async def test_early_gate_checks_other_scanner( monkeypatch, scanners, metadata_provider, tmp_path ): scanners.other.exists = True manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=101, save_dir=str(tmp_path) ) assert result["success"] is False assert result["error"] == "Model version already exists in other library" assert scanners.other.calls == [101] @pytest.mark.asyncio async def test_scanner_dispatch_has_no_lora_fall_through(scanners): """The Phase-2 trap: 'other' must reach the other scanner explicitly, and unknown types must raise instead of silently deduping against loras.""" manager = DownloadManager() scanner = await manager._get_scanner_for_model_type("other") assert scanner is scanners.other scanner = await manager._get_scanner_for_model_type("lora") assert scanner is scanners.lora with pytest.raises(ValueError): await manager._get_scanner_for_model_type("bogus") @pytest.mark.asyncio async def test_explicit_file_gate_uses_other_scanner_not_lora( monkeypatch, scanners, metadata_provider, tmp_path ): """A matching local entry in the LORA library must not block an 'other' download; only the other scanner's library is consulted.""" local_entry = { "file_name": "file", "sha256": "deadbeef", "civitai": {"id": 42}, } scanners.lora._cache = SimpleNamespace(raw_data=[dict(local_entry)]) scanners.other._cache = SimpleNamespace(raw_data=[]) metadata_provider.payload = _other_payload( "VAE", files=[ { "id": 7, "type": "Model", "primary": True, "name": "file.safetensors", "hashes": {"SHA256": "deadbeef"}, "downloadUrl": "https://example.invalid/file.safetensors", } ], ) captured = {} _capture_execute(monkeypatch, captured) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=42, save_dir=str(tmp_path), file_params={"id": 7, "type": "Model"}, ) assert result["success"] is True assert captured["model_type"] == "other" @pytest.mark.asyncio async def test_explicit_file_gate_blocks_when_other_scanner_has_file( monkeypatch, scanners, metadata_provider, tmp_path ): local_entry = { "file_name": "file", "sha256": "deadbeef", "civitai": {"id": 42}, } scanners.other._cache = SimpleNamespace(raw_data=[dict(local_entry)]) metadata_provider.payload = _other_payload( "VAE", files=[ { "id": 7, "type": "Model", "primary": True, "name": "file.safetensors", "hashes": {"SHA256": "deadbeef"}, "downloadUrl": "https://example.invalid/file.safetensors", } ], ) execute_mock = AsyncMock(return_value={"success": True}) monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=42, save_dir=str(tmp_path), file_params={"id": 7, "type": "Model"}, ) assert result["success"] is False assert "already exists in other library" in result["error"] assert execute_mock.await_count == 0 @pytest.mark.asyncio async def test_version_level_fallback_gate_checks_other_scanner( monkeypatch, scanners, metadata_provider, tmp_path ): """file_params that resolve to nothing fall back to the version-level gate, which must consult the other scanner.""" scanners.other.exists = True execute_mock = AsyncMock(return_value={"success": True}) monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=101, save_dir=str(tmp_path), file_params={"id": 999999, "type": "Model"}, ) assert result["success"] is False assert result["error"] == "Model version already exists in other library" assert scanners.other.calls == [101] assert execute_mock.await_count == 0 @pytest.mark.asyncio async def test_default_paths_use_per_sub_type_root( monkeypatch, scanners, metadata_provider, tmp_path ): """model.type VAE -> default_other_roots['vae'].""" captured = {} _capture_execute(monkeypatch, captured) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=99, use_default_paths=True ) assert result["success"] is True assert str(tmp_path / "vae") in str(captured["save_dir"]) @pytest.mark.asyncio async def test_default_paths_file_type_fallback_for_unmapped_model_type( monkeypatch, scanners, metadata_provider, tmp_path ): """model.type 'Other' maps to nothing; a 'Upscaler' file type decides.""" metadata_provider.payload = _other_payload( "Other", files=[ { "type": "Upscaler", "primary": True, "downloadUrl": "https://example.invalid/upscaler.safetensors", "name": "upscaler.safetensors", } ], ) captured = {} _capture_execute(monkeypatch, captured) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=99, use_default_paths=True ) assert result["success"] is True assert str(tmp_path / "upscale_models") in str(captured["save_dir"]) @pytest.mark.asyncio async def test_default_paths_explicit_file_pick_wins( monkeypatch, scanners, metadata_provider, tmp_path ): """An explicit pick of a bundled VAE component file routes to the vae root even though model.type maps to upscaler.""" metadata_provider.payload = _other_payload( "Upscaler", files=[ { "id": 1, "type": "Model", "primary": True, "downloadUrl": "https://example.invalid/model.safetensors", "name": "model.safetensors", }, { "id": 2, "type": "VAE", "downloadUrl": "https://example.invalid/bundled-vae.safetensors", "name": "bundled-vae.safetensors", }, ], ) captured = {} _capture_execute(monkeypatch, captured) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=99, use_default_paths=True, file_params={"id": 2, "type": "VAE"}, ) assert result["success"] is True assert str(tmp_path / "vae") in str(captured["save_dir"]) assert captured["download_urls"] == [ "https://example.invalid/bundled-vae.safetensors" ] @pytest.mark.asyncio async def test_default_paths_errors_when_sub_type_root_unconfigured( monkeypatch, scanners, metadata_provider, tmp_path ): """controlnet has no configured default root in the fixture settings.""" metadata_provider.payload = _other_payload("Controlnet") execute_mock = AsyncMock(return_value={"success": True}) monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=99, use_default_paths=True ) assert result["success"] is False assert "controlnet" in result["error"] assert execute_mock.await_count == 0 @pytest.mark.asyncio async def test_default_paths_errors_when_sub_type_undecidable( monkeypatch, scanners, metadata_provider, tmp_path ): """model.type 'Other' with only plain 'Model' files: never silently default to the vae folder — error and ask for an explicit folder.""" metadata_provider.payload = _other_payload("Other") execute_mock = AsyncMock(return_value={"success": True}) monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=99, use_default_paths=True ) assert result["success"] is False assert "sub-type" in result["error"] assert execute_mock.await_count == 0 @pytest.mark.asyncio async def test_civarchive_source_same_payload_shape( monkeypatch, scanners, metadata_provider, tmp_path ): """CivArchive downloads walk the same path with the same payload shape.""" metadata_provider.payload = _other_payload("TextEncoder") captured = {} _capture_execute(monkeypatch, captured) manager = DownloadManager() result = await manager.download_from_civitai( model_version_id=99, save_dir=str(tmp_path), source="civarchive" ) assert result["success"] is True assert captured["model_type"] == "other" def test_build_metadata_for_resume_uses_other_metadata(): manager = DownloadManager() metadata = manager._build_metadata_for_resume( model_type="other", version_info={"model": {"type": "VAE"}}, file_info={"name": "file.safetensors"}, save_path="/tmp/file.safetensors", ) assert metadata.metadata_class == "OtherModelMetadata" def test_other_extension_set_matches_checkpoint(): manager = DownloadManager() extensions = manager._get_supported_extensions_for_type("other") assert extensions == manager._get_supported_extensions_for_type("checkpoint") assert ".gguf" in extensions assert ".safetensors" in extensions @pytest.mark.asyncio async def test_sync_downloaded_version_uses_other_scanner(monkeypatch, scanners): """Update tracking for a downloaded other-model version consults the other scanner for local versions.""" class FakeUpdateService: def __init__(self): self.calls = [] async def update_in_library_versions( self, model_type, model_id, version_ids, version_info=None ): self.calls.append((model_type, model_id, version_ids)) update_service = FakeUpdateService() monkeypatch.setattr( ServiceRegistry, "get_model_update_service", AsyncMock(return_value=update_service), ) manager = DownloadManager() await manager._sync_downloaded_version( "other", 7, {"id": 42, "model": {"id": 7}} ) assert update_service.calls == [("other", 7, [42])]