sanguo_vnpy_v2/tests/api/test_llm_client.py

# tests/api/test_llm_client.py
"""LLM 薄 client:重试/参数自愈/strict 重试.全部 MockTransport,零真调用."""
import asyncio

import httpx
import pytest

from sanguo_api.llm.client import LLMClient
from sanguo_api.llm.provider import LLMConfig

CFG = LLMConfig("glm_cn", "https://open.bigmodel.cn/api/paas/v4/",
                "glm-5.3-flash", "sk-test")
MSGS = [{"role": "user", "content": "hi"}]


def _ok(content='{"a": 1}'):
    return httpx.Response(200, json={"choices": [
        {"message": {"content": content}}],
        "usage": {"prompt_tokens": 10, "completion_tokens": 5}})


def _make(handler):
    calls, sleeps = [], []

    async def _rec(s):
        sleeps.append(s)

    transport = httpx.MockTransport(handler)
    client = LLMClient(CFG, transport=transport, sleep=_rec)
    return client, calls, sleeps


class TestHappyPath:
    def test_chat_json_returns_dict_and_bears_model(self):
        def h(request):
            calls.append(request)
            import json as j
            body = j.loads(request.content)
            assert body["model"] == "glm-5.3-flash"
            assert body["response_format"] == {"type": "json_object"}
            return _ok()
        client, calls, _ = _make(h)
        out = asyncio.run(client.chat_json(MSGS))
        assert out == {"a": 1}
        assert calls[0].headers["authorization"] == "Bearer sk-test"


class TestRetry:
    def test_429_honors_retry_after_then_succeeds(self):
        def h(request):
            calls.append(request)
            if len(calls) == 1:
                return httpx.Response(429, headers={"Retry-After": "2"},
                                      json={"error": "rate"})
            return _ok()
        client, calls, sleeps = _make(h)
        assert asyncio.run(client.chat_json(MSGS)) == {"a": 1}
        assert len(calls) == 2 and sleeps == [2.0]

    def test_5xx_exhausts_attempts_raises_llm_error(self):
        def h(request):
            calls.append(request)
            return httpx.Response(502, json={"error": "bad gw"})
        client, calls, sleeps = _make(h)
        from sanguo_api.llm import LLMError
        with pytest.raises(LLMError) as e:
            asyncio.run(client.chat_json(MSGS))
        assert "502" in str(e.value)
        assert len(calls) == 3  # max_attempts=3

    def test_connect_error_retries_then_succeeds(self):
        def h(request):
            calls.append(request)
            if len(calls) == 1:
                raise httpx.ConnectError("refused")
            return _ok()
        client, calls, sleeps = _make(h)
        assert asyncio.run(client.chat_json(MSGS)) == {"a": 1}
        assert sleeps and len(calls) == 2

    def test_connect_error_exhausts_raises_llm_error(self):
        def h(request):
            calls.append(request)
            raise httpx.ConnectError("refused")
        client, calls, _ = _make(h)
        from sanguo_api.llm import LLMError
        with pytest.raises(LLMError, match="连接失败"):
            asyncio.run(client.chat_json(MSGS))
        assert len(calls) == 3  # max_attempts=3

    def test_401_raises_immediately(self):
        def h(request):
            calls.append(request)
            return httpx.Response(401, json={"error": "bad key"})
        client, calls, _ = _make(h)
        from sanguo_api.llm import LLMError
        with pytest.raises(LLMError):
            asyncio.run(client.chat_json(MSGS))
        assert len(calls) == 1


class TestParamSelfHeal:
    def test_400_rejecting_response_format_strips_and_retries(self):
        def h(request):
            calls.append(request)
            import json as j
            body = j.loads(request.content)
            if "response_format" in body:
                return httpx.Response(400, text='{"error":"response_format not supported"}')
            return _ok()
        client, calls, _ = _make(h)
        assert asyncio.run(client.chat_json(MSGS)) == {"a": 1}
        assert len(calls) == 2

    def test_400_unknown_raises_with_upstream_detail(self):
        def h(request):
            calls.append(request)
            return httpx.Response(400, text="quota exceeded 余额不足")
        client, calls, _ = _make(h)
        from sanguo_api.llm import LLMError
        with pytest.raises(LLMError) as e:
            asyncio.run(client.chat_json(MSGS))
        assert "余额不足" in str(e.value)


class TestStrictRetry:
    def test_non_json_content_appends_strict_and_retries_once(self):
        def h(request):
            calls.append(request)
            import json as j
            body = j.loads(request.content)
            if len(calls) == 1:
                return _ok(content="我觉得应该说……")
            assert any("Strictly output valid JSON" in m["content"]
                       for m in body["messages"])
            return _ok()
        client, calls, _ = _make(h)
        assert asyncio.run(client.chat_json(MSGS)) == {"a": 1}
        assert len(calls) == 2

    def test_non_json_twice_raises(self):
        def h(request):
            calls.append(request)
            return _ok(content="还是不行")
        client, calls, _ = _make(h)
        from sanguo_api.llm import LLMError
        with pytest.raises(LLMError):
            asyncio.run(client.chat_json(MSGS))
        assert len(calls) == 2


class TestWantUsage:
    """P4-3 漏斗 token 计量: want_usage=True 返 (obj, usage), 默认零变化."""

    def test_want_usage_returns_tuple(self):
        client, _, _ = _make(lambda r: _ok())
        obj, usage = asyncio.run(client.chat_json(MSGS, want_usage=True))
        assert obj == {"a": 1}
        assert usage["prompt_tokens"] == 10
        assert usage["completion_tokens"] == 5

    def test_default_backcompat_returns_dict(self):
        client, _, _ = _make(lambda r: _ok())
        out = asyncio.run(client.chat_json(MSGS))
        assert out == {"a": 1} and not isinstance(out, tuple)