mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40: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:
120
tests/api/test_monitoring_user_rate_limit.py
Normal file
120
tests/api/test_monitoring_user_rate_limit.py
Normal file
@@ -0,0 +1,120 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from src.api.monitoring.user import UserRateLimitStatusAdapter
|
||||
|
||||
|
||||
def _build_query_with_keys(keys: list[object]) -> MagicMock:
|
||||
query = MagicMock()
|
||||
query.filter.return_value.order_by.return_value.all.return_value = keys
|
||||
return query
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_status_adapter_reports_user_and_key_layers(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
now = datetime(2026, 3, 13, 12, 0, tzinfo=timezone.utc)
|
||||
db = MagicMock()
|
||||
user = SimpleNamespace(id="user-1", rate_limit=None)
|
||||
key = SimpleNamespace(id="key-1", name="Primary", is_standalone=False, rate_limit=10)
|
||||
standalone = SimpleNamespace(
|
||||
id="skey-1", name="Standalone", is_standalone=True, rate_limit=None
|
||||
)
|
||||
|
||||
db.query.return_value = _build_query_with_keys([key, standalone])
|
||||
|
||||
limiter = MagicMock()
|
||||
limiter.bucket_seconds = 60
|
||||
limiter.get_reset_at.return_value = now
|
||||
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:bucket"
|
||||
limiter.get_standalone_rpm_key.return_value = "rpm:ukey:skey-1:bucket"
|
||||
limiter.get_key_rpm_key.side_effect = lambda key_id: f"rpm:key:{key_id}:bucket"
|
||||
|
||||
async def _get_scope_count(scope_key: str) -> int:
|
||||
counts = {
|
||||
"rpm:user:user-1:bucket": 55,
|
||||
"rpm:key:key-1:bucket": 7,
|
||||
"rpm:ukey:skey-1:bucket": 12,
|
||||
}
|
||||
return counts[scope_key]
|
||||
|
||||
limiter.get_scope_count = AsyncMock(side_effect=_get_scope_count)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.monitoring.user.get_user_rpm_limiter",
|
||||
AsyncMock(return_value=limiter),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.monitoring.user.SystemConfigService.get_config",
|
||||
lambda *_a, **_k: 60,
|
||||
)
|
||||
|
||||
context = SimpleNamespace(db=db, user=user)
|
||||
result = await UserRateLimitStatusAdapter().handle(context)
|
||||
|
||||
assert result["user_id"] == "user-1"
|
||||
assert result["api_keys"][0] == {
|
||||
"api_key_name": "Primary",
|
||||
"limit": 10,
|
||||
"remaining": 3,
|
||||
"scope": "key",
|
||||
"reset_time": now.isoformat(),
|
||||
"window": "60s",
|
||||
"user_limit": 60,
|
||||
"user_remaining": 5,
|
||||
"key_limit": 10,
|
||||
"key_remaining": 3,
|
||||
}
|
||||
assert result["api_keys"][1] == {
|
||||
"api_key_name": "Standalone",
|
||||
"limit": 60,
|
||||
"remaining": 48,
|
||||
"scope": "user",
|
||||
"reset_time": now.isoformat(),
|
||||
"window": "60s",
|
||||
"user_limit": 60,
|
||||
"user_remaining": 48,
|
||||
"key_limit": None,
|
||||
"key_remaining": None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_status_adapter_reports_unlimited_key_without_counts(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
db = MagicMock()
|
||||
user = SimpleNamespace(id="user-1", rate_limit=0)
|
||||
key = SimpleNamespace(id="key-1", name="Unlimited", is_standalone=False, rate_limit=0)
|
||||
|
||||
db.query.return_value = _build_query_with_keys([key])
|
||||
|
||||
limiter = MagicMock()
|
||||
limiter.bucket_seconds = 60
|
||||
limiter.get_reset_at.return_value = datetime.now(timezone.utc)
|
||||
limiter.get_user_rpm_key.return_value = "rpm:user:user-1:bucket"
|
||||
limiter.get_key_rpm_key.return_value = "rpm:key:key-1:bucket"
|
||||
limiter.get_scope_count = AsyncMock()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.api.monitoring.user.get_user_rpm_limiter",
|
||||
AsyncMock(return_value=limiter),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"src.api.monitoring.user.SystemConfigService.get_config",
|
||||
lambda *_a, **_k: 60,
|
||||
)
|
||||
|
||||
context = SimpleNamespace(db=db, user=user)
|
||||
result = await UserRateLimitStatusAdapter().handle(context)
|
||||
|
||||
assert result["api_keys"][0]["limit"] is None
|
||||
assert result["api_keys"][0]["remaining"] is None
|
||||
assert result["api_keys"][0]["scope"] is None
|
||||
limiter.get_scope_count.assert_not_awaited()
|
||||
Reference in New Issue
Block a user