f958138e83
Telemetry degradation used to be a single warning and a private boolean. In a long-running process that is indistinguishable from telemetry working: issue #15 was only found by hand-reconciling milestone log lines against llm_calls rows, after 19 calls had silently gone unrecorded. The SQLite side was worse — once init failed, every write returned without even a log line. Degradation now has one shared owner. TelemetryStatusTracker holds the state machine (enter/recover/drop/should-retry), announces entry and recovery once each, and repeats the drop count under a row-and-time double threshold so a degraded backend neither floods the log nor goes quiet. Both recorders hold one; both count the rows they drop. For programmatic consumers, TelemetryStatus is a frozen snapshot exposed as telemetry_status on all three clients, resolved through a single isinstance check. It is a separate optional port rather than a member of TelemetryRecorder: that protocol is @runtime_checkable, so adding an attribute would make every implementation that only defines record_llm_call stop satisfying it — downstream isinstance assertions would break on upgrade. The existing assertion in test_ports.py is what keeps that decision honest. Failure criteria are deliberately untouched here: Postgres still treats a pool failure as permanent, only now visibly. `_failed` and the tracker therefore both carry the verdict for the span of this one change; the cooldown rework collapses them into the tracker alone.
875 lines
35 KiB
Python
875 lines
35 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.structured.native_schema import NativeSchemaStrategy
|
|
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",
|
|
}
|
|
|
|
_OCR_ENV = {
|
|
"OCR__MONKEY__1__BASE_URL": "http://10.77.0.20:7866",
|
|
"OCR__MONKEY__1__API_KEY": "none",
|
|
"OCR__MONKEY__1__MODEL": "monkey-ocr",
|
|
"OCR__MONKEY__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 TestSamplingOverlay:
|
|
"""调用级采样参数入口(issue #4 Task 3)。"""
|
|
|
|
def _capturing_client(self, captured, **overrides):
|
|
def handler(request):
|
|
captured.append(json.loads(request.content))
|
|
return _sse()
|
|
|
|
return _client(handler=handler, **overrides)
|
|
|
|
async def test_overlay_reaches_request_body(self):
|
|
captured = []
|
|
async with self._capturing_client(captured) as client:
|
|
await client.chat([{"role": "user", "content": "hi"}], overlay={"seed": 42})
|
|
assert captured[0]["seed"] == 42
|
|
|
|
async def test_call_level_beats_config_level(self):
|
|
"""优先级: 调用级 > 配置级(设计决策 A)。"""
|
|
captured = []
|
|
source = _source(extra_body={"temperature": 0, "top_p": 0.9})
|
|
async with self._capturing_client(captured, sources=[source]) as client:
|
|
await client.chat([{"role": "user", "content": "hi"}], overlay={"temperature": 1})
|
|
assert captured[0]["temperature"] == 1 # 调用级覆盖
|
|
assert captured[0]["top_p"] == 0.9 # 配置级未被顶掉的键保留
|
|
|
|
async def test_structured_injection_beats_call_level(self):
|
|
"""结构化注入优先级最高: 它关系到响应能否被解析(设计决策 A)。"""
|
|
captured = []
|
|
client = self._capturing_client(captured, structured_strategy=NativeSchemaStrategy())
|
|
async with client:
|
|
await client.chat(
|
|
[{"role": "user", "content": "hi"}],
|
|
structured="json",
|
|
overlay={"response_format": {"type": "text"}},
|
|
)
|
|
assert captured[0]["response_format"] != {"type": "text"}
|
|
|
|
async def test_protected_key_rejected_before_onion(self):
|
|
"""保护键在进洋葱之前就报错,transport 一次都不该被碰到。"""
|
|
captured = []
|
|
async with self._capturing_client(captured) as client:
|
|
with pytest.raises(ValueError, match="stream"):
|
|
await client.chat([{"role": "user", "content": "hi"}], overlay={"stream": False})
|
|
assert captured == []
|
|
|
|
async def test_unserializable_value_rejected_before_onion(self):
|
|
"""裸 TypeError 会在 CacheMW 的降级 try 之外炸且无遥测(设计决策 B)。"""
|
|
captured = []
|
|
async with self._capturing_client(captured) as client:
|
|
with pytest.raises(ValueError, match="JSON"):
|
|
await client.chat(
|
|
[{"role": "user", "content": "hi"}], overlay={"temperature": object()}
|
|
)
|
|
assert captured == []
|
|
|
|
async def test_caller_dict_mutation_does_not_leak(self):
|
|
"""调用方逐次改 seed 复用同一 dict 是预期模式(设计决策 E)。"""
|
|
captured = []
|
|
caller_overlay = {"seed": 1}
|
|
async with self._capturing_client(captured) as client:
|
|
await client.chat([{"role": "user", "content": "hi"}], overlay=caller_overlay)
|
|
caller_overlay["seed"] = 2
|
|
await client.chat([{"role": "user", "content": "hi"}], overlay=caller_overlay)
|
|
assert [c["seed"] for c in captured] == [1, 2]
|
|
|
|
|
|
class _MemoryRecorder:
|
|
"""收下遥测行原样存起来;断言"哪些行被写了"必须能看到零行的情形。"""
|
|
|
|
def __init__(self):
|
|
self.rows = []
|
|
|
|
async def record_llm_call(self, **fields):
|
|
self.rows.append(fields)
|
|
|
|
|
|
class TestCallerDimensions:
|
|
"""调用方自定义维度进遥测(issue #11 Task 4)。"""
|
|
|
|
_MSG = [{"role": "user", "content": "hi"}]
|
|
|
|
async def test_dimensions_reach_telemetry_row(self):
|
|
recorder = _MemoryRecorder()
|
|
async with _client(telemetry=recorder) as client:
|
|
await client.chat(self._MSG, tenant_id="t1", meta={"batch": "b-42"})
|
|
row = recorder.rows[-1]
|
|
assert row["tenant_id"] == "t1"
|
|
assert json.loads(row["meta"]) == {"batch": "b-42"}
|
|
|
|
async def test_default_path_writes_sentinels(self):
|
|
"""不传两参数时落哨兵值而非 NULL(§4.4: NULL 在 RLS 下是永久不可见的黑洞)。"""
|
|
recorder = _MemoryRecorder()
|
|
async with _client(telemetry=recorder) as client:
|
|
await client.chat(self._MSG)
|
|
row = recorder.rows[-1]
|
|
assert row["tenant_id"] == "" and row["meta"] == "{}"
|
|
|
|
async def test_invalid_meta_key_rejected_before_any_telemetry(self):
|
|
"""校验早于遥测(§4.2 核心承诺): 放进洋葱就会被降级成 warning 而调用照常发出。"""
|
|
recorder = _MemoryRecorder()
|
|
async with _client(telemetry=recorder) as client:
|
|
with pytest.raises(ValueError, match="meta"):
|
|
await client.chat(self._MSG, meta={"BAD-KEY": 1})
|
|
assert recorder.rows == []
|
|
|
|
async def test_non_finite_float_rejected_before_any_telemetry(self):
|
|
"""nan 产出的是 PG 拒收的非法 JSON;放行等于把调用方 bug 变成静默丢遥测(§6)。"""
|
|
recorder = _MemoryRecorder()
|
|
async with _client(telemetry=recorder) as client:
|
|
with pytest.raises(ValueError, match="nan"):
|
|
await client.chat(self._MSG, meta={"k": float("nan")})
|
|
assert recorder.rows == []
|
|
|
|
async def test_meta_does_not_enter_cache_key(self):
|
|
"""仅 meta 不同必须仍命中缓存(F1): 进 key 会让存量缓存全量冷启动且不报错。
|
|
|
|
带对照组: 只断言"命中"的话,缓存 key 退化成常量(忽略一切输入)时本用例
|
|
照样绿——那是恒真断言。故再改一个**确实进 key** 的维度(namespace)断言
|
|
miss,证明 key 仍在区分输入,"meta 不进 key"才是被测出来的结论。
|
|
"""
|
|
cache = InMemoryCache() # 两个 client 共用一份存储,否则对照组的 miss 是白来的
|
|
client = _client(cache=cache, cache_namespace="proj", cache_ttl_s=3600)
|
|
async with client:
|
|
first = await client.chat(self._MSG, meta={"batch": "b-1"})
|
|
second = await client.chat(self._MSG, meta={"batch": "b-2"})
|
|
assert first.cache_hit is False and second.cache_hit is True
|
|
|
|
other_ns = _client(cache=cache, cache_namespace="other", cache_ttl_s=3600)
|
|
async with other_ns:
|
|
assert (await other_ns.chat(self._MSG, meta={"batch": "b-1"})).cache_hit is False
|
|
|
|
|
|
class TestModelFingerprint:
|
|
"""配置级采样参数须进缓存身份,否则改 temperature 后仍读旧缓存(决策 C)。"""
|
|
|
|
def test_empty_extra_body_keeps_legacy_fingerprint(self):
|
|
"""全源无 extra_body 时字面量与旧实现逐字相同,不触发存量缓存冷启动。"""
|
|
from polygateway.client import build_model_fingerprint
|
|
|
|
sources = [_source(), _source(name="qwen_2", model="qwen-plus")]
|
|
assert build_model_fingerprint(sources) == "qwen-max,qwen-plus"
|
|
|
|
def test_extra_body_changes_fingerprint(self):
|
|
from polygateway.client import build_model_fingerprint
|
|
|
|
plain = build_model_fingerprint([_source()])
|
|
tuned = build_model_fingerprint([_source(extra_body={"temperature": 0})])
|
|
assert plain != tuned
|
|
assert tuned.startswith("qwen-max|") # 旧字面量仍是前缀,便于人眼辨认
|
|
|
|
def test_enable_thinking_changes_fingerprint(self):
|
|
"""issue #5 配套: thinking 一旦真正改变请求体,就必须进缓存身份。
|
|
|
|
否则"关掉推理后重启"会读到开着推理时缓存的旧响应——issue #4 为
|
|
temperature 写过逐字相同的理由。
|
|
"""
|
|
from polygateway.client import build_model_fingerprint
|
|
|
|
plain = build_model_fingerprint([_source()])
|
|
off = build_model_fingerprint([_source(enable_thinking=False)])
|
|
on = build_model_fingerprint([_source(enable_thinking=True)])
|
|
assert len({plain, off, on}) == 3
|
|
|
|
def test_extra_body_only_fingerprint_is_byte_identical_to_before(self):
|
|
"""只配 extra_body、不表态 thinking 的存量源不得触发冷启动。
|
|
|
|
字面量在此硬编码: 这条断言的价值全在"逐字相同",改实现时必须先看见它红。
|
|
"""
|
|
import hashlib
|
|
import json
|
|
|
|
from polygateway.client import build_model_fingerprint
|
|
|
|
mark = json.dumps(["qwen-max", {"temperature": 0}], sort_keys=True, ensure_ascii=False)
|
|
expected = "qwen-max|" + hashlib.sha256(mark.encode("utf-8")).hexdigest()
|
|
assert build_model_fingerprint([_source(extra_body={"temperature": 0})]) == expected
|
|
|
|
def test_source_rename_does_not_change_fingerprint(self):
|
|
"""指纹按 (model, extra_body) 而非源名: 改名不该误触全量冷启动。"""
|
|
from polygateway.client import build_model_fingerprint
|
|
|
|
a = build_model_fingerprint([_source(name="qwen_1", extra_body={"temperature": 0})])
|
|
b = build_model_fingerprint([_source(name="renamed", extra_body={"temperature": 0})])
|
|
assert a == b
|
|
|
|
def test_differing_extra_body_across_sources_is_distinguished(self):
|
|
from polygateway.client import build_model_fingerprint
|
|
|
|
a = build_model_fingerprint([_source(extra_body={"temperature": 0})])
|
|
b = build_model_fingerprint([_source(extra_body={"temperature": 1})])
|
|
assert a != b
|
|
|
|
|
|
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 TestTelemetryTextCapWiring:
|
|
"""`PGW_TELEMETRY_TEXT_CAP` 必须走通全部三条 `from_settings` 装配路(issue #12)。
|
|
|
|
三条链路写的是**同一张** `llm_calls` 表:只接通 chat,embed 与 OCR 的行就
|
|
永远不受 cap 约束,同表内一半受控一半不受控——那正是本 issue 要消灭的状态。
|
|
"""
|
|
|
|
_CAP_ENV = dict(_ENV, PGW_TELEMETRY_TEXT_CAP="8")
|
|
_OCR_CAP_ENV = dict(_OCR_ENV, PGW_TELEMETRY_TEXT_CAP="8")
|
|
|
|
def test_gateway_from_settings_wires_the_cap(self):
|
|
settings = GatewaySettings.from_env("LLM", env=self._CAP_ENV)
|
|
client = GatewayClient.from_settings(settings, telemetry=_MemoryRecorder())
|
|
assert client._terminal._emitter._text_cap == 8
|
|
# 对照组: 不设该键时 emitter 拿到的必须是 None,否则 8 可能是硬编码来的
|
|
unset = GatewayClient.from_settings(
|
|
GatewaySettings.from_env("LLM", env=_ENV), telemetry=_MemoryRecorder()
|
|
)
|
|
assert unset._terminal._emitter._text_cap is None
|
|
|
|
def test_embedding_from_settings_wires_the_cap(self):
|
|
from polygateway.config import EmbeddingSettings
|
|
from polygateway.embedding import EmbeddingClient
|
|
|
|
gateway = GatewaySettings.from_env("LLM", env=self._CAP_ENV)
|
|
client = EmbeddingClient.from_settings(
|
|
EmbeddingSettings(gateway=gateway, batch_size=2), telemetry=_MemoryRecorder()
|
|
)
|
|
assert client._emitter._text_cap == 8
|
|
unset = EmbeddingClient.from_settings(
|
|
EmbeddingSettings(gateway=GatewaySettings.from_env("LLM", env=_ENV), batch_size=2),
|
|
telemetry=_MemoryRecorder(),
|
|
)
|
|
assert unset._emitter._text_cap is None
|
|
|
|
def test_ocr_from_settings_wires_the_cap(self):
|
|
from polygateway.config import OcrSettings
|
|
from polygateway.ocr import OcrClient
|
|
|
|
settings = OcrSettings.from_env("OCR", env=dict(self._OCR_CAP_ENV))
|
|
client = OcrClient.from_settings(settings, telemetry=_MemoryRecorder())
|
|
assert client._emitter._text_cap == 8
|
|
no_cap = dict(self._OCR_CAP_ENV)
|
|
no_cap.pop("PGW_TELEMETRY_TEXT_CAP")
|
|
unset = OcrClient.from_settings(
|
|
OcrSettings.from_env("OCR", env=no_cap), telemetry=_MemoryRecorder()
|
|
)
|
|
assert unset._emitter._text_cap is None
|
|
|
|
async def test_capped_body_reaches_the_recorder_end_to_end(self, monkeypatch):
|
|
"""装配路通了还不够: 真跑一次 chat,落库的 messages 与 response 确已截断。
|
|
|
|
`from_settings` 自建 transport(没有 client_factory 入口),故在装配点
|
|
换掉该类以接上 MockTransport——洋葱其余各层仍是 `from_settings` 装的真件。
|
|
"""
|
|
recorder = _MemoryRecorder()
|
|
long_text = "甲乙丙丁戊己庚辛壬癸" # 10 字,cap=8 → 略 2 字
|
|
monkeypatch.setattr(
|
|
"polygateway.client.OpenAICompatTransport",
|
|
lambda **kwargs: OpenAICompatTransport(
|
|
client_factory=lambda source: httpx.AsyncClient(
|
|
transport=httpx.MockTransport(lambda request: _sse(content=long_text))
|
|
)
|
|
),
|
|
)
|
|
settings = GatewaySettings.from_env("LLM", env=self._CAP_ENV)
|
|
async with GatewayClient.from_settings(settings, telemetry=recorder) as client:
|
|
await client.chat([{"role": "user", "content": long_text}])
|
|
row = recorder.rows[-1]
|
|
assert json.loads(row["messages"])[0]["content"] == "甲乙丙丁戊己庚辛…(略 2 字)"
|
|
assert row["response"] == "甲乙丙丁戊己庚辛…(略 2 字)"
|
|
|
|
def test_non_positive_cap_rejected_on_the_direct_construction_path(self):
|
|
"""直接构造是库承诺的另一条公共装配路;cap=0 会让每条正文只剩省略标记。"""
|
|
with pytest.raises(ValueError, match="text_cap"):
|
|
_client(telemetry=_MemoryRecorder(), text_cap=0)
|
|
|
|
|
|
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)
|
|
|
|
|
|
# —— 资源所有权纪律(issue #15 D 组): 谁建的谁关,注入的一律不碰 ——
|
|
|
|
_CACHE_ENV = dict(
|
|
_ENV, PGW_CACHE_BACKEND="memory", PGW_CACHE_NAMESPACE="proj", PGW_CACHE_TTL_S="3600"
|
|
)
|
|
|
|
|
|
class _Closable:
|
|
"""记 close 次数的假组件;所有权纪律的唯一观测点。"""
|
|
|
|
def __init__(self):
|
|
self.closed = 0
|
|
|
|
async def aclose(self):
|
|
self.closed += 1
|
|
|
|
|
|
class _SyncClosable:
|
|
"""只有同步 close 的假 recorder(SQLiteRecorder 形态,收敛后的 helper 须探测到)。"""
|
|
|
|
def __init__(self):
|
|
self.closed = 0
|
|
|
|
def close(self):
|
|
self.closed += 1
|
|
|
|
|
|
def _parts(*names):
|
|
return {name: _Closable() for name in names}
|
|
|
|
|
|
def _patch_builders(monkeypatch, built, *, transport_path):
|
|
"""把工厂的自建点换成可计数假件;transport 无注入入口,故恒自建。"""
|
|
monkeypatch.setattr(transport_path, lambda **kwargs: built["transport"])
|
|
monkeypatch.setattr("polygateway.client._build_limiter", lambda s, src: built["limiter"])
|
|
monkeypatch.setattr("polygateway.client._build_breaker", lambda s: built["breaker"])
|
|
monkeypatch.setattr("polygateway.client._build_telemetry", lambda s: built["telemetry"])
|
|
if "cache" in built:
|
|
monkeypatch.setattr("polygateway.client._build_cache", lambda s: built["cache"])
|
|
|
|
|
|
class TestGatewayClientOwnership:
|
|
"""`__init__` 是全量注入路径,经它传入的一切都归调用方(设计 §3.4)。"""
|
|
|
|
_GATEWAY_TRANSPORT = "polygateway.client.OpenAICompatTransport"
|
|
|
|
async def test_injected_components_are_never_closed(self):
|
|
"""共享 recorder/transport 被第一个关闭的 client 弄死,正是 R5 显式共享走不通的原因。"""
|
|
injected = _parts(*("transport", "telemetry", "cache", "limiter", "breaker"))
|
|
client = _client(
|
|
transport=injected["transport"],
|
|
telemetry=injected["telemetry"],
|
|
cache=injected["cache"],
|
|
cache_namespace="proj",
|
|
cache_ttl_s=3600,
|
|
limiter=injected["limiter"],
|
|
breaker=injected["breaker"],
|
|
)
|
|
await client.aclose()
|
|
assert {name: part.closed for name, part in injected.items()} == {
|
|
"transport": 0,
|
|
"telemetry": 0,
|
|
"cache": 0,
|
|
"limiter": 0,
|
|
"breaker": 0,
|
|
}
|
|
|
|
async def test_factory_closes_every_component_it_built(self, monkeypatch):
|
|
"""泄漏钉子: 自建的 redis limiter/breaker 今天没人关,连引用都没留。"""
|
|
built = _parts("transport", "telemetry", "cache", "limiter", "breaker")
|
|
_patch_builders(monkeypatch, built, transport_path=self._GATEWAY_TRANSPORT)
|
|
client = GatewayClient.from_settings(GatewaySettings.from_env("LLM", env=_CACHE_ENV))
|
|
await client.aclose()
|
|
assert {name: part.closed for name, part in built.items()} == {
|
|
"transport": 1,
|
|
"telemetry": 1,
|
|
"cache": 1,
|
|
"limiter": 1,
|
|
"breaker": 1,
|
|
}
|
|
|
|
async def test_factory_keeps_hands_off_injected_components(self, monkeypatch):
|
|
built = _parts("transport", "telemetry", "cache", "limiter", "breaker")
|
|
_patch_builders(monkeypatch, built, transport_path=self._GATEWAY_TRANSPORT)
|
|
injected = _parts("telemetry", "cache", "limiter", "breaker")
|
|
client = GatewayClient.from_settings(
|
|
GatewaySettings.from_env("LLM", env=_CACHE_ENV),
|
|
limiter=injected["limiter"],
|
|
breaker=injected["breaker"],
|
|
cache=injected["cache"],
|
|
telemetry=injected["telemetry"],
|
|
)
|
|
await client.aclose()
|
|
assert all(part.closed == 0 for part in injected.values())
|
|
assert built["transport"].closed == 1 # 工厂恒自建 transport,归 client
|
|
|
|
async def test_aclose_is_idempotent(self, monkeypatch):
|
|
built = _parts("transport", "telemetry", "cache", "limiter", "breaker")
|
|
_patch_builders(monkeypatch, built, transport_path=self._GATEWAY_TRANSPORT)
|
|
client = GatewayClient.from_settings(GatewaySettings.from_env("LLM", env=_CACHE_ENV))
|
|
await client.aclose()
|
|
await client.aclose()
|
|
assert all(part.closed == 1 for part in built.values())
|
|
|
|
async def test_sync_only_recorder_is_closed(self, monkeypatch):
|
|
"""SQLiteRecorder 只有同步 `close()`;收敛成 helper 之后这条分支不得丢。"""
|
|
built = _parts("transport", "cache", "limiter", "breaker")
|
|
recorder = _SyncClosable()
|
|
built["telemetry"] = recorder
|
|
_patch_builders(monkeypatch, built, transport_path=self._GATEWAY_TRANSPORT)
|
|
client = GatewayClient.from_settings(GatewaySettings.from_env("LLM", env=_CACHE_ENV))
|
|
await client.aclose()
|
|
assert recorder.closed == 1
|
|
|
|
|
|
def _embedding_client(**overrides):
|
|
from polygateway.embedding import EmbeddingClient
|
|
|
|
defaults = {
|
|
"scope": "embed",
|
|
"sources": [_source()],
|
|
"selector": RoundRobinSelector(),
|
|
"limiter": _Closable(),
|
|
"breaker": _Closable(),
|
|
"transport": _Closable(),
|
|
"retry": RetryPolicy(3, 2.0, 30.0),
|
|
"backpressure": BackpressurePolicy(300.0, 0.01),
|
|
"batch_size": 2,
|
|
}
|
|
defaults.update(overrides)
|
|
return EmbeddingClient(**defaults)
|
|
|
|
|
|
class TestEmbeddingClientOwnership:
|
|
"""三处必须各钉一次: 收敛成 helper 之后,有人把逻辑复制回去也得当场被发现。"""
|
|
|
|
async def test_injected_components_are_never_closed(self):
|
|
injected = _parts("transport", "telemetry", "limiter", "breaker")
|
|
client = _embedding_client(
|
|
transport=injected["transport"],
|
|
telemetry=injected["telemetry"],
|
|
limiter=injected["limiter"],
|
|
breaker=injected["breaker"],
|
|
)
|
|
await client.aclose()
|
|
assert all(part.closed == 0 for part in injected.values())
|
|
|
|
async def test_factory_closes_every_component_it_built(self, monkeypatch):
|
|
from polygateway.config import EmbeddingSettings
|
|
from polygateway.embedding import EmbeddingClient
|
|
|
|
built = _parts("transport", "telemetry", "limiter", "breaker")
|
|
_patch_builders(
|
|
monkeypatch,
|
|
built,
|
|
transport_path="polygateway.transports.openai_compat.OpenAICompatTransport",
|
|
)
|
|
settings = EmbeddingSettings(
|
|
gateway=GatewaySettings.from_env("LLM", env=_ENV), batch_size=2
|
|
)
|
|
client = EmbeddingClient.from_settings(settings)
|
|
await client.aclose()
|
|
assert all(part.closed == 1 for part in built.values())
|
|
|
|
async def test_factory_keeps_hands_off_injected_components(self, monkeypatch):
|
|
from polygateway.config import EmbeddingSettings
|
|
from polygateway.embedding import EmbeddingClient
|
|
|
|
built = _parts("transport", "telemetry", "limiter", "breaker")
|
|
_patch_builders(
|
|
monkeypatch,
|
|
built,
|
|
transport_path="polygateway.transports.openai_compat.OpenAICompatTransport",
|
|
)
|
|
injected = _parts("telemetry", "limiter", "breaker")
|
|
settings = EmbeddingSettings(
|
|
gateway=GatewaySettings.from_env("LLM", env=_ENV), batch_size=2
|
|
)
|
|
client = EmbeddingClient.from_settings(
|
|
settings,
|
|
limiter=injected["limiter"],
|
|
breaker=injected["breaker"],
|
|
telemetry=injected["telemetry"],
|
|
)
|
|
await client.aclose()
|
|
assert all(part.closed == 0 for part in injected.values())
|
|
assert built["transport"].closed == 1
|
|
|
|
|
|
def _ocr_client(**overrides):
|
|
from polygateway.ocr import OcrClient
|
|
|
|
defaults = {
|
|
"scope": "ocr",
|
|
"sources": [_source(name="m1", provider="monkey", model="monkey-ocr")],
|
|
"selector": RoundRobinSelector(),
|
|
"limiter": _Closable(),
|
|
"breaker": _Closable(),
|
|
"transport": _Closable(),
|
|
"retry": RetryPolicy(3, 2.0, 30.0),
|
|
"backpressure": BackpressurePolicy(300.0, 0.01),
|
|
}
|
|
defaults.update(overrides)
|
|
return OcrClient(**defaults)
|
|
|
|
|
|
class TestOcrClientOwnership:
|
|
async def test_injected_components_are_never_closed(self):
|
|
injected = _parts("transport", "telemetry", "limiter", "breaker")
|
|
client = _ocr_client(
|
|
transport=injected["transport"],
|
|
telemetry=injected["telemetry"],
|
|
limiter=injected["limiter"],
|
|
breaker=injected["breaker"],
|
|
)
|
|
await client.aclose()
|
|
assert all(part.closed == 0 for part in injected.values())
|
|
|
|
async def test_factory_closes_every_component_it_built(self, monkeypatch):
|
|
from polygateway.config import OcrSettings
|
|
from polygateway.ocr import OcrClient
|
|
|
|
built = _parts("transport", "telemetry", "limiter", "breaker")
|
|
_patch_builders(
|
|
monkeypatch,
|
|
built,
|
|
transport_path="polygateway.transports.monkey_ocr.MonkeyOcrTransport",
|
|
)
|
|
client = OcrClient.from_settings(OcrSettings.from_env("OCR", env=dict(_OCR_ENV)))
|
|
await client.aclose()
|
|
assert all(part.closed == 1 for part in built.values())
|
|
|
|
async def test_factory_keeps_hands_off_injected_components(self, monkeypatch):
|
|
from polygateway.config import OcrSettings
|
|
from polygateway.ocr import OcrClient
|
|
|
|
built = _parts("transport", "telemetry", "limiter", "breaker")
|
|
_patch_builders(
|
|
monkeypatch,
|
|
built,
|
|
transport_path="polygateway.transports.monkey_ocr.MonkeyOcrTransport",
|
|
)
|
|
injected = _parts("telemetry", "limiter", "breaker")
|
|
client = OcrClient.from_settings(
|
|
OcrSettings.from_env("OCR", env=dict(_OCR_ENV)),
|
|
limiter=injected["limiter"],
|
|
breaker=injected["breaker"],
|
|
telemetry=injected["telemetry"],
|
|
)
|
|
await client.aclose()
|
|
assert all(part.closed == 0 for part in injected.values())
|
|
assert built["transport"].closed == 1
|
|
|
|
|
|
class TestRedisCacheOwnership:
|
|
"""组件内部自建的连接归组件自己;照抄 RedisLimiter._owns_client 的正确先例。"""
|
|
|
|
async def test_injected_client_is_not_closed(self):
|
|
from polygateway.backends.redis_cache import RedisCache
|
|
|
|
client = _Closable()
|
|
await RedisCache(client).aclose()
|
|
assert client.closed == 0
|
|
|
|
async def test_self_built_client_is_closed_once(self, monkeypatch):
|
|
from types import SimpleNamespace
|
|
|
|
from polygateway.backends import redis_cache
|
|
|
|
built = _Closable()
|
|
monkeypatch.setattr(
|
|
redis_cache, "aioredis", SimpleNamespace(from_url=lambda url, **kwargs: built)
|
|
)
|
|
cache = redis_cache.RedisCache.from_url("redis://localhost:6379/0")
|
|
await cache.aclose()
|
|
await cache.aclose() # 幂等: 不重复关
|
|
assert built.closed == 1
|
|
|
|
|
|
class TestTelemetryStatusExposure:
|
|
"""降级状态的只读出口: 一处 isinstance 判定,三个 client 各钉一次(设计 §3.3)。"""
|
|
|
|
def _recorder(self, tmp_path):
|
|
from polygateway.telemetry.sqlite import SQLiteRecorder
|
|
|
|
return SQLiteRecorder(tmp_path / "telemetry.db", auto_migrate=True)
|
|
|
|
def _assert_snapshot(self, status):
|
|
from polygateway.types import TelemetryStatus
|
|
|
|
assert isinstance(status, TelemetryStatus)
|
|
assert status.degraded is False
|
|
|
|
def test_gateway_client_without_telemetry_reports_none(self):
|
|
assert _client().telemetry_status is None
|
|
|
|
def test_gateway_client_with_foreign_recorder_reports_none(self):
|
|
"""注入的第三方 recorder 不提供状态 → None,绝不得抛 AttributeError。"""
|
|
assert _client(telemetry=_Closable()).telemetry_status is None
|
|
|
|
def test_gateway_client_with_builtin_recorder_reports_snapshot(self, tmp_path):
|
|
self._assert_snapshot(_client(telemetry=self._recorder(tmp_path)).telemetry_status)
|
|
|
|
def test_embedding_client_exposes_the_same_outlet(self, tmp_path):
|
|
assert _embedding_client().telemetry_status is None
|
|
assert _embedding_client(telemetry=_Closable()).telemetry_status is None
|
|
self._assert_snapshot(
|
|
_embedding_client(telemetry=self._recorder(tmp_path)).telemetry_status
|
|
)
|
|
|
|
def test_ocr_client_exposes_the_same_outlet(self, tmp_path):
|
|
assert _ocr_client().telemetry_status is None
|
|
assert _ocr_client(telemetry=_Closable()).telemetry_status is None
|
|
self._assert_snapshot(_ocr_client(telemetry=self._recorder(tmp_path)).telemetry_status)
|