mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
1d72a8f9c1
Close #233 Co-authored-by: AAEE86 <ppk0227@hotmail.com> - cli_monitor_mixin: 引入 STREAM_IDLE_TIMEOUT_SECONDS（可通过环境变量配置）， 流传输开始后若超出空闲窗口无新 chunk 则提前取消并返回 504，避免长时间挂起 - stream_context: 新增 managed_recorded_bodies 上下文管理器，确保 chunks 在 telemetry 完成后及时释放；stream_telemetry 使用该接口统一管理 response body 构建 - health endpoint: 将状态聚合改为 GROUP BY 直接统计，事件列表按 api_format 单独查询，避免单次 limit 拉取大量记录导致的遗漏与性能问题；同时过滤不活跃 provider/endpoint，与公开健康接口保持一致 - endpoint health service: 修正时间线数据按 endpoint_id 而非 key_id 聚合 - token_bucket: 引入 max_buckets/bucket_expiry 上限与定时清理，防止内存无限增长； 修复 refill_rate=0 时 get_reset_time 除零异常；新增 _is_unlimited_rate_limit 判断 - maintenance_scheduler: 调整清理顺序（先删整行再按窗口清理），新增 newer_than 边界参数，避免同一行在同一轮中被重复改写 - sync_execute: 新增 create_pending_usage 开关，允许已预创建记录的调用方跳过重复创建 - quota_reader / provider_ops balance: 小幅修复与健壮性提升 - Dockerfile: 添加 MALLOC_ARENA_MAX=2 环境变量以降低 gunicorn worker RSS - 补充相关测试覆盖
112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime, timezone
|
|
|
|
import pytest
|
|
|
|
import src.plugins.rate_limit.token_bucket as token_bucket_module
|
|
from src.plugins.rate_limit.token_bucket import RedisTokenBucketBackend, TokenBucketStrategy
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_token_bucket_cleans_up_expired_buckets(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
|
|
strategy = TokenBucketStrategy()
|
|
strategy.configure({"bucket_expiry": 1, "cleanup_interval": 0})
|
|
|
|
await strategy.check_limit("api_key:stale")
|
|
strategy.buckets["api_key:stale"].last_access_time -= 3600
|
|
|
|
await strategy.check_limit("api_key:fresh")
|
|
|
|
assert "api_key:stale" not in strategy.buckets
|
|
assert "api_key:fresh" in strategy.buckets
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_token_bucket_reconfigures_existing_bucket_when_rate_limit_changes(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
|
|
strategy = TokenBucketStrategy()
|
|
|
|
await strategy.check_limit("user:42", rate_limit=120)
|
|
bucket = strategy.buckets["user:42"]
|
|
bucket.tokens = 90
|
|
|
|
await strategy.check_limit("user:42", rate_limit=30)
|
|
|
|
updated_bucket = strategy.buckets["user:42"]
|
|
assert updated_bucket.capacity == 30
|
|
assert updated_bucket.refill_rate == 0.5
|
|
assert updated_bucket.tokens <= 30
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_token_bucket_treats_non_positive_dynamic_rate_limit_as_unlimited(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
|
|
strategy = TokenBucketStrategy()
|
|
|
|
result = await strategy.check_limit("public_ip:test", rate_limit=0)
|
|
consumed = await strategy.consume("public_ip:test", amount=1, rate_limit=0)
|
|
|
|
assert result.allowed is True
|
|
assert consumed is True
|
|
assert "public_ip:test" not in strategy.buckets
|
|
|
|
|
|
class _FakeRedisClient:
|
|
async def hmget(self, _key: str, *_fields: str) -> list[None]:
|
|
return [None, None]
|
|
|
|
def register_script(self, _script: str): # type: ignore[no-untyped-def]
|
|
async def _runner(*args, **kwargs): # type: ignore[no-untyped-def]
|
|
return [1, 0, 0]
|
|
|
|
return _runner
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_token_bucket_retries_redis_backend_probe_after_initial_miss(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("RATE_LIMIT_BACKEND", "auto")
|
|
strategy = TokenBucketStrategy()
|
|
strategy._redis_retry_interval = 0
|
|
|
|
fake_redis = _FakeRedisClient()
|
|
calls = {"count": 0}
|
|
|
|
def _fake_get_redis_client_sync(): # type: ignore[no-untyped-def]
|
|
calls["count"] += 1
|
|
if calls["count"] == 1:
|
|
return None
|
|
return fake_redis
|
|
|
|
monkeypatch.setattr(
|
|
token_bucket_module,
|
|
"get_redis_client_sync",
|
|
_fake_get_redis_client_sync,
|
|
)
|
|
|
|
await strategy.check_limit("public_ip:first")
|
|
assert strategy._redis_backend is None
|
|
|
|
await strategy.check_limit("public_ip:second")
|
|
assert strategy._redis_backend is not None
|
|
assert calls["count"] == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redis_token_bucket_missing_bucket_reports_reset_now() -> None:
|
|
backend = RedisTokenBucketBackend(_FakeRedisClient())
|
|
|
|
result = await backend.peek("public_ip:test", capacity=60, refill_rate=1.0, amount=1)
|
|
|
|
assert result.allowed is True
|
|
assert result.remaining == 60
|
|
assert result.reset_at is not None
|
|
assert abs((result.reset_at - datetime.now(timezone.utc)).total_seconds()) < 2
|