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

from __future__ import annotations

import json
from types import SimpleNamespace
from typing import Any, Dict, List, Tuple

import pytest

from py.utils import example_images_metadata as metadata_module


class StubScanner:
    def __init__(self, cache_items: List[Dict[str, Any]]) -> None:
        self.cache = SimpleNamespace(raw_data=cache_items)
        self.updates: List[Tuple[str, str, Dict[str, Any]]] = []

    async def get_cached_data(self):
        return self.cache

    async def update_single_model_cache(self, old_path: str, new_path: str, metadata: Dict[str, Any]) -> bool:
        self.updates.append((old_path, new_path, metadata))
        return True


@pytest.fixture(autouse=True)
def patch_metadata_manager(monkeypatch: pytest.MonkeyPatch):
    saved: List[Tuple[str, Dict[str, Any]]] = []

    async def fake_save(path: str, metadata: Dict[str, Any]) -> bool:
        saved.append((path, metadata.copy()))
        return True

    monkeypatch.setattr(metadata_module.MetadataManager, "save_metadata", staticmethod(fake_save))
    return saved


async def test_update_metadata_after_import_enriches_entries(monkeypatch: pytest.MonkeyPatch, tmp_path, patch_metadata_manager):
    model_hash = "a" * 64
    model_file = tmp_path / "model.safetensors"
    model_file.write_text("content", encoding="utf-8")
    model_data = {
        "model_name": "Example",
        "file_path": str(model_file),
        "civitai": {},
    }
    scanner = StubScanner([model_data])

    image_path = tmp_path / "custom.png"
    image_path.write_bytes(b"fakepng")

    monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: "Prompt text Negative prompt: bad Steps: 20, Sampler: Euler"))
    monkeypatch.setattr(metadata_module.MetadataUpdater, "_parse_image_metadata", staticmethod(lambda payload: {"prompt": "Prompt text", "negativePrompt": "bad", "parameters": {"Steps": "20"}}))

    regular, custom = await metadata_module.MetadataUpdater.update_metadata_after_import(
        model_hash,
        model_data,
        scanner,
        [(str(image_path), "short-id")],
    )

    assert isinstance(custom, list)
    assert custom[0]["id"] == "short-id"
    assert custom[0]["meta"]["prompt"] == "Prompt text"
    assert custom[0]["hasMeta"] is True
    assert custom[0]["type"] == "image"

    assert patch_metadata_manager[0][0] == str(model_file)
    assert scanner.updates


async def test_refresh_model_metadata_records_failures(monkeypatch: pytest.MonkeyPatch, tmp_path):
    model_hash = "b" * 64
    model_file = tmp_path / "model.safetensors"
    model_file.write_text("content", encoding="utf-8")
    cache_item = {"sha256": model_hash, "file_path": str(model_file)}
    scanner = StubScanner([cache_item])

    class StubMetadataSync:
        async def fetch_and_update_model(self, **_kwargs):
            return True, None

    monkeypatch.setattr(metadata_module, "_metadata_sync_service", StubMetadataSync())

    result = await metadata_module.MetadataUpdater.refresh_model_metadata(
        model_hash,
        "Example",
        "lora",
        scanner,
        {"refreshed_models": set(), "errors": [], "last_error": None},
    )
    assert result is True


async def test_update_metadata_from_local_examples_generates_entries(monkeypatch: pytest.MonkeyPatch, tmp_path):
    model_hash = "c" * 64
    model_dir = tmp_path / model_hash
    model_dir.mkdir()
    (model_dir / "image.png").write_text("data", encoding="utf-8")
    model_data = {"model_name": "Local", "civitai": {}, "file_path": str(tmp_path / "model.safetensors")}

    async def fake_save(path, metadata):
        return True

    monkeypatch.setattr(metadata_module.MetadataManager, "save_metadata", staticmethod(fake_save))
    monkeypatch.setattr(metadata_module.ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None))

    success = await metadata_module.MetadataUpdater.update_metadata_from_local_examples(
        model_hash,
        model_data,
        "lora",
        StubScanner([model_data]),
        str(model_dir),
    )
    assert success is True
    assert model_data["civitai"]["images"]