273 lines
9.6 KiB
Python
273 lines
9.6 KiB
Python
"""GatewayClient 装配与端到端(fake 后端 + MockTransport)测试(设计 §2.4)。"""
|
|
|
|
import asyncio
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from polygateway import (
|
|
AllSourcesExhausted,
|
|
GatewayClient,
|
|
GatewaySettings,
|
|
gather_bounded,
|
|
)
|
|
from polygateway.backends.memory.breaker import InMemoryGate
|
|
from polygateway.backends.memory.cache import InMemoryCache
|
|
from polygateway.backends.memory.limiter import InMemoryLimiter
|
|
from polygateway.sources import RoundRobinSelector
|
|
from polygateway.structured.json_repair import JsonRepairStrategy
|
|
from polygateway.transports.openai_compat import OpenAICompatTransport
|
|
from polygateway.types import (
|
|
BackpressurePolicy,
|
|
BreakerConfig,
|
|
GlobalLimits,
|
|
RetryPolicy,
|
|
SourceConfig,
|
|
)
|
|
|
|
_REPO = Path(__file__).resolve().parents[2]
|
|
|
|
_ENV = {
|
|
"LLM__QWEN__1__BASE_URL": "https://gw.example/v1",
|
|
"LLM__QWEN__1__API_KEY": "sk-a",
|
|
"LLM__QWEN__1__MODEL": "qwen-max",
|
|
"LLM__QWEN__1__TIMEOUT_S": "120",
|
|
"LLM_MAX_RETRIES": "3",
|
|
"LLM_RETRY_BASE_DELAY": "2.0",
|
|
"LLM_RETRY_MAX_DELAY": "30.0",
|
|
"LLM_CIRCUIT_BREAKER_THRESHOLD": "5",
|
|
"LLM_CIRCUIT_BREAKER_COOLDOWN": "60",
|
|
"PGW_CACHE_BACKEND": "none",
|
|
"PGW_TELEMETRY_BACKEND": "none",
|
|
}
|
|
|
|
|
|
def _sse(content='{"answer": 1}'):
|
|
chunk = json.dumps({"choices": [{"delta": {"content": content}}]})
|
|
usage = json.dumps({"choices": [], "usage": {"prompt_tokens": 3, "completion_tokens": 4}})
|
|
body = f"data: {chunk}\n\ndata: {usage}\n\ndata: [DONE]\n\n"
|
|
return httpx.Response(200, content=body.encode(), headers={"content-type": "text/event-stream"})
|
|
|
|
|
|
def _source(name="qwen_1", **overrides):
|
|
base = {
|
|
"name": name,
|
|
"provider": "qwen",
|
|
"base_url": "https://gw.example/v1",
|
|
"api_key": "sk",
|
|
"model": "qwen-max",
|
|
"timeout_s": 10.0,
|
|
}
|
|
base.update(overrides)
|
|
return SourceConfig(**base)
|
|
|
|
|
|
def _client(sources=None, handler=None, *, limiter=None, quota_full="wait", **overrides):
|
|
sources = sources or [_source()]
|
|
handler = handler or (lambda request: _sse())
|
|
transport = OpenAICompatTransport(
|
|
client_factory=lambda src: httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
|
)
|
|
defaults = {
|
|
"scope": "llm",
|
|
"sources": sources,
|
|
"selector": RoundRobinSelector(),
|
|
"limiter": limiter
|
|
or InMemoryLimiter(
|
|
scope="llm",
|
|
sources={s.name: s for s in sources},
|
|
global_limits=GlobalLimits(0, 0, 0),
|
|
),
|
|
"breaker": InMemoryGate(config=BreakerConfig(5, 60.0, 120.0)),
|
|
"transport": transport,
|
|
"retry": RetryPolicy(3, 2.0, 30.0),
|
|
"backpressure": BackpressurePolicy(300.0, 0.01),
|
|
"quota_full": quota_full,
|
|
"structured_strategy": JsonRepairStrategy(),
|
|
}
|
|
defaults.update(overrides)
|
|
return GatewayClient(**defaults)
|
|
|
|
|
|
class TestChatEndToEnd:
|
|
async def test_plain_chat(self):
|
|
async with _client() as client:
|
|
resp = await client.chat([{"role": "user", "content": "hi"}])
|
|
assert resp.content == '{"answer": 1}'
|
|
assert resp.source_name == "qwen_1" and resp.prompt_tokens == 3
|
|
|
|
async def test_structured_json_tier(self):
|
|
async with _client() as client:
|
|
resp = await client.chat([{"role": "user", "content": "hi"}], structured="json")
|
|
assert resp.structured_data == {"answer": 1}
|
|
|
|
async def test_cache_hit_roundtrip(self):
|
|
client = _client(cache=InMemoryCache(), cache_namespace="proj", cache_ttl_s=3600)
|
|
async with client:
|
|
first = await client.chat([{"role": "user", "content": "hi"}])
|
|
second = await client.chat([{"role": "user", "content": "hi"}])
|
|
assert first.cache_hit is False and second.cache_hit is True
|
|
|
|
async def test_structured_unavailable_fails_loudly(self):
|
|
async with _client(structured_strategy=None) as client:
|
|
with pytest.raises(ImportError, match="structured"):
|
|
await client.chat([{"role": "user", "content": "hi"}], structured="json")
|
|
|
|
|
|
class TestFactories:
|
|
def test_from_env_assembles(self):
|
|
client = GatewayClient.from_env("LLM", env=_ENV)
|
|
assert isinstance(client, GatewayClient)
|
|
|
|
def test_from_env_unknown_provider_fails_at_assembly(self):
|
|
env = dict(_ENV)
|
|
for key in list(env):
|
|
if key.startswith("LLM__QWEN__"):
|
|
env[key.replace("QWEN", "GLM")] = env.pop(key)
|
|
with pytest.raises(ValueError, match="glm"):
|
|
GatewayClient.from_env("LLM", env=env)
|
|
|
|
def test_from_env_builds_redis_governance_backends(self):
|
|
"""M2: 配置取 redis 时装配出 Redis 后端(构造不连库,unit 可测)。"""
|
|
from polygateway.backends.redis.breaker import RedisGate
|
|
from polygateway.backends.redis.limiter import RedisLimiter
|
|
from polygateway.client import _build_breaker, _build_limiter
|
|
|
|
env = dict(
|
|
_ENV,
|
|
PGW_LIMITER_BACKEND="redis",
|
|
PGW_BREAKER_BACKEND="redis",
|
|
REDIS_URL="redis://:pw@10.0.0.1:6379/3",
|
|
)
|
|
settings = GatewaySettings.from_env("LLM", env=env)
|
|
assert isinstance(_build_limiter(settings, list(settings.sources)), RedisLimiter)
|
|
assert isinstance(_build_breaker(settings), RedisGate)
|
|
|
|
def test_from_settings_respects_injection(self):
|
|
settings = GatewaySettings.from_env("LLM", env=_ENV)
|
|
shared = InMemoryLimiter(
|
|
scope="shared",
|
|
sources={s.name: s for s in settings.sources},
|
|
global_limits=GlobalLimits(0, 0, 0),
|
|
)
|
|
client = GatewayClient.from_settings(settings, limiter=shared)
|
|
assert isinstance(client, GatewayClient)
|
|
|
|
|
|
class TestSharedBackend:
|
|
async def test_two_clients_share_global_concurrency_gate(self):
|
|
"""VT R5: 两个逻辑角色显式注入同一 limiter → 共享全局并发闸。"""
|
|
src_a, src_b = _source("role_a_1"), _source("role_b_1")
|
|
shared = InMemoryLimiter(
|
|
scope="shared",
|
|
sources={"role_a_1": src_a, "role_b_1": src_b}, # 共享后端持源并集
|
|
global_limits=GlobalLimits(max_concurrency=1, rpm=0, tpm=0),
|
|
)
|
|
started = asyncio.Event()
|
|
|
|
async def slow_handler(request):
|
|
started.set()
|
|
await asyncio.sleep(0.2)
|
|
return _sse()
|
|
|
|
client_a = _client([src_a], slow_handler, limiter=shared)
|
|
client_b = _client([src_b], limiter=shared, quota_full="fail_fast")
|
|
task = asyncio.ensure_future(client_a.chat([{"role": "user", "content": "x"}]))
|
|
await started.wait()
|
|
with pytest.raises(AllSourcesExhausted) as ei:
|
|
await client_b.chat([{"role": "user", "content": "y"}])
|
|
assert ei.value.reason == "quota_exhausted" # 全局闸被 A 占满 → B 立即失败
|
|
await task
|
|
|
|
async def test_aclose_idempotent(self):
|
|
client = _client()
|
|
await client.aclose()
|
|
await client.aclose()
|
|
|
|
|
|
class TestGatherBounded:
|
|
async def test_order_preserved_and_concurrency_capped(self):
|
|
peak = {"now": 0, "max": 0}
|
|
|
|
async def work(i):
|
|
peak["now"] += 1
|
|
peak["max"] = max(peak["max"], peak["now"])
|
|
await asyncio.sleep(0.01)
|
|
peak["now"] -= 1
|
|
return i
|
|
|
|
results = await gather_bounded((work(i) for i in range(10)), concurrency=3)
|
|
assert results == list(range(10))
|
|
assert peak["max"] <= 3
|
|
|
|
async def test_exception_propagates(self):
|
|
async def boom():
|
|
raise RuntimeError("x")
|
|
|
|
async def ok():
|
|
return 1
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await gather_bounded([ok(), boom()], concurrency=2)
|
|
|
|
async def test_invalid_concurrency(self):
|
|
with pytest.raises(ValueError):
|
|
await gather_bounded([], concurrency=0)
|
|
|
|
|
|
def _load_reference_protocol(insert_path: str, module: str):
|
|
sys.path.insert(0, str(_REPO / insert_path))
|
|
try:
|
|
import importlib
|
|
|
|
return importlib.import_module(module).LLMProvider
|
|
finally:
|
|
sys.path.pop(0)
|
|
|
|
|
|
class TestReferenceProtocolCompat:
|
|
"""结构兼容断言(只读 import reference;失败即公共承诺破裂)。"""
|
|
|
|
def test_satisfies_govdoc_llm_provider(self):
|
|
try:
|
|
proto = _load_reference_protocol(
|
|
"reference/GovDoc-SaaS/packages/docagent-core/src", "docagent_core.protocols"
|
|
)
|
|
except ImportError:
|
|
# 兜底: 按 protocols.py:15-25 逐字复制的结构断言
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
@runtime_checkable
|
|
class proto(Protocol): # noqa: N801 — 复制自 GovDoc protocols.py:15-25
|
|
async def chat(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
*,
|
|
session_id: str | None = None,
|
|
parent_call_id: str | None = None,
|
|
): ...
|
|
|
|
assert isinstance(_client(), proto)
|
|
|
|
def test_satisfies_videotree_llm_provider(self):
|
|
try:
|
|
proto = _load_reference_protocol("reference/Video-Tree-TRM5", "core.protocols")
|
|
except ImportError:
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
@runtime_checkable
|
|
class proto(Protocol): # noqa: N801 — 复制自 VT core/protocols.py:18-29
|
|
async def chat(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
*,
|
|
session_id: str | None = None,
|
|
parent_call_id: str | None = None,
|
|
cache_salt: str | None = None,
|
|
): ...
|
|
|
|
assert isinstance(_client(), proto)
|