willmiao-ComfyUI-Lora-Manager/tests/services/test_workflow_base_model.py

"""Tests for py.services.workflow_base_model: base model resolution."""

from py.services.model_scanner import ModelScanner
from py.services.workflow_base_model import resolve_workflow_base_model


class _FakeModelScanner:
    """Stand-in for the lora/checkpoint scanners backed by fixed items."""

    def __init__(self, items):
        self._items = items

    async def find_models_by_name(self, name, *, base_model=None):
        return ModelScanner.find_matching_models(
            self._items, name, base_model=base_model
        )


class _FakeProvider:
    def __init__(self, infos):
        self._infos = infos
        self.calls = []

    async def get_model_version_info(self, version_id):
        self.calls.append(version_id)
        return self._infos.get(version_id), None


def _checkpoint(file_name, base_model, folder=""):
    return {"file_name": file_name, "folder": folder, "base_model": base_model}


async def test_resolves_from_checkpoint_loader_refs():
    checkpoint_scanner = _FakeModelScanner(
        [_checkpoint("sd_xl_base_1.0.safetensors", "SDXL 1.0")]
    )
    result = await resolve_workflow_base_model(
        primary_model_refs=["sd_xl_base_1.0.safetensors"],
        model_refs=["sd_xl_base_1.0.safetensors"],
        checkpoint_scanner=checkpoint_scanner,
        lora_scanner=_FakeModelScanner([]),
    )
    assert result == "SDXL 1.0"


async def test_checkpoint_refs_accept_spaced_and_foldered_names():
    checkpoint_scanner = _FakeModelScanner(
        [_checkpoint("My Model v2.safetensors", "SDXL 1.0", folder="SDXL")]
    )
    result = await resolve_workflow_base_model(
        primary_model_refs=["SDXL/My Model v2.safetensors"],
        model_refs=[],
        checkpoint_scanner=checkpoint_scanner,
        lora_scanner=_FakeModelScanner([]),
    )
    assert result == "SDXL 1.0"


async def test_checkpoint_majority_vote_wins():
    checkpoint_scanner = _FakeModelScanner(
        [
            _checkpoint("base.safetensors", "SDXL 1.0"),
            _checkpoint("refiner.safetensors", "SDXL 1.0"),
            _checkpoint("other.safetensors", "SD 1.5"),
        ]
    )
    result = await resolve_workflow_base_model(
        primary_model_refs=[
            "base.safetensors",
            "refiner.safetensors",
            "other.safetensors",
        ],
        model_refs=[],
        checkpoint_scanner=checkpoint_scanner,
        lora_scanner=_FakeModelScanner([]),
    )
    assert result == "SDXL 1.0"


async def test_unknown_checkpoint_base_model_never_votes():
    checkpoint_scanner = _FakeModelScanner(
        [_checkpoint("base.safetensors", "Unknown")]
    )
    lora_scanner = _FakeModelScanner(
        [_checkpoint("style.safetensors", "Pony")]
    )
    result = await resolve_workflow_base_model(
        primary_model_refs=["base.safetensors"],
        model_refs=["style.safetensors"],
        checkpoint_scanner=checkpoint_scanner,
        lora_scanner=lora_scanner,
    )
    assert result == "Pony"


async def test_lora_majority_vote_is_the_fallback():
    checkpoint_scanner = _FakeModelScanner([])
    lora_scanner = _FakeModelScanner(
        [
            _checkpoint("a.safetensors", "Illustrious"),
            _checkpoint("b.safetensors", "Illustrious"),
            _checkpoint("c.safetensors", "SDXL 1.0"),
        ]
    )
    result = await resolve_workflow_base_model(
        primary_model_refs=["not-in-library.safetensors"],
        model_refs=["a.safetensors", "b.safetensors", "c.safetensors"],
        checkpoint_scanner=checkpoint_scanner,
        lora_scanner=lora_scanner,
    )
    assert result == "Illustrious"


async def test_civitai_urn_fallback_requires_opt_in():
    provider = _FakeProvider({"67890": {"baseModel": "SDXL 1.0"}})
    kwargs = dict(
        primary_model_refs=[],
        model_refs=[],
        civitai_refs=["civitai:12345@67890"],
        checkpoint_scanner=_FakeModelScanner([]),
        lora_scanner=_FakeModelScanner([]),
        metadata_provider=provider,
    )

    assert await resolve_workflow_base_model(**kwargs) is None
    assert provider.calls == []

    assert await resolve_workflow_base_model(**kwargs, include_remote=True) == "SDXL 1.0"
    assert provider.calls == ["67890"]


async def test_civitai_urn_fallback_ignores_unknown_answers():
    provider = _FakeProvider(
        {"111": {"baseModel": ""}, "222": {"baseModel": "Flux.1 D"}}
    )
    result = await resolve_workflow_base_model(
        primary_model_refs=[],
        model_refs=[],
        civitai_refs=["civitai:1@111", "civitai:2@222"],
        checkpoint_scanner=_FakeModelScanner([]),
        lora_scanner=_FakeModelScanner([]),
        include_remote=True,
        metadata_provider=provider,
    )
    assert result == "Flux.1 D"


async def test_checkpoint_name_hint_outranks_static_votes():
    checkpoint_scanner = _FakeModelScanner(
        [
            _checkpoint("real_primary.safetensors", "Illustrious"),
            _checkpoint("voted.safetensors", "SDXL 1.0"),
        ]
    )
    result = await resolve_workflow_base_model(
        primary_model_refs=["voted.safetensors"],
        model_refs=["voted.safetensors"],
        checkpoint_scanner=checkpoint_scanner,
        lora_scanner=_FakeModelScanner([]),
        checkpoint_name_hint="real_primary.safetensors",
    )
    assert result == "Illustrious"


async def test_checkpoint_name_hint_falls_back_to_votes_when_unresolved():
    checkpoint_scanner = _FakeModelScanner(
        [_checkpoint("voted.safetensors", "SDXL 1.0")]
    )
    result = await resolve_workflow_base_model(
        primary_model_refs=["voted.safetensors"],
        model_refs=["voted.safetensors"],
        checkpoint_scanner=checkpoint_scanner,
        lora_scanner=_FakeModelScanner([]),
        checkpoint_name_hint="deleted-model.safetensors",
    )
    assert result == "SDXL 1.0"


async def test_returns_none_when_nothing_resolves():
    result = await resolve_workflow_base_model(
        primary_model_refs=["missing.safetensors"],
        model_refs=["missing.safetensors"],
        checkpoint_scanner=_FakeModelScanner([]),
        lora_scanner=_FakeModelScanner([]),
    )
    assert result is None