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=45,
            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=45,
            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 == 45
        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)


class TestChatCompletionStream:
    """chat_completion_stream 方法测试"""

    def _fake_sse(self, chunks: list[str]):
        """构造一个伪装的 httpx stream 响应，按 chunks 逐行返回 SSE。"""
        import json as _json

        lines = []
        for piece in chunks:
            evt = {"choices": [{"delta": {"content": piece}, "finish_reason": None}]}
            lines.append("data: " + _json.dumps(evt, ensure_ascii=False))
        lines.append("data: " + _json.dumps({"choices": [{"delta": {}, "finish_reason": "stop"}]}))
        lines.append("data: [DONE]")

        m = MagicMock()
        m.status_code = 200
        m.__enter__ = MagicMock(return_value=m)
        m.__exit__ = MagicMock(return_value=False)
        m.raise_for_status = MagicMock()
        m.iter_lines = MagicMock(return_value=iter(lines))
        return m

    def test_stream_yields_delta_content(self, client_with_key):
        """流式调用逐段 yield delta.content"""
        import json as _json

        fake = self._fake_sse(["你好", "，", "世界"])
        with patch("packages.shared.ai_client.httpx.stream", return_value=fake) as mock_stream:
            out = list(
                client_with_key.chat_completion_stream(
                    messages=[{"role": "user", "content": "hi"}], max_tokens=100, timeout=30
                )
            )
        assert out == ["你好", "，", "世界"]
        call_kwargs = mock_stream.call_args[1]
        assert call_kwargs["json"]["stream"] is True
        assert call_kwargs["json"]["max_tokens"] == 100

    def test_stream_unavailable_returns_empty(self, client_without_key):
        """不可用时返回空生成器（不调用 httpx.stream）"""
        with patch("packages.shared.ai_client.httpx.stream") as mock_stream:
            out = list(client_without_key.chat_completion_stream(messages=[{"role": "user", "content": "hi"}]))
        assert out == []
        mock_stream.assert_not_called()

    def test_stream_handles_done_marker(self, client_with_key):
        """遇到 [DONE] 正确终止，不把它当内容"""
        import json as _json

        lines = [
            "data: " + _json.dumps({"choices": [{"delta": {"content": "A"}}]}),
            "data: [DONE]",
            "data: " + _json.dumps({"choices": [{"delta": {"content": "NEVER"}}]}),  # 应被忽略
        ]
        m = MagicMock()
        m.status_code = 200
        m.__enter__ = MagicMock(return_value=m)
        m.__exit__ = MagicMock(return_value=False)
        m.raise_for_status = MagicMock()
        m.iter_lines = MagicMock(return_value=iter(lines))
        with patch("packages.shared.ai_client.httpx.stream", return_value=m):
            out = list(
                client_with_key.chat_completion_stream(
                    messages=[{"role": "user", "content": "hi"}], max_tokens=50, timeout=30
                )
            )
        assert out == ["A"]

    def test_stream_skips_empty_delta(self, client_with_key):
        """空 delta（role 等 metadata）不应产出内容"""
        import json as _json

        lines = [
            "data: " + _json.dumps({"choices": [{"delta": {"role": "assistant"}}]}),
            "data: " + _json.dumps({"choices": [{"delta": {"content": "hi"}}]}),
            "data: " + _json.dumps({"choices": [{"delta": {}}]}),
            "data: [DONE]",
        ]
        m = MagicMock()
        m.status_code = 200
        m.__enter__ = MagicMock(return_value=m)
        m.__exit__ = MagicMock(return_value=False)
        m.raise_for_status = MagicMock()
        m.iter_lines = MagicMock(return_value=iter(lines))
        with patch("packages.shared.ai_client.httpx.stream", return_value=m):
            out = list(
                client_with_key.chat_completion_stream(
                    messages=[{"role": "user", "content": "x"}], max_tokens=50, timeout=30
                )
            )
        assert out == ["hi"]

    def test_stream_invalid_json_lines_skipped(self, client_with_key):
        """SSE 行里脏数据/非 JSON 不应中断流"""
        import json as _json

        lines = [
            "data: " + _json.dumps({"choices": [{"delta": {"content": "ok"}}]}),
            "data: not-a-json",
            ":comment line",
            "",
            "event: ping",
            "data: " + _json.dumps({"choices": [{"delta": {"content": "2"}}]}),
            "data: [DONE]",
        ]
        m = MagicMock()
        m.status_code = 200
        m.__enter__ = MagicMock(return_value=m)
        m.__exit__ = MagicMock(return_value=False)
        m.raise_for_status = MagicMock()
        m.iter_lines = MagicMock(return_value=iter(lines))
        with patch("packages.shared.ai_client.httpx.stream", return_value=m):
            out = list(
                client_with_key.chat_completion_stream(
                    messages=[{"role": "user", "content": "x"}], max_tokens=50, timeout=30
                )
            )
        assert out == ["ok", "2"]

    def test_stream_network_error_yields_empty(self, client_with_key):
        """网络错误（重试耗尽）yield 空，不抛异常给调用方"""
        import httpx as _httpx

        # max_retries=2 → 3 次总尝试
        client = client_with_key
        client.max_retries = 1  # 只重试 1 次，缩短测试
        with patch("packages.shared.ai_client.httpx.stream", side_effect=_httpx.ConnectError("boom")):
            out = list(
                client.chat_completion_stream(messages=[{"role": "user", "content": "x"}], max_tokens=50, timeout=5)
            )
        assert out == []

    def test_stream_records_finish_reason(self, client_with_key):
        """流式结束后 last_finish_reason 被正确记录"""
        import json as _json

        lines = [
            "data: " + _json.dumps({"choices": [{"delta": {"content": "x"}, "finish_reason": None}]}),
            "data: " + _json.dumps({"choices": [{"delta": {}, "finish_reason": "stop"}]}),
            "data: [DONE]",
        ]
        m = MagicMock()
        m.status_code = 200
        m.__enter__ = MagicMock(return_value=m)
        m.__exit__ = MagicMock(return_value=False)
        m.raise_for_status = MagicMock()
        m.iter_lines = MagicMock(return_value=iter(lines))
        with patch("packages.shared.ai_client.httpx.stream", return_value=m):
            list(
                client_with_key.chat_completion_stream(
                    messages=[{"role": "user", "content": "x"}], max_tokens=50, timeout=30
                )
            )
        assert client_with_key.last_finish_reason == "stop"