iomgaa
ee90fb8e88
- _stream_request + _consume_stream 合并为 _call_streaming， 在 async with self._http.stream() 内直接消费流， 确保看门狗作用于真实 HTTP 流而非内存列表（Critical 1） - breaker 检查移到 call_id 生成之前（Critical 2） - cache/ttft_timeout_s/inter_token_timeout_s 支持 None（Important 1） - 重试失败路径遥测使用统一 call_id（Important 2） Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
280 lines
9.0 KiB
Python
280 lines
9.0 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)
|
|
|
|
|
|
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"}]
|
|
|
|
async def fake_call_streaming(
|
|
_messages: Any,
|
|
) -> tuple[str, str, float | None, float | None, dict]:
|
|
return (
|
|
"Hello world",
|
|
"Let me think",
|
|
1.5, # ttft_ms
|
|
0.8, # max_inter_token_ms
|
|
{"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30},
|
|
)
|
|
|
|
with patch.object(client, "_call_streaming", side_effect=fake_call_streaming):
|
|
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"}]
|
|
|
|
call_count = 0
|
|
|
|
async def flaky_call_streaming(
|
|
_messages: Any,
|
|
) -> tuple[str, str, float | None, float | None, dict]:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count <= 2:
|
|
raise httpx.ConnectError("connection refused")
|
|
return (
|
|
"OK",
|
|
"",
|
|
1.0,
|
|
None,
|
|
{"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
|
)
|
|
|
|
with patch.object(client, "_call_streaming", side_effect=flaky_call_streaming):
|
|
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
|
|
|
|
# 所有遥测记录应使用同一个 call_id（Important 2 修复验证）
|
|
call_ids = {c["call_id"] for c in telemetry.calls}
|
|
assert len(call_ids) == 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,
|
|
) -> tuple[str, str, float | None, float | None, dict]:
|
|
raise httpx.HTTPStatusError(
|
|
"Unauthorized", request=mock_response.request, response=mock_response
|
|
)
|
|
|
|
with (
|
|
patch.object(client, "_call_streaming", 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"}]
|
|
|
|
async def fake_call_streaming(
|
|
_messages: Any,
|
|
) -> tuple[str, str, float | None, float | None, dict]:
|
|
return (
|
|
"result",
|
|
"",
|
|
1.0,
|
|
None,
|
|
{"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
|
)
|
|
|
|
with patch.object(client, "_call_streaming", side_effect=fake_call_streaming):
|
|
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"
|