219 lines
7.3 KiB
Python
219 lines
7.3 KiB
Python
"""限频器单元测试
|
|
|
|
hexi_core 包 __init__ 会触发 NoneBot 运行时初始化(get_plugin_config),
|
|
无法裸导入,故用 importlib 按文件路径直接加载 rate_limit 模块
|
|
(该模块只依赖 nonebot.log,可独立于 NoneBot 运行)。
|
|
"""
|
|
|
|
import importlib.util
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
_MODULE_PATH = (
|
|
Path(__file__).resolve().parents[1]
|
|
/ "hexi"
|
|
/ "core"
|
|
/ "rate_limit.py"
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def rl():
|
|
"""以独立模块名加载 rate_limit.py,避免与包 __init__ 的依赖链冲突"""
|
|
spec = importlib.util.spec_from_file_location("hexi_rate_limit_test", _MODULE_PATH)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
yield module
|
|
sys.modules.pop(spec.name, None)
|
|
|
|
|
|
def test_capacity_limits_burst(rl):
|
|
limiter = rl.RateLimiter(rate=10, capacity=3)
|
|
assert limiter.try_acquire() is True
|
|
assert limiter.try_acquire() is True
|
|
assert limiter.try_acquire() is True
|
|
assert limiter.try_acquire() is False # 突发额度耗尽
|
|
|
|
|
|
def test_default_capacity_at_least_one(rl):
|
|
assert rl.RateLimiter(rate=5)._capacity == 5 # 默认容量 = rate
|
|
assert rl.RateLimiter(rate=0.5)._capacity == 1 # rate < 1 时保底 1
|
|
|
|
|
|
def test_invalid_params(rl):
|
|
with pytest.raises(ValueError):
|
|
rl.RateLimiter(rate=0)
|
|
with pytest.raises(ValueError):
|
|
rl.RateLimiter(rate=1, capacity=0)
|
|
|
|
|
|
def test_lazy_refill_by_backdating(rl):
|
|
"""把 last_refill 回拨模拟时间流逝,验证惰性补充与容量上限"""
|
|
limiter = rl.RateLimiter(rate=2, capacity=2)
|
|
assert limiter.try_acquire() is True
|
|
assert limiter.try_acquire() is True
|
|
assert limiter.try_acquire() is False # 桶空
|
|
limiter._last_refill = time.monotonic() - 1.0 # 假装 1 秒前补充过
|
|
assert limiter.try_acquire() is True # 1 秒补 2 个
|
|
assert limiter.try_acquire() is True # 桶已补满,再取 1 个
|
|
assert limiter.try_acquire() is False # 再次耗尽
|
|
|
|
|
|
def test_capacity_caps_refill(rl):
|
|
limiter = rl.RateLimiter(rate=10, capacity=2)
|
|
limiter._last_refill = time.monotonic() - 10.0 # 10 秒应补 100 个,但桶上限 2
|
|
assert limiter.try_acquire() is True
|
|
assert limiter.try_acquire() is True
|
|
assert limiter.try_acquire() is False
|
|
|
|
|
|
def test_acquire_sync_waits_for_token(rl):
|
|
limiter = rl.RateLimiter(rate=1, capacity=1)
|
|
assert limiter.try_acquire() is True # 消耗唯一令牌
|
|
start = time.monotonic()
|
|
assert limiter.acquire_sync(timeout=5) is True # 约 1 秒后补出令牌
|
|
assert 0.6 <= time.monotonic() - start < 3
|
|
|
|
|
|
def test_acquire_sync_timeout(rl):
|
|
limiter = rl.RateLimiter(rate=0.1, capacity=1) # 10 秒才补 1 个
|
|
assert limiter.try_acquire() is True
|
|
start = time.monotonic()
|
|
assert limiter.acquire_sync(timeout=0.2) is False
|
|
assert time.monotonic() - start < 2
|
|
|
|
|
|
def test_acquire_sync_zero_timeout_fails_fast(rl):
|
|
limiter = rl.RateLimiter(rate=0.001, capacity=1)
|
|
assert limiter.try_acquire() is True
|
|
assert limiter.acquire_sync(timeout=0) is False # 立即失败,不等
|
|
|
|
|
|
async def test_acquire_async_waits_for_token(rl):
|
|
limiter = rl.RateLimiter(rate=1, capacity=1)
|
|
assert limiter.try_acquire() is True
|
|
start = time.monotonic()
|
|
assert await limiter.acquire(timeout=5) is True
|
|
assert 0.6 <= time.monotonic() - start < 3
|
|
|
|
|
|
async def test_acquire_async_timeout(rl):
|
|
limiter = rl.RateLimiter(rate=0.1, capacity=1)
|
|
assert limiter.try_acquire() is True
|
|
assert await limiter.acquire(timeout=0.2) is False
|
|
|
|
|
|
def test_get_limiter_shared_by_name(rl):
|
|
a = rl.get_limiter("shared_test", rate=1)
|
|
b = rl.get_limiter("shared_test", rate=1)
|
|
assert a is b # 同名复用同一实例
|
|
assert rl.get_limiter("other_test", rate=1) is not a
|
|
|
|
|
|
async def test_acquire_helper(rl):
|
|
# 首次取令牌成功,立即再取失败(不等待)
|
|
assert await rl.acquire("helper_failfast", rate=0.001, wait=False) is True
|
|
assert await rl.acquire("helper_failfast", rate=0.001, wait=False) is False
|
|
# timeout=0 等模式同样立即失败
|
|
assert await rl.acquire("helper_failfast", rate=0.001, timeout=0) is False
|
|
|
|
|
|
def test_acquire_sync_helper(rl):
|
|
assert rl.acquire_sync("helper_sync", rate=0.001, wait=False) is True
|
|
assert rl.acquire_sync("helper_sync", rate=0.001, wait=False) is False
|
|
|
|
|
|
def test_is_ratelimited(rl):
|
|
assert rl.is_ratelimited(403) is True
|
|
assert rl.is_ratelimited(429) is True
|
|
assert rl.is_ratelimited(200) is False
|
|
assert rl.is_ratelimited(404) is False
|
|
assert rl.is_ratelimited(None) is False
|
|
|
|
|
|
# ---------- steam_info 场景模拟 ----------
|
|
|
|
|
|
class _FakeTime:
|
|
"""可手动推进的 time 替身:验证限频时序无需真实等待"""
|
|
|
|
def __init__(self):
|
|
self._t = 1000.0
|
|
|
|
def monotonic(self) -> float:
|
|
return self._t
|
|
|
|
def advance(self, secs: float) -> None:
|
|
self._t += secs
|
|
|
|
|
|
class _FakeAsyncio:
|
|
"""asyncio 替身:sleep 立即完成并推进回拨时钟(模拟"等待配额"过程)"""
|
|
|
|
def __init__(self, fake_time: _FakeTime):
|
|
self._ft = fake_time
|
|
|
|
async def sleep(self, secs: float) -> None:
|
|
self._ft.advance(secs)
|
|
|
|
|
|
async def test_steam_quota_simulation(rl):
|
|
"""模拟 steam_info 场景:播报 1/min(等待)+ check 高频插入(快速失败)
|
|
|
|
参数与 steam.py 的 STEAM_API_RATE/STEAM_API_CAPACITY 保持一致
|
|
(rate=1/2 → 相邻请求间隔 >= 2s,双 key 合计 150 req/5min)。
|
|
回拨时钟跑 600 模拟秒,验证:
|
|
- 相邻成功请求间隔 >= 2s(短窗口突发被杜绝)
|
|
- 等待模式的播报在 check 抢配额下仍全部成功
|
|
- 快速失败的 check 在配额耗尽时被拒绝
|
|
"""
|
|
fake = _FakeTime()
|
|
real_time, real_asyncio = rl.time, rl.asyncio
|
|
rl.time = fake
|
|
rl.asyncio = _FakeAsyncio(fake)
|
|
try:
|
|
rate, capacity = 1 / 2, 1
|
|
gaps: list[float] = []
|
|
last_ok = fake.monotonic()
|
|
broadcast_ok = 0
|
|
check_ok = 0
|
|
check_rejected = 0
|
|
|
|
next_broadcast, next_check = 60.0, 2.0
|
|
t = 0.0
|
|
while t <= 600:
|
|
if t >= next_broadcast:
|
|
next_broadcast += 60
|
|
if await rl.acquire(
|
|
"steam_api", rate=rate, capacity=capacity, timeout=60
|
|
):
|
|
gaps.append(fake.monotonic() - last_ok)
|
|
last_ok = fake.monotonic()
|
|
broadcast_ok += 1
|
|
if t >= next_check:
|
|
next_check += 2
|
|
if await rl.acquire(
|
|
"steam_api", rate=rate, capacity=capacity, wait=False
|
|
):
|
|
gaps.append(fake.monotonic() - last_ok)
|
|
last_ok = fake.monotonic()
|
|
check_ok += 1
|
|
else:
|
|
check_rejected += 1
|
|
fake.advance(1.0)
|
|
t += 1.0
|
|
|
|
# 等待模式的播报 10 次全部成功(最多等 60s,不会被 check 挤掉)
|
|
assert broadcast_ok == 10
|
|
# 快速失败的 check 在配额耗尽时被拒(check 300 次尝试 vs 配额 300)
|
|
assert check_rejected > 0
|
|
# 相邻成功请求间隔 >= 2s(首请求无前驱,不计入)
|
|
assert check_ok + broadcast_ok <= 300
|
|
assert min(gaps[1:]) >= 1.9
|
|
finally:
|
|
rl.time, rl.asyncio = real_time, real_asyncio
|