签入版本

tests/unit/middleware/test_rate_limit.py

@@ -0,0 +1,105 @@
+"""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