iomgaa
f1a0414be8
流式 SSE + 三层看门狗 + 重试退避 + 熔断 + Redis 缓存 + 遥测。 provider 差异处理（DeepSeek reasoning_content vs Qwen think 标签）。 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
307 lines
9.8 KiB
Python
307 lines
9.8 KiB
Python
"""GovernedLLMClient 单元测试 — 四层治理栈核心路径验证。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from adapters.breaker import CircuitBreaker
|
|
from adapters.llm import CircuitOpenError, GovernedLLMClient
|
|
from core.types import LLMResponse
|
|
|
|
# ── 辅助类 ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class _FakeRedisCache:
|
|
"""内存字典替身，模拟 RedisResponseCache 的 get/set 接口。"""
|
|
|
|
def __init__(self) -> None:
|
|
self._store: dict[str, LLMResponse] = {}
|
|
|
|
async def get(
|
|
self, model: str, messages: list[dict[str, str]]
|
|
) -> LLMResponse | None:
|
|
key = f"{model}:{json.dumps(messages, sort_keys=True)}"
|
|
return self._store.get(key)
|
|
|
|
async def set(
|
|
self,
|
|
model: str,
|
|
messages: list[dict[str, str]],
|
|
response: LLMResponse,
|
|
) -> None:
|
|
key = f"{model}:{json.dumps(messages, sort_keys=True)}"
|
|
self._store[key] = response
|
|
|
|
|
|
class _FakeTelemetry:
|
|
"""记录所有 record_llm_call 调用的替身。"""
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[dict[str, Any]] = []
|
|
|
|
async def record_llm_call(self, **kwargs: Any) -> None:
|
|
self.calls.append(kwargs)
|
|
|
|
|
|
# ── 工具函数 ─────────────────────────────────────────────────────────
|
|
|
|
|
|
async def _async_line_iter(lines: list[str]):
|
|
"""将字符串列表转为异步行迭代器。"""
|
|
for line in lines:
|
|
yield line
|
|
|
|
|
|
def _make_sse_lines(
|
|
content_chunks: list[str],
|
|
*,
|
|
reasoning_chunks: list[str] | None = None,
|
|
model: str = "test-model",
|
|
usage: dict[str, int] | None = None,
|
|
) -> list[str]:
|
|
"""构造 SSE 文本行序列，模拟 OpenAI 兼容流式响应。"""
|
|
lines: list[str] = []
|
|
|
|
# 先产出 reasoning 帧
|
|
for chunk in reasoning_chunks or []:
|
|
frame = {
|
|
"choices": [{"delta": {"reasoning_content": chunk}}],
|
|
"model": model,
|
|
}
|
|
lines.append(f"data: {json.dumps(frame)}\n")
|
|
|
|
# 再产出 content 帧
|
|
for chunk in content_chunks:
|
|
frame = {
|
|
"choices": [{"delta": {"content": chunk}}],
|
|
"model": model,
|
|
}
|
|
lines.append(f"data: {json.dumps(frame)}\n")
|
|
|
|
# usage 帧
|
|
if usage is None:
|
|
usage = {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
|
|
frame_usage = {"choices": [], "usage": usage, "model": model}
|
|
lines.append(f"data: {json.dumps(frame_usage)}\n")
|
|
|
|
# 终止标记
|
|
lines.append("data: [DONE]\n")
|
|
return lines
|
|
|
|
|
|
def _build_client(
|
|
*,
|
|
breaker: CircuitBreaker | None = None,
|
|
cache: _FakeRedisCache | None = None,
|
|
telemetry: _FakeTelemetry | None = None,
|
|
provider: str = "test-provider",
|
|
model: str = "test-model",
|
|
thinking: bool = False,
|
|
) -> GovernedLLMClient:
|
|
"""构造标准测试客户端。"""
|
|
return GovernedLLMClient(
|
|
model=model,
|
|
base_url="http://localhost:8000/v1",
|
|
api_key="test-key",
|
|
provider=provider,
|
|
thinking=thinking,
|
|
breaker=breaker or CircuitBreaker(fail_threshold=3, cooldown_s=60),
|
|
cache=cache or _FakeRedisCache(),
|
|
telemetry=telemetry or _FakeTelemetry(),
|
|
timeout_s=30.0,
|
|
ttft_timeout_s=10.0,
|
|
inter_token_timeout_s=5.0,
|
|
max_retries=3,
|
|
retry_base_delay_s=0.0, # 测试不等待
|
|
retry_max_delay_s=0.0,
|
|
)
|
|
|
|
|
|
# ── 测试用例 ─────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_circuit_open_raises():
|
|
"""熔断器开启 → 直接抛 CircuitOpenError，不发请求。"""
|
|
breaker = CircuitBreaker(fail_threshold=1, cooldown_s=600)
|
|
breaker.force_open("test-provider", time.monotonic())
|
|
client = _build_client(breaker=breaker)
|
|
|
|
with pytest.raises(CircuitOpenError):
|
|
await client.chat([{"role": "user", "content": "hello"}])
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cache_hit_returns_cached_and_records_telemetry():
|
|
"""缓存命中 → 返回缓存响应 + 新 call_id + 遥测记录 cache_hit=True。"""
|
|
cache = _FakeRedisCache()
|
|
telemetry = _FakeTelemetry()
|
|
client = _build_client(cache=cache, telemetry=telemetry)
|
|
|
|
# 预填缓存
|
|
cached_resp = LLMResponse(
|
|
content="cached answer",
|
|
thinking="",
|
|
model="test-model",
|
|
provider="test-provider",
|
|
prompt_tokens=10,
|
|
completion_tokens=5,
|
|
latency_ms=0,
|
|
ttft_ms=None,
|
|
max_inter_token_ms=None,
|
|
cache_hit=True,
|
|
call_id="old-id",
|
|
)
|
|
messages = [{"role": "user", "content": "hello"}]
|
|
await cache.set("test-model", messages, cached_resp)
|
|
|
|
result = await client.chat(messages)
|
|
|
|
assert result.content == "cached answer"
|
|
assert result.cache_hit is True
|
|
# 每次调用都生成新 call_id
|
|
assert result.call_id != "old-id"
|
|
|
|
# 遥测已记录
|
|
assert len(telemetry.calls) == 1
|
|
assert telemetry.calls[0]["cache_hit"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_successful_streaming_call():
|
|
"""正常流式调用 → 正确累积 content + thinking + 采集 ttft_ms + tokens。"""
|
|
telemetry = _FakeTelemetry()
|
|
client = _build_client(telemetry=telemetry)
|
|
messages = [{"role": "user", "content": "hello"}]
|
|
|
|
sse_lines = _make_sse_lines(
|
|
["Hello", " world"],
|
|
reasoning_chunks=["Let me", " think"],
|
|
usage={"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30},
|
|
)
|
|
|
|
async def fake_stream_request(_messages: Any) -> list[str]:
|
|
return sse_lines
|
|
|
|
with patch.object(client, "_stream_request", side_effect=fake_stream_request):
|
|
result = await client.chat(messages)
|
|
|
|
assert result.content == "Hello world"
|
|
assert result.thinking == "Let me think"
|
|
assert result.prompt_tokens == 20
|
|
assert result.completion_tokens == 10
|
|
assert result.ttft_ms is not None
|
|
assert result.ttft_ms >= 0
|
|
assert result.cache_hit is False
|
|
|
|
# 遥测已记录
|
|
assert len(telemetry.calls) == 1
|
|
assert telemetry.calls[0]["cache_hit"] is False
|
|
assert telemetry.calls[0]["error"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transient_error_retries_and_records_telemetry():
|
|
"""ConnectError 重试 2 次后成功 → 结果正确 + 遥测记录成功。"""
|
|
telemetry = _FakeTelemetry()
|
|
client = _build_client(telemetry=telemetry)
|
|
messages = [{"role": "user", "content": "hello"}]
|
|
|
|
sse_lines = _make_sse_lines(["OK"])
|
|
call_count = 0
|
|
|
|
async def flaky_stream(_messages: Any) -> list[str]:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count <= 2:
|
|
raise httpx.ConnectError("connection refused")
|
|
return sse_lines
|
|
|
|
with patch.object(client, "_stream_request", side_effect=flaky_stream):
|
|
result = await client.chat(messages)
|
|
|
|
assert result.content == "OK"
|
|
assert call_count == 3
|
|
|
|
# 前 2 次失败也应记遥测
|
|
error_calls = [c for c in telemetry.calls if c.get("error") is not None]
|
|
assert len(error_calls) == 2
|
|
# 最后 1 次成功也记了遥测
|
|
success_calls = [c for c in telemetry.calls if c.get("error") is None]
|
|
assert len(success_calls) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fatal_error_force_opens_breaker():
|
|
"""401 → force_open 熔断器 + 遥测记录 error。"""
|
|
breaker = CircuitBreaker(fail_threshold=3, cooldown_s=60)
|
|
telemetry = _FakeTelemetry()
|
|
client = _build_client(breaker=breaker, telemetry=telemetry)
|
|
messages = [{"role": "user", "content": "hello"}]
|
|
|
|
mock_response = httpx.Response(
|
|
status_code=401,
|
|
request=httpx.Request("POST", "http://localhost:8000/v1/chat/completions"),
|
|
)
|
|
|
|
async def auth_fail(_messages: Any) -> list[str]:
|
|
raise httpx.HTTPStatusError(
|
|
"Unauthorized", request=mock_response.request, response=mock_response
|
|
)
|
|
|
|
with patch.object(client, "_stream_request", side_effect=auth_fail), \
|
|
pytest.raises(httpx.HTTPStatusError):
|
|
await client.chat(messages)
|
|
|
|
# 熔断器已被 force_open
|
|
assert breaker.is_open("test-provider", time.monotonic())
|
|
|
|
# 遥测已记录 error
|
|
assert len(telemetry.calls) == 1
|
|
assert telemetry.calls[0]["error"] is not None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_qwen_thinking_stripped():
|
|
"""Qwen 模型 <think> 标签被正确剥离到 thinking 字段。"""
|
|
from adapters.llm import _strip_qwen_thinking
|
|
|
|
content_with_think = "<think>这是思考过程</think>这是正式输出"
|
|
content, thinking = _strip_qwen_thinking(content_with_think)
|
|
assert content == "这是正式输出"
|
|
assert thinking == "这是思考过程"
|
|
|
|
# 无 think 标签时原样返回
|
|
plain = "纯文本输出"
|
|
content2, thinking2 = _strip_qwen_thinking(plain)
|
|
assert content2 == "纯文本输出"
|
|
assert thinking2 == ""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_parent_call_id_forwarded_to_telemetry():
|
|
"""parent_call_id 和 session_id 正确传递到遥测记录。"""
|
|
telemetry = _FakeTelemetry()
|
|
client = _build_client(telemetry=telemetry)
|
|
messages = [{"role": "user", "content": "hello"}]
|
|
|
|
sse_lines = _make_sse_lines(["result"])
|
|
|
|
async def fake_stream(_messages: Any) -> list[str]:
|
|
return sse_lines
|
|
|
|
with patch.object(client, "_stream_request", side_effect=fake_stream):
|
|
await client.chat(
|
|
messages, session_id="sess-42", parent_call_id="parent-99"
|
|
)
|
|
|
|
assert len(telemetry.calls) == 1
|
|
assert telemetry.calls[0]["session_id"] == "sess-42"
|
|
assert telemetry.calls[0]["parent_call_id"] == "parent-99"
|