"""LLM 封装测试。"""

from unittest.mock import MagicMock, patch

from promptgrad.llm import call_llm


class TestCallLLM:
    @patch("promptgrad.llm.completion")
    def test_call_returns_output_and_usage(self, mock_completion):
        mock_response = MagicMock()
        mock_response.choices = [MagicMock()]
        mock_response.choices[0].message.content = "Hello world"
        mock_response.usage.prompt_tokens = 10
        mock_response.usage.completion_tokens = 5
        mock_completion.return_value = mock_response

        result = call_llm(
            model="claude-sonnet-4-20250514",
            messages=[{"role": "user", "content": "Hi"}],
            temperature=0.3,
            max_tokens=100,
        )
        assert result["output"] == "Hello world"
        assert result["token_usage"]["input"] == 10
        assert result["token_usage"]["output"] == 5
        assert "latency_ms" in result

    @patch("promptgrad.llm.completion")
    def test_call_with_system_message(self, mock_completion):
        mock_response = MagicMock()
        mock_response.choices = [MagicMock()]
        mock_response.choices[0].message.content = "response"
        mock_response.usage.prompt_tokens = 20
        mock_response.usage.completion_tokens = 10
        mock_completion.return_value = mock_response

        call_llm(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "You are helpful."},
                {"role": "user", "content": "Hi"},
            ],
        )
        mock_completion.assert_called_once()
        call_args = mock_completion.call_args
        assert call_args.kwargs["model"] == "gpt-4o"
