ComfyUI-Lora-Manager/tests/utils/test_metadata_manager.py

from __future__ import annotations

import json

import pytest

from py.utils.metadata_manager import MetadataManager
from py.utils.models import LoraMetadata


@pytest.mark.asyncio
async def test_save_metadata_fills_missing_file_facts(tmp_path) -> None:
    """A payload missing file_name/size/modified is healed on write."""
    model_path = tmp_path / "MyModel.safetensors"
    model_path.write_bytes(b"fake model data")
    payload = {"file_path": str(model_path), "model_name": "My Model"}

    result = await MetadataManager.save_metadata(str(model_path), payload)
    assert result is True

    metadata_path = tmp_path / "MyModel.metadata.json"
    saved = json.loads(metadata_path.read_text(encoding="utf-8"))
    stat_result = model_path.stat()
    assert saved["file_name"] == "MyModel"
    assert saved["size"] == stat_result.st_size
    assert saved["modified"] == stat_result.st_mtime


@pytest.mark.asyncio
async def test_save_metadata_keeps_existing_file_facts(tmp_path) -> None:
    """Existing file facts are never overwritten on write."""
    model_path = tmp_path / "MyModel.safetensors"
    model_path.write_bytes(b"fake model data")
    payload = {
        "file_path": str(model_path),
        "file_name": "CustomName",
        "size": 123,
        "modified": 456.0,
    }

    result = await MetadataManager.save_metadata(str(model_path), payload)
    assert result is True

    metadata_path = tmp_path / "MyModel.metadata.json"
    saved = json.loads(metadata_path.read_text(encoding="utf-8"))
    assert saved["file_name"] == "CustomName"
    assert saved["size"] == 123
    assert saved["modified"] == 456.0


@pytest.mark.asyncio
async def test_load_metadata_payload_restores_file_facts_when_sidecar_missing(
    tmp_path,
) -> None:
    """Missing sidecar: the payload is rebuilt with local file facts."""
    model_path = tmp_path / "MyModel.safetensors"
    model_path.write_bytes(b"fake model data")

    payload = await MetadataManager.load_metadata_payload(str(model_path))

    stat_result = model_path.stat()
    assert payload["file_path"] == str(model_path)
    assert payload["file_name"] == "MyModel"
    assert payload["size"] == stat_result.st_size
    assert payload["modified"] == stat_result.st_mtime


@pytest.mark.asyncio
async def test_hydrate_model_data_restores_required_fields_when_sidecar_missing(
    tmp_path,
) -> None:
    """Self-heal repro: sidecar deleted, cache entry hydrated, base fields kept."""
    model_path = tmp_path / "MyModel.safetensors"
    model_path.write_bytes(b"fake model data")
    model_data = {
        "file_path": str(model_path),
        "folder": "extra_loras",
        "file_name": "MyModel",
        "model_name": "My Model",
        "size": 10,
        "modified": 100.0,
        "sha256": "abc123",
        "base_model": "Illustrious",
        "preview_url": "",
        "civitai": {"id": 123},
    }

    await MetadataManager.hydrate_model_data(model_data)

    stat_result = model_path.stat()
    assert model_data["file_name"] == "MyModel"
    assert model_data["model_name"] == "My Model"
    assert model_data["size"] == stat_result.st_size
    # `modified` is the import timestamp per schema; the cache value wins.
    assert model_data["modified"] == 100.0
    assert model_data["sha256"] == "abc123"
    assert model_data["base_model"] == "Illustrious"
    assert model_data["folder"] == "extra_loras"
    assert model_data["civitai"] == {"id": 123}


@pytest.mark.asyncio
async def test_hydrate_model_data_disk_wins_when_sidecar_exists(tmp_path) -> None:
    """Existing sidecar stays authoritative; no cache key is resurrected."""
    model_path = tmp_path / "MyModel.safetensors"
    model_path.write_bytes(b"fake model data")
    stat_result = model_path.stat()
    sidecar = {
        "file_path": str(model_path),
        "file_name": "MyModel",
        "model_name": "Disk Name",
        "size": stat_result.st_size,
        "modified": stat_result.st_mtime,
        "sha256": "diskhash",
        "base_model": "SDXL",
        "preview_url": "",
    }
    (tmp_path / "MyModel.metadata.json").write_text(
        json.dumps(sidecar), encoding="utf-8"
    )
    model_data = {
        "file_path": str(model_path),
        "folder": "extra_loras",
        "model_name": "Cache Name",
        "sha256": "cachehash",
        "civitai": {"id": 999},
    }

    await MetadataManager.hydrate_model_data(model_data)

    assert model_data["model_name"] == "Disk Name"
    assert model_data["sha256"] == "diskhash"
    # civitai is present only as the dataclass default {}; the cache's value
    # ({"id": 999}) is not resurrected.
    assert model_data["civitai"] == {}
    assert model_data["folder"] == "extra_loras"


@pytest.mark.asyncio
async def test_hydrate_model_data_keeps_sha256_missing_for_caller_persist_fix(
    tmp_path,
) -> None:
    """Sidecar exists but lacks sha256: hydrate leaves it missing so the
    caller's self-heal persist block still has work to do."""
    model_path = tmp_path / "MyModel.safetensors"
    model_path.write_bytes(b"fake model data")
    stat_result = model_path.stat()
    sidecar = {
        "file_path": str(model_path),
        "file_name": "MyModel",
        "model_name": "My Model",
        "size": stat_result.st_size,
        "modified": stat_result.st_mtime,
        "base_model": "SDXL",
        "preview_url": "",
        # sha256 deliberately absent
    }
    (tmp_path / "MyModel.metadata.json").write_text(
        json.dumps(sidecar), encoding="utf-8"
    )
    model_data = {
        "file_path": str(model_path),
        "sha256": "cachehash",
        "model_name": "Cache Name",
    }

    await MetadataManager.hydrate_model_data(model_data)

    assert "sha256" not in model_data


@pytest.mark.asyncio
async def test_load_metadata_payload_tolerates_missing_model_file(tmp_path) -> None:
    """A nonexistent model file must not crash payload loading."""
    missing_path = tmp_path / "Ghost.safetensors"
    payload = await MetadataManager.load_metadata_payload(str(missing_path))
    assert payload["file_path"] == str(missing_path)
    assert "file_name" not in payload
    assert "size" not in payload
    assert "modified" not in payload


@pytest.mark.asyncio
async def test_self_healed_sidecar_is_parseable(tmp_path) -> None:
    """End-to-end repro: deleted sidecar + refresh recreates a parseable file."""
    model_path = tmp_path / "MyModel.safetensors"
    model_path.write_bytes(b"fake model data")
    model_data = {
        "file_path": str(model_path),
        "folder": "extra_loras",
        "file_name": "MyModel",
        "model_name": "My Model",
        "size": 10,
        "modified": 100.0,
        "sha256": "abc123",
        "base_model": "Illustrious",
        "preview_url": "",
        "civitai": {"id": 123},
    }

    # Simulate the self-heal flow: hydrate from (missing) sidecar, then persist.
    await MetadataManager.hydrate_model_data(model_data)
    data_to_save = model_data.copy()
    data_to_save.pop("folder", None)
    await MetadataManager.save_metadata(str(model_path), data_to_save)

    metadata, should_skip = await MetadataManager.load_metadata(
        str(model_path), LoraMetadata
    )
    assert should_skip is False
    assert metadata is not None
    assert metadata.file_name == "MyModel"
    assert metadata.model_name == "My Model"
    assert metadata.sha256 == "abc123"