285 lines
8.3 KiB
Python
285 lines
8.3 KiB
Python
"""命令冷却单元测试
|
|
|
|
hexi_core 包 __init__ 会触发 NoneBot 运行时初始化(get_plugin_config),
|
|
无法裸导入,故用 importlib 以真实模块名加载 rate_limit / cooldown 并注册进
|
|
sys.modules —— cooldown 内部对 rate_limit 的 import 会直接命中缓存,
|
|
不触发包 __init__ 的导入链。
|
|
|
|
NoneBot 的依赖注入会对参数做类型校验,故 guard 测试使用真实
|
|
GroupMessageEvent / Matcher 实例,并通过 ContextVar 提供 finish 所需环境。
|
|
"""
|
|
|
|
import importlib.util
|
|
import re
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from nonebot.adapters.onebot.v11 import (
|
|
Bot,
|
|
GroupMessageEvent,
|
|
Message,
|
|
PrivateMessageEvent,
|
|
)
|
|
from nonebot.adapters.onebot.v11.event import Sender
|
|
from nonebot.dependencies import Dependent
|
|
from nonebot.exception import FinishedException
|
|
from nonebot.internal.matcher.matcher import current_bot, current_event
|
|
from nonebot.matcher import Matcher
|
|
|
|
_BASE = (
|
|
Path(__file__).resolve().parents[1]
|
|
/ "hexi"
|
|
/ "core"
|
|
)
|
|
_RL_NAME = "hexi.core.rate_limit"
|
|
_CD_NAME = "hexi.core.cooldown"
|
|
|
|
|
|
def _load_module(name: str, path: Path):
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def cd():
|
|
_load_module(_RL_NAME, _BASE / "rate_limit.py")
|
|
module = _load_module(_CD_NAME, _BASE / "cooldown.py")
|
|
yield module
|
|
sys.modules.pop(_CD_NAME, None)
|
|
sys.modules.pop(_RL_NAME, None)
|
|
|
|
|
|
def _event(uid: int = 123456, group_id: int = 2) -> GroupMessageEvent:
|
|
return GroupMessageEvent(
|
|
time=0,
|
|
self_id=1,
|
|
post_type="message",
|
|
message_type="group",
|
|
sub_type="normal",
|
|
user_id=uid,
|
|
message_id=1,
|
|
message=Message("hi"),
|
|
raw_message="hi",
|
|
font=0,
|
|
sender=Sender(user_id=uid),
|
|
group_id=group_id,
|
|
)
|
|
|
|
|
|
def _private_event(uid: int = 123456) -> PrivateMessageEvent:
|
|
return PrivateMessageEvent(
|
|
time=0,
|
|
self_id=1,
|
|
post_type="message",
|
|
message_type="private",
|
|
sub_type="friend",
|
|
user_id=uid,
|
|
message_id=1,
|
|
message=Message("hi"),
|
|
raw_message="hi",
|
|
font=0,
|
|
sender=Sender(user_id=uid),
|
|
)
|
|
|
|
|
|
class FakeBot:
|
|
"""极简 bot:send 记录消息"""
|
|
|
|
def __init__(self):
|
|
self.sent = []
|
|
|
|
async def send(self, event, message, **kwargs):
|
|
self.sent.append(str(message))
|
|
return 0
|
|
|
|
|
|
def _make_matcher(cd, seconds=10, **kwargs):
|
|
mcls = Matcher.new("message")
|
|
return cd.cooldown(seconds, **kwargs)(mcls)
|
|
|
|
|
|
async def _call_guard(mcls, bot, event):
|
|
"""模拟 NoneBot 调用 guard handler(Dependent 全关键字调用)"""
|
|
m = mcls()
|
|
dep = mcls.handlers[0]
|
|
tok_bot = current_bot.set(bot)
|
|
tok_event = current_event.set(event)
|
|
try:
|
|
return await dep(
|
|
matcher=m,
|
|
bot=bot,
|
|
event=event,
|
|
state={},
|
|
stack=None,
|
|
dependency_cache=None,
|
|
)
|
|
finally:
|
|
current_bot.reset(tok_bot)
|
|
current_event.reset(tok_event)
|
|
|
|
|
|
# ---------- Cooldown 类 ----------
|
|
|
|
|
|
def test_trigger_and_query(cd):
|
|
c = cd.Cooldown(3)
|
|
assert c.try_trigger("u1") is True
|
|
assert c.in_cd("u1") is True
|
|
assert 0 < c.remaining("u1") <= 3.01 # 浮点容差
|
|
assert c.try_trigger("u1") is False # CD 中再次触发失败
|
|
|
|
|
|
def test_instances_isolated(cd):
|
|
c1 = cd.Cooldown(3)
|
|
c2 = cd.Cooldown(3)
|
|
assert c1.try_trigger("u1") is True
|
|
assert c2.try_trigger("u1") is True # 无 name 的不同实例互不影响
|
|
|
|
|
|
def test_remaining_after_window(cd):
|
|
c = cd.Cooldown(3)
|
|
assert c.try_trigger("u1") is True
|
|
limiter = c._limiter("u1")
|
|
limiter._last_refill = time.monotonic() - 3.1 # 模拟 3.1 秒后
|
|
assert c.in_cd("u1") is False
|
|
assert c.remaining("u1") == 0
|
|
assert c.try_trigger("u1") is True # 冷却结束可再次触发
|
|
|
|
|
|
def test_independent_per_key(cd):
|
|
c = cd.Cooldown(3)
|
|
assert c.try_trigger("u1") is True
|
|
assert c.in_cd("u2") is False # 其他用户不受影响
|
|
assert c.try_trigger("u2") is True
|
|
|
|
|
|
def test_invalid_seconds(cd):
|
|
with pytest.raises(ValueError):
|
|
cd.Cooldown(0)
|
|
|
|
|
|
def test_rate_limiter_remaining_secs(cd):
|
|
rl = sys.modules[_RL_NAME]
|
|
limiter = rl.RateLimiter(rate=1, capacity=1)
|
|
assert limiter.remaining_secs() == 0 # 满桶无等待
|
|
assert limiter.try_acquire() is True
|
|
assert 0 < limiter.remaining_secs() <= 1
|
|
limiter._last_refill = time.monotonic() - 2.0 # 2 秒后已补满
|
|
assert limiter.remaining_secs() == 0
|
|
|
|
|
|
# ---------- cooldown 装饰器 ----------
|
|
|
|
|
|
def test_guard_registered_first(cd):
|
|
mcls = _make_matcher(cd, 10)
|
|
assert len(mcls.handlers) == 1
|
|
|
|
async def real_handler(bot: Bot, ev: GroupMessageEvent):
|
|
pass
|
|
|
|
mcls.handlers.append(
|
|
Dependent.parse(call=real_handler, allow_types=Matcher.HANDLER_PARAM_TYPES)
|
|
) # 模拟 @matcher.handle() 追加真 handler
|
|
assert mcls.handlers[0].call is not real_handler # guard 仍在最前
|
|
|
|
|
|
async def test_guard_passes_and_blocks_with_hint(cd):
|
|
mcls = _make_matcher(cd, 10, hint="慢点,{secs} 秒后")
|
|
bot = FakeBot()
|
|
ev = _event(1001)
|
|
|
|
# 首次触发:放行,不发提示
|
|
assert await _call_guard(mcls, bot, ev) is None
|
|
assert bot.sent == []
|
|
|
|
# 冷却中触发:finish 提示并终止
|
|
with pytest.raises(FinishedException):
|
|
await _call_guard(mcls, bot, ev)
|
|
assert len(bot.sent) == 1
|
|
msg = bot.sent[0]
|
|
assert "慢点," in msg and "秒后" in msg
|
|
secs = int(re.search(r"\d+", msg).group())
|
|
assert 1 <= secs <= 10 # 剩余秒数向上取整
|
|
|
|
|
|
async def test_guard_blocks_only_same_user(cd):
|
|
mcls = _make_matcher(cd, 10)
|
|
bot = FakeBot()
|
|
await _call_guard(mcls, bot, _event(1001))
|
|
await _call_guard(mcls, bot, _event(1002)) # 其他用户不受影响
|
|
assert bot.sent == []
|
|
with pytest.raises(FinishedException):
|
|
await _call_guard(mcls, bot, _event(1001))
|
|
assert len(bot.sent) == 1
|
|
|
|
|
|
async def test_guard_custom_key_shares_cd(cd):
|
|
mcls = _make_matcher(cd, 10, key=lambda e: "global")
|
|
bot = FakeBot()
|
|
await _call_guard(mcls, bot, _event(1001))
|
|
# 自定义 key 下不同用户共享同一冷却
|
|
with pytest.raises(FinishedException):
|
|
await _call_guard(mcls, bot, _event(1002))
|
|
assert len(bot.sent) == 1
|
|
|
|
|
|
# ---------- scope 作用域 ----------
|
|
|
|
|
|
async def test_scope_group_shared_by_group(cd):
|
|
mcls = _make_matcher(cd, 10, scope="group")
|
|
bot = FakeBot()
|
|
await _call_guard(mcls, bot, _event(1001, group_id=2))
|
|
# 同群其他用户也处于 CD 中
|
|
with pytest.raises(FinishedException):
|
|
await _call_guard(mcls, bot, _event(1002, group_id=2))
|
|
assert len(bot.sent) == 1
|
|
|
|
|
|
async def test_scope_group_isolated_between_groups(cd):
|
|
mcls = _make_matcher(cd, 10, scope="group")
|
|
bot = FakeBot()
|
|
await _call_guard(mcls, bot, _event(1001, group_id=2))
|
|
# 不同群的用户不受影响
|
|
await _call_guard(mcls, bot, _event(1001, group_id=3))
|
|
assert bot.sent == []
|
|
|
|
|
|
async def test_scope_group_private_falls_back_to_user(cd):
|
|
mcls = _make_matcher(cd, 10, scope="group")
|
|
bot = FakeBot()
|
|
await _call_guard(mcls, bot, _event(1001, group_id=2))
|
|
# 私聊没有群,退化为按用户:不同用户私聊互不影响
|
|
await _call_guard(mcls, bot, _private_event(1002))
|
|
# 同一用户私聊与群聊各自独立
|
|
await _call_guard(mcls, bot, _private_event(1001))
|
|
assert bot.sent == []
|
|
with pytest.raises(FinishedException):
|
|
await _call_guard(mcls, bot, _private_event(1001))
|
|
assert len(bot.sent) == 1
|
|
|
|
|
|
async def test_scope_global_shared(cd):
|
|
mcls = _make_matcher(cd, 10, scope="global")
|
|
bot = FakeBot()
|
|
await _call_guard(mcls, bot, _event(1001, group_id=2))
|
|
# 任何用户、任何群都共享
|
|
with pytest.raises(FinishedException):
|
|
await _call_guard(mcls, bot, _event(1002, group_id=3))
|
|
assert len(bot.sent) == 1
|
|
|
|
|
|
async def test_key_takes_priority_over_scope(cd):
|
|
mcls = _make_matcher(cd, 10, scope="global", key=lambda e: e.get_user_id())
|
|
bot = FakeBot()
|
|
await _call_guard(mcls, bot, _event(1001))
|
|
# 传了 key 则忽略 scope:不同用户互不影响
|
|
await _call_guard(mcls, bot, _event(1002))
|
|
assert bot.sent == []
|