"""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_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)