mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat(rate-limit): 实现分层 RPM 限速,支持系统默认/用户/独立Key三级配置
- 新增用户级 rate_limit 字段,支持系统默认/用户自定义/不限制三种模式 - 独立 Key 的 rate_limit 语义调整:null=跟随系统默认,0=不限制,>0=自定义 - 实现 UserRpmLimiter 基于 Redis sliding window 的 RPM 限速引擎 - Pipeline 请求流程集成用户级 RPM 检查 - 管理后台和用户面板新增 RPM 限速配置与实时状态查看 - 系统设置新增全局默认 RPM 配置项 - 迁移脚本回填现有 API Key 的 rate_limit 默认值 - 新增用户/Key RPM 状态监控 API 和前端展示 Closes #231 Co-authored-by: LewisPen <LewisPen@nyadoo.com>
This commit is contained in:
148
tests/services/test_user_rpm_limiter.py
Normal file
148
tests/services/test_user_rpm_limiter.py
Normal file
@@ -0,0 +1,148 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.services.rate_limit.user_rpm_limiter import RpmCheckResult, UserRpmLimiter
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def limiter(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[UserRpmLimiter]:
|
||||
limiter = UserRpmLimiter()
|
||||
limiter._redis = None
|
||||
limiter._memory_counts.clear()
|
||||
await limiter.close()
|
||||
# 阻止 initialize() 连接真实 Redis,确保纯内存模式
|
||||
monkeypatch.setattr(limiter, "initialize", AsyncMock())
|
||||
yield limiter
|
||||
limiter._memory_counts.clear()
|
||||
limiter._redis = None
|
||||
await limiter.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_and_consume_enforces_user_scope_in_memory(
|
||||
limiter: UserRpmLimiter,
|
||||
) -> None:
|
||||
user_key = limiter.get_user_rpm_key("user-1")
|
||||
key_key = limiter.get_key_rpm_key("key-1")
|
||||
|
||||
first = await limiter.check_and_consume(
|
||||
user_rpm_key=user_key,
|
||||
user_rpm_limit=1,
|
||||
key_rpm_key=key_key,
|
||||
key_rpm_limit=0,
|
||||
)
|
||||
second = await limiter.check_and_consume(
|
||||
user_rpm_key=user_key,
|
||||
user_rpm_limit=1,
|
||||
key_rpm_key=key_key,
|
||||
key_rpm_limit=0,
|
||||
)
|
||||
|
||||
assert first == RpmCheckResult(allowed=True, remaining=0)
|
||||
assert second.allowed is False
|
||||
assert second.scope == "user"
|
||||
assert second.limit == 1
|
||||
assert second.remaining == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_and_consume_enforces_key_scope_in_memory(
|
||||
limiter: UserRpmLimiter,
|
||||
) -> None:
|
||||
user_key = limiter.get_user_rpm_key("user-1")
|
||||
key_key = limiter.get_key_rpm_key("key-1")
|
||||
|
||||
first = await limiter.check_and_consume(
|
||||
user_rpm_key=user_key,
|
||||
user_rpm_limit=3,
|
||||
key_rpm_key=key_key,
|
||||
key_rpm_limit=1,
|
||||
)
|
||||
second = await limiter.check_and_consume(
|
||||
user_rpm_key=user_key,
|
||||
user_rpm_limit=3,
|
||||
key_rpm_key=key_key,
|
||||
key_rpm_limit=1,
|
||||
)
|
||||
|
||||
assert first.allowed is True
|
||||
assert second.allowed is False
|
||||
assert second.scope == "key"
|
||||
assert second.limit == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_and_consume_redis_failure_falls_back_to_memory_when_fail_close(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
limiter: UserRpmLimiter,
|
||||
) -> None:
|
||||
fake_redis = MagicMock()
|
||||
fake_redis.eval = AsyncMock(side_effect=RuntimeError("redis down"))
|
||||
limiter._redis = fake_redis
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.services.rate_limit.user_rpm_limiter.config.rate_limit_fail_open", False
|
||||
)
|
||||
|
||||
user_key = limiter.get_user_rpm_key("user-1")
|
||||
key_key = limiter.get_key_rpm_key("key-1")
|
||||
|
||||
first = await limiter.check_and_consume(
|
||||
user_rpm_key=user_key,
|
||||
user_rpm_limit=1,
|
||||
key_rpm_key=key_key,
|
||||
key_rpm_limit=0,
|
||||
)
|
||||
second = await limiter.check_and_consume(
|
||||
user_rpm_key=user_key,
|
||||
user_rpm_limit=1,
|
||||
key_rpm_key=key_key,
|
||||
key_rpm_limit=0,
|
||||
)
|
||||
|
||||
assert first.allowed is True
|
||||
assert second.allowed is False
|
||||
assert second.scope == "user"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_and_consume_redis_failure_fail_open_allows_request(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
limiter: UserRpmLimiter,
|
||||
) -> None:
|
||||
fake_redis = MagicMock()
|
||||
fake_redis.eval = AsyncMock(side_effect=RuntimeError("redis down"))
|
||||
limiter._redis = fake_redis
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.services.rate_limit.user_rpm_limiter.config.rate_limit_fail_open", True
|
||||
)
|
||||
|
||||
result = await limiter.check_and_consume(
|
||||
user_rpm_key=limiter.get_user_rpm_key("user-1"),
|
||||
user_rpm_limit=1,
|
||||
key_rpm_key=limiter.get_key_rpm_key("key-1"),
|
||||
key_rpm_limit=0,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_and_consume_skips_when_all_limits_are_zero(
|
||||
limiter: UserRpmLimiter,
|
||||
) -> None:
|
||||
result = await limiter.check_and_consume(
|
||||
user_rpm_key=limiter.get_user_rpm_key("user-1"),
|
||||
user_rpm_limit=0,
|
||||
key_rpm_key=limiter.get_key_rpm_key("key-1"),
|
||||
key_rpm_limit=0,
|
||||
)
|
||||
|
||||
assert result.allowed is True
|
||||
assert result.scope is None
|
||||
assert limiter._memory_counts == {}
|
||||
Reference in New Issue
Block a user