feat(relink): accept CivitArchive URLs when linking models

2026-08-28 08:21:27 -03:00 · 2026-08-26 21:31:30 +08:00
parent 3025c64fea
commit 641a61f804
21 changed files with 756 additions and 121 deletions
@@ -560,6 +560,131 @@ async def test_relink_metadata_raises_when_version_missing():
            model_version_id=None,
        )

+
+@pytest.mark.asyncio
+async def test_relink_metadata_uses_named_civarchive_provider(tmp_path):
+    default_provider = SimpleNamespace(
+        get_model_by_hash=AsyncMock(),
+        get_model_version=AsyncMock(),
+    )
+    civarchive_provider = SimpleNamespace(
+        get_model_by_hash=AsyncMock(),
+        get_model_version=AsyncMock(
+            return_value={
+                "files": [
+                    {
+                        "primary": True,
+                        "type": "Model",
+                        "hashes": {"SHA256": "ABCDEF"},
+                    }
+                ],
+                "model": {"name": "Archived"},
+                "images": [],
+            }
+        ),
+    )
+
+    async def select_provider(name: str):
+        return civarchive_provider if name == "civarchive_api" else default_provider
+
+    provider_selector = AsyncMock(side_effect=select_provider)
+    helpers = build_service(
+        default_provider=default_provider,
+        provider_selector=provider_selector,
+    )
+
+    metadata = {"model_name": "Local", "sha256": "original"}
+    result = await helpers.service.relink_metadata(
+        file_path=str(tmp_path / "model.safetensors"),
+        metadata=metadata,
+        model_id=1,
+        model_version_id=2,
+        provider_name="civarchive_api",
+    )
+
+    assert result["model_name"] == "Archived"
+    assert result["sha256"] == "original"
+    provider_selector.assert_awaited_with("civarchive_api")
+    civarchive_provider.get_model_version.assert_awaited_once_with(1, 2)
+    helpers.default_provider_factory.assert_not_awaited()
+    helpers.metadata_manager.save_metadata.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_relink_metadata_raises_when_version_missing_with_civarchive():
+    default_provider = SimpleNamespace(
+        get_model_by_hash=AsyncMock(),
+        get_model_version=AsyncMock(),
+    )
+    civarchive_provider = SimpleNamespace(
+        get_model_by_hash=AsyncMock(),
+        get_model_version=AsyncMock(return_value=None),
+    )
+
+    async def select_provider(name: str):
+        return civarchive_provider if name == "civarchive_api" else default_provider
+
+    provider_selector = AsyncMock(side_effect=select_provider)
+    helpers = build_service(
+        default_provider=default_provider,
+        provider_selector=provider_selector,
+    )
+
+    with pytest.raises(ValueError, match="CivitArchive"):
+        await helpers.service.relink_metadata(
+            file_path="/tmp/model.safetensors",
+            metadata={},
+            model_id=9,
+            model_version_id=None,
+            provider_name="civarchive_api",
+        )
+
+
+@pytest.mark.asyncio
+async def test_relink_metadata_raises_friendly_error_when_provider_unavailable():
+    provider_selector = AsyncMock(
+        side_effect=ValueError("Provider 'civarchive_api' is not registered")
+    )
+    helpers = build_service(provider_selector=provider_selector)
+
+    with pytest.raises(ValueError, match="CivitArchive is not available or not enabled"):
+        await helpers.service.relink_metadata(
+            file_path="/tmp/model.safetensors",
+            metadata={},
+            model_id=9,
+            model_version_id=None,
+            provider_name="civarchive_api",
+        )
+
+
+@pytest.mark.asyncio
+async def test_relink_metadata_default_call_uses_default_provider_factory(tmp_path):
+    helpers = build_service()
+    helpers.default_provider.get_model_version.return_value = {
+        "files": [
+            {
+                "primary": True,
+                "type": "Model",
+                "hashes": {"SHA256": "ABCDEF"},
+            }
+        ],
+        "model": {"name": "Remote"},
+        "images": [],
+    }
+
+    result = await helpers.service.relink_metadata(
+        file_path=str(tmp_path / "model.safetensors"),
+        metadata={"model_name": "Local", "sha256": "original"},
+        model_id=1,
+        model_version_id=None,
+    )
+
+    assert result["model_name"] == "Remote"
+    assert result["sha256"] == "original"
+    helpers.default_provider_factory.assert_awaited_once()
+    helpers.provider_selector.assert_not_awaited()
+    helpers.metadata_manager.save_metadata.assert_awaited_once()
+
@pytest.mark.asyncio
 async def test_fetch_and_update_model_persists_db_checked_when_sqlite_fails(tmp_path):
    """