ai-robot-core/ai-service/tests/test_template_engine.py

"""
Unit tests for TemplateEngine.
[AC-IDS-06, AC-IDS-11] Test template variable filling.
"""

import pytest

from app.services.flow.template_engine import TemplateEngine


class MockLLMClient:
    """Mock LLM client for testing."""

    def __init__(self, response: str = "先生"):
        self._response = response

    async def generate_text(self, prompt: str) -> str:
        return self._response

    async def generate(self, messages: list) -> "MockResponse":
        return MockResponse(self._response)


class MockResponse:
    """Mock LLM response."""
    def __init__(self, content: str):
        self.content = content


class TestTemplateEngine:
    """[AC-IDS-06, AC-IDS-11] Test cases for TemplateEngine."""

    @pytest.mark.asyncio
    async def test_fill_template_no_variables(self):
        """Test template without variables returns as-is."""
        engine = TemplateEngine(llm_client=None)

        result = await engine.fill_template(
            template="您好，请问有什么可以帮您？",
            context=None,
            history=None,
        )

        assert result == "您好，请问有什么可以帮您？"

    @pytest.mark.asyncio
    async def test_fill_template_from_context(self):
        """Test filling variables from context."""
        engine = TemplateEngine(llm_client=None)

        result = await engine.fill_template(
            template="您好{user_name}，请问有什么可以帮您？",
            context={"user_name": "张先生"},
            history=None,
        )

        assert result == "您好张先生，请问有什么可以帮您？"

    @pytest.mark.asyncio
    async def test_fill_template_from_inputs(self):
        """Test filling variables from context inputs."""
        engine = TemplateEngine(llm_client=None)

        result = await engine.fill_template(
            template="您好，您咨询的是{product}相关的问题吗？",
            context={
                "inputs": [
                    {"variable": "product", "input": "手机"},
                ]
            },
            history=None,
        )

        assert result == "您好，您咨询的是手机相关的问题吗？"

    @pytest.mark.asyncio
    async def test_fill_template_with_llm(self):
        """Test filling variables using LLM generation."""
        llm_client = MockLLMClient(response="先生")
        engine = TemplateEngine(llm_client=llm_client)

        result = await engine.fill_template(
            template="您好{greeting_style}，请问您{inquiry_style}？",
            context=None,
            history=[{"role": "user", "content": "你好"}],
        )

        assert "先生" in result

    @pytest.mark.asyncio
    async def test_fill_template_multiple_variables(self):
        """Test filling multiple variables."""
        engine = TemplateEngine(llm_client=None)

        result = await engine.fill_template(
            template="您好{name}，您订购的{product}已发货，预计{date}送达。",
            context={
                "name": "李女士",
                "product": "iPhone 15",
                "date": "明天",
            },
            history=None,
        )

        assert result == "您好李女士，您订购的iPhone 15已发货，预计明天送达。"

    @pytest.mark.asyncio
    async def test_fill_template_missing_variable(self):
        """Test handling missing variables with placeholder."""
        engine = TemplateEngine(llm_client=None)

        result = await engine.fill_template(
            template="您好{unknown_var}，请问有什么可以帮您？",
            context=None,
            history=None,
        )

        assert result == "您好[unknown_var]，请问有什么可以帮您？"

    def test_extract_variables(self):
        """Test extracting variable names from template."""
        engine = TemplateEngine(llm_client=None)

        variables = engine.extract_variables(
            "您好{name}，您订购的{product}预计{date}送达。"
        )

        assert variables == ["name", "product", "date"]

    def test_extract_variables_empty(self):
        """Test extracting from template without variables."""
        engine = TemplateEngine(llm_client=None)

        variables = engine.extract_variables("您好，请问有什么可以帮您？")

        assert variables == []

    def test_extract_variables_adjacent(self):
        """Test extracting adjacent variables."""
        engine = TemplateEngine(llm_client=None)

        variables = engine.extract_variables("{a}{b}{c}")

        assert variables == ["a", "b", "c"]

    @pytest.mark.asyncio
    async def test_fill_template_with_history_context(self):
        """Test that history is used for LLM prompt."""
        llm_client = MockLLMClient(response="贵姓")
        engine = TemplateEngine(llm_client=llm_client)

        result = await engine.fill_template(
            template="您好，请问您{inquiry_style}？",
            context=None,
            history=[
                {"role": "user", "content": "我想咨询一下"},
                {"role": "assistant", "content": "好的，请问您想咨询什么？"},
            ],
        )

        assert "贵姓" in result

    @pytest.mark.asyncio
    async def test_fill_template_exception_handling(self):
        """Test that exceptions are handled gracefully."""
        class FailingLLMClient:
            async def generate_text(self, prompt: str) -> str:
                raise RuntimeError("LLM service unavailable")

        engine = TemplateEngine(llm_client=FailingLLMClient())

        result = await engine.fill_template(
            template="您好{greeting}，请问有什么可以帮您？",
            context=None,
            history=None,
        )

        assert result == "您好[greeting]，请问有什么可以帮您？"