mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +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>
121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
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()
|