106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
"""middleware/rate_limit.py 请求限流中间件单元测试。"""
|
||
|
||
from contextlib import asynccontextmanager
|
||
|
||
import fakeredis.aioredis
|
||
from fastapi import FastAPI
|
||
from fastapi.testclient import TestClient
|
||
|
||
import middleware.rate_limit as rate_limit_module
|
||
from middleware.rate_limit import rate_limit_middleware
|
||
|
||
|
||
def _make_app():
|
||
@asynccontextmanager
|
||
async def lifespan(app):
|
||
# 在应用运行的事件循环内创建 FakeRedis，避免跨循环绑定问题
|
||
app.state.redis = fakeredis.aioredis.FakeRedis(decode_responses=True)
|
||
yield
|
||
await app.state.redis.aclose()
|
||
|
||
app = FastAPI(lifespan=lifespan)
|
||
|
||
@app.get("/")
|
||
async def home():
|
||
return {"ok": True}
|
||
|
||
app.middleware("http")(rate_limit_middleware)
|
||
|
||
return app
|
||
|
||
|
||
def test_rate_limit_disabled_passes(monkeypatch):
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_ENABLED", False)
|
||
with TestClient(_make_app()) as client:
|
||
for _ in range(5):
|
||
r = client.get("/")
|
||
assert r.status_code == 200
|
||
assert r.json()["ok"] is True
|
||
|
||
|
||
async def _fake_sliding_window(redis, current_key, prev_key, window_seconds, window_index, now):
|
||
"""fakeredis 不支持 Lua eval，用等价实现打桩滑动窗口计数/加权语义。"""
|
||
cur = await redis.incr(current_key)
|
||
if cur == 1:
|
||
await redis.expire(current_key, window_seconds)
|
||
prev = int(await redis.get(prev_key) or 0)
|
||
elapsed = now - window_index * window_seconds
|
||
f = elapsed / window_seconds
|
||
f = max(0.0, min(1.0, f))
|
||
return prev * (1 - f) + cur
|
||
|
||
|
||
def test_rate_limit_exceeds_returns_business_error(monkeypatch):
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_ENABLED", True)
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_LIMIT", 3)
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_WINDOW_SECONDS", 60)
|
||
monkeypatch.setattr(rate_limit_module, "sliding_window_incr", _fake_sliding_window)
|
||
with TestClient(_make_app()) as client:
|
||
for _ in range(3):
|
||
r = client.get("/")
|
||
assert r.status_code == 200
|
||
assert r.json()["ok"] is True
|
||
r = client.get("/")
|
||
assert r.status_code == 200
|
||
assert r.json()["code"] == 1
|
||
assert "过于频繁" in r.json()["msg"]
|
||
|
||
|
||
def test_rate_limit_window_accepts_after_reset(monkeypatch):
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_ENABLED", True)
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_LIMIT", 2)
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_WINDOW_SECONDS", 60)
|
||
monkeypatch.setattr(rate_limit_module, "sliding_window_incr", _fake_sliding_window)
|
||
with TestClient(_make_app()) as client:
|
||
for _ in range(2):
|
||
r = client.get("/")
|
||
assert r.json()["ok"] is True
|
||
r = client.get("/")
|
||
assert r.json()["code"] == 1
|
||
# 同一窗口内持续限流，不会因后续请求解除
|
||
r = client.get("/")
|
||
assert r.json()["code"] == 1
|
||
|
||
|
||
def test_rate_limit_sliding_window_prevents_boundary_burst(monkeypatch):
|
||
"""窗口边界处请求数按时间加权，不发生固定窗口的 2× 突发。"""
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_ENABLED", True)
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_LIMIT", 3)
|
||
monkeypatch.setattr(rate_limit_module, "RATE_LIMIT_WINDOW_SECONDS", 100)
|
||
monkeypatch.setattr(rate_limit_module, "sliding_window_incr", _fake_sliding_window)
|
||
|
||
clock = {"t": 50.0}
|
||
monkeypatch.setattr(rate_limit_module.time, "time", lambda: clock["t"])
|
||
|
||
with TestClient(_make_app()) as client:
|
||
# 窗口0（t=50）：2 个请求，配额充足，放行
|
||
for _ in range(2):
|
||
assert client.get("/").json()["ok"] is True
|
||
# 进入窗口1（t=150，f=0.5）：上一窗口 2 个请求加权计入
|
||
clock["t"] = 150.0
|
||
# effective: 2*0.5+1=2 放行；2*0.5+2=3 放行；2*0.5+3=4 > 3 拒绝
|
||
for _ in range(2):
|
||
assert client.get("/").json()["ok"] is True
|
||
r = client.get("/")
|
||
assert r.json()["code"] == 1
|