e69ca4c82c
A client used to close whatever transport, recorder or cache it happened to hold, injected or not, so the first client to shut down killed the backend its siblings were still using. That is why the explicit-sharing path the architecture prescribes was unusable in practice and downstream projects fell back to one private instance per client. The mirror image of the same gap: the redis clients the factories build for the limiter and the breaker were never closed at all, because nobody kept a reference to them once they were handed to the retry middleware. Ownership is now stated once, the way RedisLimiter already stated it: whoever builds a resource closes it, injected ones are left alone. The constructor is the full-injection path, so it owns nothing by default and only the factories mark what they built. RedisCache gains the same rule for its own client, and the three copies of the "probe for aclose, fall back to close" dance collapse into a single helper so the next correction cannot land in only one of them.
838 lines
34 KiB
Python
838 lines
34 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
|