Files
PolyGateway/tests/contracts/test_limiter_contract.py
T
iomgaa ab496bb298 feat: let tpm be configured without an est_tokens companion
The gate check forced operators to guess a per-call token size before
they could enable the TPM gate at all; est_tokens is now an optional
tuning override and effective_est_tokens() derives the reservation from
the provider quota. Reservation and settlement already read the same
derived value, so the deposit still nets to zero on both the success
path and the non-dead transient failure path.

The rest of _validate_gates is untouched, and the est_tokens field plus
its EST_TOKENS env key stay put for migration compatibility.
2026-07-30 10:57:40 -04:00

121 lines
5.5 KiB
Python

"""限流后端契约测试(CHS tests/contracts_limiter.py 5 项 + M1 设计 §4.2 补强)。
任何 RateLimiter 后端都必须逐条通过;M2 的 Redis 实现复用本套件。
"""
from polygateway.types import GlobalLimits
from tests.contracts.conftest import make_source
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)
class TestConcurrencyGate:
async def test_concurrency_caps_and_zero_side_effect(self, limiter_factory):
src = make_source(max_concurrency=2)
limiter = limiter_factory([src], _NO_GLOBAL)
p1 = await limiter.try_acquire("s1", 0)
p2 = await limiter.try_acquire("s1", 0)
assert p1 is not None and p2 is not None
# 满员 → None,且失败的 acquire 零副作用
assert await limiter.try_acquire("s1", 0) is None
stats = await limiter.source_stats("s1")
assert stats.inflight == 2
# 释放一个后又能进
await p1.release()
assert (await limiter.source_stats("s1")).inflight == 1
p3 = await limiter.try_acquire("s1", 0)
assert p3 is not None
await p2.release()
await p3.release()
assert (await limiter.source_stats("s1")).inflight == 0
async def test_release_idempotent(self, limiter_factory):
limiter = limiter_factory([make_source(max_concurrency=1)], _NO_GLOBAL)
permit = await limiter.try_acquire("s1", 0)
await permit.release()
await permit.release()
assert (await limiter.source_stats("s1")).inflight == 0
async def test_global_concurrency_across_sources(self, limiter_factory):
sources = [make_source("s1"), make_source("s2")]
limiter = limiter_factory(sources, GlobalLimits(max_concurrency=2, rpm=0, tpm=0))
assert await limiter.try_acquire("s1", 0) is not None
assert await limiter.try_acquire("s2", 0) is not None
assert await limiter.try_acquire("s1", 0) is None # 全局闸挡住第三个
async def test_lease_expiry_reclaims_slot(self, limiter_factory, clock):
"""permit 持有者死亡(未 release)→ 租约过期后并发槽自动回收。"""
limiter = limiter_factory([make_source(max_concurrency=1)], _NO_GLOBAL, lease_ttl_s=30.0)
_leaked = await limiter.try_acquire("s1", 0)
assert await limiter.try_acquire("s1", 0) is None
clock.advance(31.0)
assert await limiter.try_acquire("s1", 0) is not None
class TestRpmGate:
async def test_rpm_not_refunded_by_release(self, limiter_factory):
src = make_source(rpm=3)
limiter = limiter_factory([src], _NO_GLOBAL)
for _ in range(3):
permit = await limiter.try_acquire("s1", 0)
assert permit is not None
await permit.release() # 释放并发,但 RPM 计数不归还
assert await limiter.try_acquire("s1", 0) is None
assert (await limiter.source_stats("s1")).rpm_used == 3
class TestTpmGate:
async def test_prededuct_and_settle_refund(self, limiter_factory):
src = make_source(tpm=1000, est_tokens=400)
limiter = limiter_factory([src], _NO_GLOBAL)
p1 = await limiter.try_acquire("s1", 400)
p2 = await limiter.try_acquire("s1", 400)
assert p1 is not None and p2 is not None
assert await limiter.try_acquire("s1", 400) is None # 预扣用满
# 实际 0 tokens → 全额退款
await p1.settle(0)
await p1.release()
assert (await limiter.source_stats("s1")).tpm_used == 400
# settle 幂等: 第二次调用无副作用
await p1.settle(0)
assert (await limiter.source_stats("s1")).tpm_used == 400
# 多退少补: 实际超预扣则补记
await p2.settle(600)
await p2.release()
assert (await limiter.source_stats("s1")).tpm_used == 600
async def test_settle_equal_to_prededuct_keeps_deposit(self, limiter_factory):
"""预扣量与结算量同为派生值时,双后端都必须留存押金(delta==0)。
这里只锁**后端算术**: 相等的两个数进出,窗口残留量恰为该值。
"调用点是否真的取了派生值"是编排行为,由 tests/unit/test_retry.py
经 RetryMW 端到端覆盖,不在本契约文件重复(否则只是自证同一个入参)。
"""
src = make_source(tpm=6000, est_tokens=0) # 派生值 = max(1, 6000 // 60) = 100
derived = src.effective_est_tokens()
assert derived == 100
limiter = limiter_factory([src], _NO_GLOBAL)
permit = await limiter.try_acquire("s1", derived)
assert permit is not None
await permit.settle(derived)
await permit.release()
assert (await limiter.source_stats("s1")).tpm_used == derived
async def test_failed_acquire_leaves_no_tpm_trace(self, limiter_factory):
src = make_source(tpm=500, est_tokens=400)
limiter = limiter_factory([src], _NO_GLOBAL)
p1 = await limiter.try_acquire("s1", 400)
assert p1 is not None
assert await limiter.try_acquire("s1", 400) is None
assert (await limiter.source_stats("s1")).tpm_used == 400 # 失败尝试零痕迹
class TestProgress:
async def test_progress_marks_fresh(self, limiter_factory, clock):
limiter = limiter_factory([make_source()], _NO_GLOBAL)
assert await limiter.progress_age_s() == float("inf") # 从未出餐
await limiter.mark_progress()
assert await limiter.progress_age_s() < 5.0
clock.advance(42.0)
assert 41.0 < await limiter.progress_age_s() < 43.0