feat(metadata): implement model data hydration and enhance metadata handling across services, fixes #547

tests/utils/test_example_images_metadata.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import json
+import os
+from pathlib import Path
 from types import SimpleNamespace
 from typing import Any, Dict, List, Tuple
 
@@ -30,7 +32,23 @@ def patch_metadata_manager(monkeypatch: pytest.MonkeyPatch):
         saved.append((path, metadata.copy()))
         return True
 
+    class SimpleMetadata:
+        def __init__(self, payload: Dict[str, Any]) -> None:
+            self._payload = payload
+            self._unknown_fields: Dict[str, Any] = {}
+
+        def to_dict(self) -> Dict[str, Any]:
+            return self._payload.copy()
+
+    async def fake_load(path: str, *_args: Any, **_kwargs: Any):
+        metadata_path = path if path.endswith(".metadata.json") else f"{os.path.splitext(path)[0]}.metadata.json"
+        if os.path.exists(metadata_path):
+            data = json.loads(Path(metadata_path).read_text(encoding="utf-8"))
+            return SimpleMetadata(data), False
+        return None, False
+
     monkeypatch.setattr(metadata_module.MetadataManager, "save_metadata", staticmethod(fake_save))
+    monkeypatch.setattr(metadata_module.MetadataManager, "load_metadata", staticmethod(fake_load))
     return saved
 
 
@@ -64,10 +82,80 @@ async def test_update_metadata_after_import_enriches_entries(monkeypatch: pytest
     assert custom[0]["hasMeta"] is True
     assert custom[0]["type"] == "image"
 
-    assert patch_metadata_manager[0][0] == str(model_file)
+    assert Path(patch_metadata_manager[0][0]) == model_file
     assert scanner.updates
 
 
+@pytest.mark.asyncio
+async def test_update_metadata_after_import_preserves_existing_metadata(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path,
+    patch_metadata_manager,
+):
+    model_hash = "b" * 64
+    model_file = tmp_path / "preserve.safetensors"
+    model_file.write_text("content", encoding="utf-8")
+    metadata_path = tmp_path / "preserve.metadata.json"
+
+    existing_payload = {
+        "model_name": "Example",
+        "file_path": str(model_file),
+        "civitai": {
+            "id": 42,
+            "modelId": 88,
+            "name": "Example",
+            "trainedWords": ["foo"],
+            "images": [{"url": "https://example.com/default.png", "type": "image"}],
+            "customImages": [
+                {"id": "existing-id", "type": "image", "url": "", "nsfwLevel": 0}
+            ],
+        },
+        "extraField": "keep-me",
+    }
+    metadata_path.write_text(json.dumps(existing_payload), encoding="utf-8")
+
+    model_data = {
+        "sha256": model_hash,
+        "model_name": "Example",
+        "file_path": str(model_file),
+        "civitai": {
+            "id": 42,
+            "modelId": 88,
+            "name": "Example",
+            "trainedWords": ["foo"],
+            "customImages": [],
+        },
+    }
+    scanner = StubScanner([model_data])
+
+    image_path = tmp_path / "new.png"
+    image_path.write_bytes(b"fakepng")
+
+    monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))
+    monkeypatch.setattr(metadata_module.MetadataUpdater, "_parse_image_metadata", staticmethod(lambda payload: None))
+
+    regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
+        model_hash,
+        model_data,
+        scanner,
+        [(str(image_path), "new-id")],
+    )
+
+    assert regular == existing_payload["civitai"]["images"]
+    assert any(entry["id"] == "new-id" for entry in custom)
+
+    saved_path, saved_payload = patch_metadata_manager[-1]
+    assert Path(saved_path) == model_file
+    assert saved_payload["extraField"] == "keep-me"
+    assert saved_payload["civitai"]["images"] == existing_payload["civitai"]["images"]
+    assert saved_payload["civitai"]["trainedWords"] == ["foo"]
+    assert {entry["id"] for entry in saved_payload["civitai"]["customImages"]} == {"existing-id", "new-id"}
+
+    assert scanner.updates
+    updated_metadata = scanner.updates[-1][2]
+    assert updated_metadata["civitai"]["images"] == existing_payload["civitai"]["images"]
+    assert {entry["id"] for entry in updated_metadata["civitai"]["customImages"]} == {"existing-id", "new-id"}
+
 async def test_refresh_model_metadata_records_failures(monkeypatch: pytest.MonkeyPatch, tmp_path):
     model_hash = "b" * 64
     model_file = tmp_path / "model.safetensors"
@@ -79,6 +167,16 @@ async def test_refresh_model_metadata_records_failures(monkeypatch: pytest.Monke
         async def fetch_and_update_model(self, **_kwargs):
             return True, None
 
+    async def fake_hydrate(model_data: Dict[str, Any]) -> Dict[str, Any]:
+        model_data["hydrated"] = True
+        return model_data
+
+    monkeypatch.setattr(
+        metadata_module.MetadataManager,
+        "hydrate_model_data",
+        staticmethod(fake_hydrate),
+    )
+
     monkeypatch.setattr(metadata_module, "_metadata_sync_service", StubMetadataSync())
 
     result = await metadata_module.MetadataUpdater.refresh_model_metadata(
@@ -89,6 +187,7 @@ async def test_refresh_model_metadata_records_failures(monkeypatch: pytest.Monke
         {"refreshed_models": set(), "errors": [], "last_error": None},
     )
     assert result is True
+    assert cache_item["hydrated"] is True
 
 
 async def test_update_metadata_from_local_examples_generates_entries(monkeypatch: pytest.MonkeyPatch, tmp_path):
@@ -112,4 +211,4 @@ async def test_update_metadata_from_local_examples_generates_entries(monkeypatch
         str(model_dir),
     )
     assert success is True
-    assert model_data["civitai"]["images"]
+    assert model_data["civitai"]["images"]