fix(metadata): hydrate cached records before saving

This commit is contained in:
pixelpaws
2025-10-10 08:30:51 +08:00
parent 76d3aa2b5b
commit ef03a2a917
3 changed files with 54 additions and 10 deletions

View File

@@ -176,6 +176,9 @@ class MetadataSyncService:
metadata_provider = await self._get_default_provider() metadata_provider = await self._get_default_provider()
civitai_metadata, error = await metadata_provider.get_model_by_hash(sha256) civitai_metadata, error = await metadata_provider.get_model_by_hash(sha256)
await self._metadata_manager.hydrate_model_data(model_data)
if not civitai_metadata: if not civitai_metadata:
if error == "Model not found": if error == "Model not found":
model_data["from_civitai"] = False model_data["from_civitai"] = False

View File

@@ -1,4 +1,5 @@
from types import SimpleNamespace from types import SimpleNamespace
from typing import Any, Dict
from unittest.mock import AsyncMock from unittest.mock import AsyncMock
import pytest import pytest
@@ -20,7 +21,10 @@ def build_service(
default_provider: SimpleNamespace | None = None, default_provider: SimpleNamespace | None = None,
provider_selector: AsyncMock | None = None, provider_selector: AsyncMock | None = None,
): ):
metadata_manager = SimpleNamespace(save_metadata=AsyncMock()) metadata_manager = SimpleNamespace(
save_metadata=AsyncMock(),
hydrate_model_data=AsyncMock(side_effect=lambda payload: payload),
)
preview_service = SimpleNamespace(ensure_preview_for_metadata=AsyncMock()) preview_service = SimpleNamespace(ensure_preview_for_metadata=AsyncMock())
settings = DummySettings(settings_values) settings = DummySettings(settings_values)
@@ -105,12 +109,24 @@ async def test_fetch_and_update_model_success_updates_cache(tmp_path):
} }
helpers.default_provider.get_model_by_hash.return_value = (civitai_payload, None) helpers.default_provider.get_model_by_hash.return_value = (civitai_payload, None)
model_data = {"model_name": "Local", "folder": "root"} model_path = tmp_path / "model.safetensors"
async def hydrate(payload: Dict[str, Any]) -> Dict[str, Any]:
payload["hydrated"] = True
return payload
helpers.metadata_manager.hydrate_model_data.side_effect = hydrate
model_data = {
"model_name": "Local",
"folder": "root",
"file_path": str(model_path),
}
update_cache = AsyncMock(return_value=True) update_cache = AsyncMock(return_value=True)
ok, error = await helpers.service.fetch_and_update_model( ok, error = await helpers.service.fetch_and_update_model(
sha256="abc", sha256="abc",
file_path=str(tmp_path / "model.safetensors"), file_path=str(model_path),
model_data=model_data, model_data=model_data,
update_cache_func=update_cache, update_cache_func=update_cache,
) )
@@ -120,10 +136,15 @@ async def test_fetch_and_update_model_success_updates_cache(tmp_path):
assert model_data["civitai_deleted"] is False assert model_data["civitai_deleted"] is False
assert "civitai" in model_data assert "civitai" in model_data
metadata_path = str(tmp_path / "model.metadata.json") helpers.metadata_manager.hydrate_model_data.assert_awaited_once()
assert model_data["hydrated"] is True
metadata_path = str(model_path.with_suffix(".metadata.json"))
await_args = helpers.metadata_manager.save_metadata.await_args_list await_args = helpers.metadata_manager.save_metadata.await_args_list
assert await_args, "expected metadata to be persisted" assert await_args, "expected metadata to be persisted"
assert await_args[-1][0][0] == metadata_path last_call = await_args[-1]
assert last_call.args[0] == metadata_path
assert last_call.args[1]["hydrated"] is True
update_cache.assert_awaited_once() update_cache.assert_awaited_once()
@@ -132,14 +153,23 @@ async def test_fetch_and_update_model_handles_missing_remote_metadata(tmp_path):
helpers = build_service() helpers = build_service()
helpers.default_provider.get_model_by_hash.return_value = (None, "Model not found") helpers.default_provider.get_model_by_hash.return_value = (None, "Model not found")
model_path = tmp_path / "model.safetensors"
async def hydrate(payload: Dict[str, Any]) -> Dict[str, Any]:
payload["hydrated"] = True
return payload
helpers.metadata_manager.hydrate_model_data.side_effect = hydrate
model_data = { model_data = {
"model_name": "Local", "model_name": "Local",
"folder": "sub", "folder": "sub",
"file_path": str(model_path),
} }
ok, error = await helpers.service.fetch_and_update_model( ok, error = await helpers.service.fetch_and_update_model(
sha256="missing", sha256="missing",
file_path=str(tmp_path / "model.safetensors"), file_path=str(model_path),
model_data=model_data, model_data=model_data,
update_cache_func=AsyncMock(), update_cache_func=AsyncMock(),
) )
@@ -149,17 +179,24 @@ async def test_fetch_and_update_model_handles_missing_remote_metadata(tmp_path):
assert model_data["from_civitai"] is False assert model_data["from_civitai"] is False
assert model_data["civitai_deleted"] is True assert model_data["civitai_deleted"] is True
helpers.metadata_manager.hydrate_model_data.assert_awaited_once()
assert model_data["hydrated"] is True
helpers.metadata_manager.save_metadata.assert_awaited_once() helpers.metadata_manager.save_metadata.assert_awaited_once()
args, _ = helpers.metadata_manager.save_metadata.await_args call_args = helpers.metadata_manager.save_metadata.await_args
assert args[0].endswith("model.safetensors") assert call_args.args[0].endswith("model.safetensors")
assert "folder" not in args[1] assert "folder" not in call_args.args[1]
assert call_args.args[1]["hydrated"] is True
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_fetch_and_update_model_respects_deleted_without_archive(): async def test_fetch_and_update_model_respects_deleted_without_archive():
helpers = build_service(settings_values={"enable_metadata_archive_db": False}) helpers = build_service(settings_values={"enable_metadata_archive_db": False})
model_data = {"civitai_deleted": True} model_data = {
"civitai_deleted": True,
"file_path": "/tmp/model.safetensors",
}
update_cache = AsyncMock() update_cache = AsyncMock()
ok, error = await helpers.service.fetch_and_update_model( ok, error = await helpers.service.fetch_and_update_model(
@@ -172,6 +209,7 @@ async def test_fetch_and_update_model_respects_deleted_without_archive():
assert not ok assert not ok
assert "metadata archive DB is not enabled" in error assert "metadata archive DB is not enabled" in error
helpers.default_provider_factory.assert_not_awaited() helpers.default_provider_factory.assert_not_awaited()
helpers.metadata_manager.hydrate_model_data.assert_not_awaited()
update_cache.assert_not_awaited() update_cache.assert_not_awaited()

View File

@@ -30,6 +30,9 @@ class RecordingMetadataManager:
Path(metadata_path).write_text(json.dumps(metadata)) Path(metadata_path).write_text(json.dumps(metadata))
return True return True
async def hydrate_model_data(self, model_data: Dict[str, Any]) -> Dict[str, Any]:
return model_data
class RecordingPreviewService: class RecordingPreviewService:
def __init__(self) -> None: def __init__(self) -> None: