ComfyUI-Lora-Manager/tests/services/test_agent_enrichment_source.py

"""Tests for source-aware AI enrichment orchestration.

Covers the fast-fail gate (:meth:`AgentService._enrichment_skip_reason`) and
the prompt-context builder for non-Hugging Face model sources.
"""

from __future__ import annotations

from unittest import mock

import pytest

from py.services.agent.agent_service import AgentService


class TestEnrichmentSkipReason:
    def test_skips_when_no_source_linked(self):
        reason = AgentService._enrichment_skip_reason({})
        assert "source_url" in reason

    def test_allows_huggingface(self):
        assert (
            AgentService._enrichment_skip_reason(
                {"hf_url": "https://huggingface.co/user/repo"}
            )
            == ""
        )

    def test_allows_modelscope(self):
        assert (
            AgentService._enrichment_skip_reason(
                {
                    "source_platform": "modelscope",
                    "source_url": "https://modelscope.cn/models/user/repo",
                }
            )
            == ""
        )

    def test_skips_tensorart_with_reason(self):
        reason = AgentService._enrichment_skip_reason(
            {
                "source_platform": "tensorart",
                "source_url": "https://tensor.art/models/827823520299086029",
            }
        )
        assert "TensorArt" in reason
        assert "not available" in reason

    def test_skips_unknown_platform(self):
        reason = AgentService._enrichment_skip_reason(
            {"source_platform": "somewhere", "source_url": "https://somewhere.example/m/1"}
        )
        assert "somewhere" in reason


class TestBuildPromptContext:
    @pytest.mark.asyncio
    async def test_modelscope_card_populates_source_variables(self):
        service = AgentService()
        readme = "---\nbase_model: krea/Krea-2-Turbo\n---\n# krea\n"

        with (
            mock.patch(
                "py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
                new=mock.AsyncMock(return_value=readme),
            ) as mock_fetch,
            mock.patch(
                "py.metadata_ops.list_base_models",
                new=mock.AsyncMock(return_value=["Krea 2 Turbo"]),
            ),
            mock.patch(
                "py.metadata_ops.identify_model_type",
                new=mock.AsyncMock(return_value="lora"),
            ),
            mock.patch(
                "py.services.settings_manager.SettingsManager.get_priority_tag_config",
                return_value={"lora": "style, subject"},
            ),
        ):
            context = await service._build_prompt_context(
                skill_name="enrich_hf_metadata",
                model_path="/models/loras/krea.safetensors",
                metadata={
                    "source_platform": "modelscope",
                    "source_url": "https://modelscope.cn/models/jj3550945163/Krea-2-LORA",
                    "file_name": "krea",
                },
                registry=mock.Mock(),
                llm=mock.Mock(),
            )

        mock_fetch.assert_awaited_once_with("jj3550945163/Krea-2-LORA")
        assert context["source_platform"] == "modelscope"
        assert context["source_id"] == "jj3550945163/Krea-2-LORA"
        assert context["source_label"] == "ModelScope"
        assert (
            context["asset_base_url"]
            == "https://modelscope.cn/models/jj3550945163/Krea-2-LORA/resolve/master"
        )
        assert readme in context["readme_content_full"]
        # Hugging Face aliases stay empty for a non-HF source.
        assert context["hf_url"] == ""
        assert context["repo"] == "jj3550945163/Krea-2-LORA"

    @pytest.mark.asyncio
    async def test_huggingface_keeps_legacy_aliases(self):
        service = AgentService()
        readme = "# card\n"

        with (
            mock.patch(
                "py.services.model_sources.huggingface.HuggingFaceSource.fetch_model_card",
                new=mock.AsyncMock(return_value=readme),
            ) as mock_fetch,
            mock.patch(
                "py.metadata_ops.list_base_models",
                new=mock.AsyncMock(return_value=[]),
            ),
            mock.patch(
                "py.metadata_ops.identify_model_type",
                new=mock.AsyncMock(return_value="lora"),
            ),
            mock.patch(
                "py.services.settings_manager.SettingsManager.get_priority_tag_config",
                return_value={},
            ),
        ):
            context = await service._build_prompt_context(
                skill_name="enrich_hf_metadata",
                model_path="/models/loras/thing.safetensors",
                metadata={"hf_url": "https://huggingface.co/user/repo"},
                registry=mock.Mock(),
                llm=mock.Mock(),
            )

        mock_fetch.assert_awaited_once_with("user/repo")
        assert context["source_platform"] == "huggingface"
        assert context["hf_url"] == "https://huggingface.co/user/repo"
        assert context["repo"] == "user/repo"

    @pytest.mark.asyncio
    async def test_tensorart_never_fetches_a_card(self):
        service = AgentService()

        with (
            mock.patch(
                "py.services.model_sources.huggingface.HuggingFaceSource.fetch_model_card",
                new=mock.AsyncMock(),
            ) as hf_fetch,
            mock.patch(
                "py.services.model_sources.modelscope.ModelScopeSource.fetch_model_card",
                new=mock.AsyncMock(),
            ) as ms_fetch,
            mock.patch(
                "py.metadata_ops.list_base_models",
                new=mock.AsyncMock(return_value=[]),
            ),
            mock.patch(
                "py.metadata_ops.identify_model_type",
                new=mock.AsyncMock(return_value="lora"),
            ),
            mock.patch(
                "py.services.settings_manager.SettingsManager.get_priority_tag_config",
                return_value={},
            ),
        ):
            context = await service._build_prompt_context(
                skill_name="enrich_hf_metadata",
                model_path="/models/loras/thing.safetensors",
                metadata={
                    "source_platform": "tensorart",
                    "source_url": "https://tensor.art/models/827823520299086029",
                },
                registry=mock.Mock(),
                llm=mock.Mock(),
            )

        hf_fetch.assert_not_awaited()
        ms_fetch.assert_not_awaited()
        assert context["readme_content"] == ""
        assert context["source_platform"] == "tensorart"