mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
f92b0943b5
- 新增用户级 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>
149 lines
4.1 KiB
Python
149 lines
4.1 KiB
Python
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 == {}
|