"""限频器单元测试 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