feat: add retry middleware with per-attempt governance orchestration
This commit is contained in:
@@ -0,0 +1,281 @@
|
||||
"""RetryMW 尝试编排测试(保真蓝本 CHS governance.py:107-268)。
|
||||
|
||||
用真实内存后端 + FakeClock + 可编程 fake transport,验证换源/退避/熔断
|
||||
写回/permit 结算/取消穿透等治理行为。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from polygateway.backends.memory.breaker import InMemoryGate
|
||||
from polygateway.backends.memory.limiter import InMemoryLimiter
|
||||
from polygateway.errors import (
|
||||
AllSourcesExhausted,
|
||||
CircuitOpenError,
|
||||
RequestRejectedError,
|
||||
ResultInvalidError,
|
||||
SourceDeadError,
|
||||
TransientError,
|
||||
)
|
||||
from polygateway.middleware.retry import RetryMW
|
||||
from polygateway.ports import GateState
|
||||
from polygateway.sources import RoundRobinSelector, SourceCooldownMemo
|
||||
from polygateway.types import (
|
||||
BackpressurePolicy,
|
||||
BreakerConfig,
|
||||
ChatRequest,
|
||||
GlobalLimits,
|
||||
RetryPolicy,
|
||||
SourceConfig,
|
||||
TransportResult,
|
||||
)
|
||||
from tests.contracts.conftest import FakeClock
|
||||
|
||||
_BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0)
|
||||
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)
|
||||
|
||||
|
||||
def _src(name, **overrides):
|
||||
base = dict(
|
||||
name=name, provider="openai", base_url="https://gw.example/v1",
|
||||
api_key="sk", model="m", timeout_s=10.0,
|
||||
)
|
||||
base.update(overrides)
|
||||
return SourceConfig(**base)
|
||||
|
||||
|
||||
def _ok(content="ok"):
|
||||
return TransportResult(
|
||||
content=content, thinking="", prompt_tokens=10, completion_tokens=5,
|
||||
usage_source="measured", ttft_ms=12.0, max_inter_token_ms=3.0, raw={},
|
||||
)
|
||||
|
||||
|
||||
class FakeTransport:
|
||||
"""按脚本逐次返回结果或抛异常;记录每次 (source_name, call_id)。"""
|
||||
|
||||
def __init__(self, script):
|
||||
self.script = list(script)
|
||||
self.calls = []
|
||||
|
||||
async def complete(self, *, messages, source, stream, overlay, call_id):
|
||||
self.calls.append((source.name, call_id))
|
||||
action = self.script.pop(0)
|
||||
if isinstance(action, Exception):
|
||||
raise action
|
||||
if action == "hang":
|
||||
await asyncio.Event().wait()
|
||||
return action
|
||||
|
||||
|
||||
class FakeSleep:
|
||||
"""记录退避时长,立即返回(不真等)。"""
|
||||
|
||||
def __init__(self):
|
||||
self.delays = []
|
||||
|
||||
async def __call__(self, seconds):
|
||||
self.delays.append(seconds)
|
||||
|
||||
|
||||
def _harness(sources, script, *, clock=None, max_attempts=3, quota_full="wait",
|
||||
global_limits=_NO_GLOBAL, rng=lambda: 0.0):
|
||||
clock = clock or FakeClock()
|
||||
limiter = InMemoryLimiter(
|
||||
scope="llm", sources={s.name: s for s in sources},
|
||||
global_limits=global_limits, lease_ttl_s=100.0, now=clock,
|
||||
)
|
||||
gate = InMemoryGate(config=_BREAKER, now=clock)
|
||||
transport = FakeTransport(script)
|
||||
sleep = FakeSleep()
|
||||
mw = RetryMW(
|
||||
scope="llm", sources=sources, selector=RoundRobinSelector(),
|
||||
limiter=limiter, gate=gate, transport=transport,
|
||||
retry=RetryPolicy(max_attempts=max_attempts, backoff_base_s=2.0, backoff_max_s=30.0),
|
||||
backpressure=BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.01),
|
||||
quota_full=quota_full, cooldown_memo=SourceCooldownMemo(now=clock),
|
||||
emitter=None, now=clock, sleep=sleep, rng=rng,
|
||||
)
|
||||
return mw, limiter, gate, transport, sleep, clock
|
||||
|
||||
|
||||
_REQ = ChatRequest(messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
|
||||
class TestSuccessPath:
|
||||
async def test_first_attempt_success_builds_response(self):
|
||||
mw, limiter, gate, transport, sleep, _ = _harness([_src("a")], [_ok("hello")])
|
||||
resp = await mw(_REQ)
|
||||
assert resp.content == "hello"
|
||||
assert resp.source_name == "a" and resp.provider == "openai"
|
||||
assert resp.cache_hit is False and resp.call_id
|
||||
assert (await limiter.source_stats("a")).inflight == 0 # permit 已释放
|
||||
assert await limiter.progress_age_s() < 5.0 # mark_progress 已调用
|
||||
assert sleep.delays == []
|
||||
|
||||
async def test_settle_uses_actual_usage(self):
|
||||
src = _src("a", tpm=1000, est_tokens=400)
|
||||
mw, limiter, *_ = _harness([src], [_ok()])
|
||||
await mw(_REQ)
|
||||
# 预扣 400,实际 15 → settle 后窗口只记 15
|
||||
assert (await limiter.source_stats("a")).tpm_used == 15
|
||||
|
||||
|
||||
class TestRetryAndFailover:
|
||||
async def test_transient_switches_source_then_succeeds(self):
|
||||
mw, _, _, transport, sleep, _ = _harness(
|
||||
[_src("a"), _src("b")], [TransientError("boom"), _ok()]
|
||||
)
|
||||
resp = await mw(_REQ)
|
||||
assert [name for name, _ in transport.calls] == ["a", "b"]
|
||||
assert resp.source_name == "b"
|
||||
assert len(sleep.delays) == 1 # 瞬时错误退避一次
|
||||
|
||||
async def test_each_attempt_gets_fresh_call_id(self):
|
||||
mw, _, _, transport, _, _ = _harness([_src("a")], [TransientError("x"), _ok()])
|
||||
await mw(_REQ)
|
||||
ids = [cid for _, cid in transport.calls]
|
||||
assert len(ids) == 2 and ids[0] != ids[1]
|
||||
|
||||
async def test_max_attempts_is_total_attempts(self):
|
||||
mw, _, _, transport, _, _ = _harness(
|
||||
[_src("a")], [TransientError("1"), TransientError("2"), TransientError("3")],
|
||||
max_attempts=3,
|
||||
)
|
||||
with pytest.raises(AllSourcesExhausted) as ei:
|
||||
await mw(_REQ)
|
||||
assert len(transport.calls) == 3 # 恰 3 次总尝试(含首次)
|
||||
assert ei.value.reason == "retry_exhausted"
|
||||
assert ei.value.per_source_reasons.get("a") == "network_error"
|
||||
|
||||
async def test_backoff_formula_and_retry_after_max(self):
|
||||
# rng=0 → jitter 因子 0.5;第一次退避 = 2*2^0*0.5 = 1.0
|
||||
mw, _, _, _, sleep, _ = _harness([_src("a")], [TransientError("x"), _ok()])
|
||||
await mw(_REQ)
|
||||
assert sleep.delays == [1.0]
|
||||
# Retry-After 提示更大时取提示值
|
||||
mw2, _, _, _, sleep2, _ = _harness(
|
||||
[_src("a")], [TransientError("x", retry_after_s=7.5), _ok()]
|
||||
)
|
||||
await mw2(_REQ)
|
||||
assert sleep2.delays == [7.5]
|
||||
|
||||
async def test_source_dead_switches_immediately_and_force_opens(self):
|
||||
mw, _, gate, transport, sleep, _ = _harness(
|
||||
[_src("a"), _src("b")], [SourceDeadError("401"), _ok()]
|
||||
)
|
||||
resp = await mw(_REQ)
|
||||
assert resp.source_name == "b"
|
||||
assert sleep.delays == [] # 源死亡不退避
|
||||
assert not (await gate.try_enter("a", "w")).allowed # a 已 force_open
|
||||
|
||||
|
||||
class TestNonRetryableOutcomes:
|
||||
async def test_request_rejected_propagates_without_retry(self):
|
||||
exc = RequestRejectedError("400", source_name="a", status_code=400)
|
||||
mw, _, gate, transport, _, _ = _harness([_src("a")], [exc])
|
||||
with pytest.raises(RequestRejectedError):
|
||||
await mw(_REQ)
|
||||
assert len(transport.calls) == 1
|
||||
# 网关已应答 → 记成功,熔断计数未增长
|
||||
assert (await gate.try_enter("a", "w")).allowed
|
||||
|
||||
async def test_result_invalid_records_success_and_propagates(self):
|
||||
mw, limiter, gate, transport, _, _ = _harness(
|
||||
[_src("a")], [ResultInvalidError("bad json", raw_text="{oops")]
|
||||
)
|
||||
with pytest.raises(ResultInvalidError):
|
||||
await mw(_REQ)
|
||||
assert len(transport.calls) == 1 # 坏结果不重试
|
||||
assert (await gate.try_enter("a", "w")).allowed # 熔断记成功
|
||||
assert (await limiter.source_stats("a")).inflight == 0
|
||||
|
||||
|
||||
class TestScopeUnavailable:
|
||||
async def test_all_sources_circuit_open(self):
|
||||
clock = FakeClock()
|
||||
script = [TransientError(str(i)) for i in range(9)]
|
||||
mw, _, gate, _, _, _ = _harness(
|
||||
[_src("a")], script, clock=clock, max_attempts=99
|
||||
)
|
||||
# 3 次失败后 a 开路 → 第 4 次尝试选不到源且 gate_rejections==全部 → CircuitOpen
|
||||
with pytest.raises(CircuitOpenError) as ei:
|
||||
await mw(_REQ)
|
||||
assert ei.value.reason == "circuit_open"
|
||||
assert ei.value.retry_after_s > 0
|
||||
assert ei.value.per_source_reasons # 携逐源原因
|
||||
|
||||
async def test_no_sources_configured(self):
|
||||
mw, *_ = _harness([], [])
|
||||
with pytest.raises(AllSourcesExhausted) as ei:
|
||||
await mw(_REQ)
|
||||
assert ei.value.reason == "no_sources"
|
||||
|
||||
async def test_quota_fail_fast(self):
|
||||
src = _src("a", max_concurrency=1)
|
||||
mw, limiter, _, _, _, _ = _harness([src], [_ok()], quota_full="fail_fast")
|
||||
held = await limiter.try_acquire("a", 0) # 外部占满并发
|
||||
assert held is not None
|
||||
with pytest.raises(AllSourcesExhausted) as ei:
|
||||
await mw(_REQ)
|
||||
assert ei.value.reason == "quota_exhausted"
|
||||
|
||||
async def test_quota_wait_polls_until_slot_frees(self):
|
||||
src = _src("a", max_concurrency=1)
|
||||
clock = FakeClock()
|
||||
limiter = InMemoryLimiter(
|
||||
scope="llm", sources={"a": src}, global_limits=_NO_GLOBAL,
|
||||
lease_ttl_s=100.0, now=clock,
|
||||
)
|
||||
held = await limiter.try_acquire("a", 0)
|
||||
released = {"done": False}
|
||||
|
||||
async def sleep_and_release(seconds):
|
||||
if not released["done"]:
|
||||
released["done"] = True
|
||||
await held.release()
|
||||
|
||||
gate = InMemoryGate(config=_BREAKER, now=clock)
|
||||
transport = FakeTransport([_ok()])
|
||||
mw = RetryMW(
|
||||
scope="llm", sources=[src], selector=RoundRobinSelector(),
|
||||
limiter=limiter, gate=gate, transport=transport,
|
||||
retry=RetryPolicy(max_attempts=3, backoff_base_s=2.0, backoff_max_s=30.0),
|
||||
backpressure=BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.01),
|
||||
quota_full="wait", cooldown_memo=SourceCooldownMemo(now=clock),
|
||||
emitter=None, now=clock, sleep=sleep_and_release, rng=lambda: 0.0,
|
||||
)
|
||||
resp = await mw(_REQ)
|
||||
assert resp.content == "ok" and released["done"]
|
||||
|
||||
|
||||
class TestCancellation:
|
||||
async def test_cancel_mid_flight_releases_permit(self):
|
||||
mw, limiter, _, _, _, _ = _harness([_src("a", max_concurrency=1)], ["hang"])
|
||||
task = asyncio.ensure_future(mw(_REQ))
|
||||
await asyncio.sleep(0.05)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert (await limiter.source_stats("a")).inflight == 0 # finally 释放
|
||||
|
||||
async def test_cancel_probe_releases_probe_lease(self):
|
||||
clock = FakeClock()
|
||||
mw, _, gate, _, _, _ = _harness(
|
||||
[_src("a")],
|
||||
[TransientError("1"), TransientError("2"), TransientError("3"), "hang"],
|
||||
clock=clock, max_attempts=99,
|
||||
)
|
||||
# 三连失败开路
|
||||
with pytest.raises(CircuitOpenError):
|
||||
await mw(_REQ)
|
||||
clock.advance(_BREAKER.cooldown_s + 1)
|
||||
task = asyncio.ensure_future(mw(_REQ)) # 半开探针 → hang
|
||||
await asyncio.sleep(0.05)
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
# 探针租约已归还: 下一 caller 立即拿到探针而非等租约过期
|
||||
nxt = await gate.try_enter("a", "w2")
|
||||
assert nxt.allowed and nxt.is_probe
|
||||
Reference in New Issue
Block a user