xiaoxia-saas/tests/unit/test_ai_client.py

"""AI Client (DoubaoClient) 单元测试"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from packages.shared.ai_client import DoubaoClient, get_doubao_client


@pytest.fixture
def mock_settings():
    """模拟配置"""
    with patch("packages.shared.ai_client.get_shared_settings") as mock:
        mock.return_value = MagicMock(
            doubao_api_key="test-api-key",
            doubao_model="doubao-pro-32k",
            doubao_base_url="https://ark.example.com/api/v3",
            doubao_timeout=30,
            doubao_max_retries=2,
        )
        yield mock


@pytest.fixture
def client_with_key(mock_settings):
    """有 API Key 的客户端"""
    return DoubaoClient()


@pytest.fixture
def client_without_key():
    """没有 API Key 的客户端"""
    with patch("packages.shared.ai_client.get_shared_settings") as mock:
        mock.return_value = MagicMock(
            doubao_api_key="",
            doubao_model="doubao-pro-32k",
            doubao_base_url="https://ark.example.com/api/v3",
            doubao_timeout=30,
            doubao_max_retries=2,
        )
        yield DoubaoClient()


class TestDoubaoClientInit:
    """初始化测试"""

    def test_init_with_api_key(self, mock_settings):
        """有 API Key 时初始化正常"""
        client = DoubaoClient()
        assert client.api_key == "test-api-key"
        assert client.model == "doubao-pro-32k"
        assert client.base_url == "https://ark.example.com/api/v3"
        assert client.timeout == 30
        assert client.max_retries == 2

    def test_base_url_strips_trailing_slash(self, mock_settings):
        """base_url 去掉末尾斜杠"""
        mock_settings.return_value.doubao_base_url = "https://ark.example.com/api/v3/"
        client = DoubaoClient()
        assert client.base_url == "https://ark.example.com/api/v3"


class TestIsAvailable:
    """is_available 属性测试"""

    def test_available_with_key(self, client_with_key):
        """有 API Key 时可用"""
        assert client_with_key.is_available is True

    def test_unavailable_without_key(self, client_without_key):
        """无 API Key 时不可用"""
        assert client_without_key.is_available is False


class TestChatCompletion:
    """chat_completion 方法测试"""

    def test_success_returns_content(self, client_with_key):
        """成功调用返回内容"""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"choices": [{"message": {"content": "  你好，我是豆包  "}}]}
        mock_response.raise_for_status = MagicMock()

        with patch("packages.shared.ai_client.httpx.post", return_value=mock_response) as mock_post:
            result = client_with_key.chat_completion(messages=[{"role": "user", "content": "你好"}])

            assert result == "你好，我是豆包"
            mock_post.assert_called_once()
            # 验证 URL
            call_args = mock_post.call_args
            assert call_args[0][0].endswith("/chat/completions")
            # 验证 header 包含 Authorization
            assert "Authorization" in call_args[1]["headers"]
            assert "Bearer test-api-key" in call_args[1]["headers"]["Authorization"]

    def test_unavailable_returns_none(self, client_without_key):
        """不可用时返回 None"""
        with patch("packages.shared.ai_client.httpx.post") as mock_post:
            result = client_without_key.chat_completion(messages=[{"role": "user", "content": "hi"}])
            assert result is None
            mock_post.assert_not_called()

    def test_with_temperature_and_max_tokens(self, client_with_key):
        """自定义 temperature 和 max_tokens"""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"choices": [{"message": {"content": "hi"}}]}
        mock_response.raise_for_status = MagicMock()

        with patch("packages.shared.ai_client.httpx.post", return_value=mock_response) as mock_post:
            client_with_key.chat_completion(
                messages=[{"role": "user", "content": "hi"}],
                temperature=0.3,
                max_tokens=512,
            )

            payload = mock_post.call_args[1]["json"]
            assert payload["temperature"] == 0.3
            assert payload["max_tokens"] == 512

    def test_retry_on_failure(self, client_with_key):
        """失败时自动重试"""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"choices": [{"message": {"content": "success"}}]}
        mock_response.raise_for_status = MagicMock()

        call_count = 0

        def side_effect(*args, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count < 3:  # 前两次失败，第三次成功
                raise Exception("temporary error")
            return mock_response

        with patch("packages.shared.ai_client.httpx.post", side_effect=side_effect):
            with patch("packages.shared.ai_client.time.sleep"):  # 跳过 sleep
                result = client_with_key.chat_completion(messages=[{"role": "user", "content": "hi"}])
                assert result == "success"
                assert call_count == 3  # 初始 1 次 + 2 次重试

    def test_all_retries_fail_returns_none(self, client_with_key):
        """所有重试都失败返回 None"""
        with patch("packages.shared.ai_client.httpx.post", side_effect=Exception("API down")):
            with patch("packages.shared.ai_client.time.sleep"):
                result = client_with_key.chat_completion(messages=[{"role": "user", "content": "hi"}])
                assert result is None

    def test_empty_choices_returns_none(self, client_with_key):
        """空 choices 返回 None 或抛异常"""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"choices": []}
        mock_response.raise_for_status = MagicMock()

        with patch("packages.shared.ai_client.httpx.post", return_value=mock_response):
            with patch("packages.shared.ai_client.time.sleep"):
                # 会因 IndexError 进入异常分支，最终返回 None
                result = client_with_key.chat_completion(messages=[{"role": "user", "content": "hi"}])
                assert result is None

    def test_messages_in_payload(self, client_with_key):
        """messages 正确传递到 payload"""
        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_response.json.return_value = {"choices": [{"message": {"content": "ok"}}]}
        mock_response.raise_for_status = MagicMock()

        messages = [
            {"role": "system", "content": "你是助手"},
            {"role": "user", "content": "你好"},
        ]

        with patch("packages.shared.ai_client.httpx.post", return_value=mock_response) as mock_post:
            client_with_key.chat_completion(messages=messages)

            payload = mock_post.call_args[1]["json"]
            assert payload["messages"] == messages
            assert payload["model"] == "doubao-pro-32k"


class TestGetDoubaoClient:
    """单例函数测试"""

    def test_returns_same_instance(self):
        """两次调用返回同一实例"""
        client1 = get_doubao_client()
        client2 = get_doubao_client()
        assert client1 is client2

    def test_returns_doubao_client_instance(self):
        """返回 DoubaoClient 实例"""
        client = get_doubao_client()
        assert isinstance(client, DoubaoClient)