139 lines
4.6 KiB
Python
139 lines
4.6 KiB
Python
"""限流：认证按 IP、AI 按用户。
|
|
|
|
限流默认在测试里关闭（conftest 设了 RATE_LIMIT_ENABLED=false），
|
|
这里手动开启 limiter 并重置计数器。
|
|
"""
|
|
import pytest
|
|
|
|
from app.ratelimit import limiter
|
|
|
|
|
|
@pytest.fixture
|
|
def rate_limit_on():
|
|
"""临时开启限流，结束后恢复并清空计数。"""
|
|
limiter.enabled = True
|
|
limiter.reset()
|
|
yield
|
|
limiter.reset()
|
|
limiter.enabled = False
|
|
|
|
|
|
async def test_login_rate_limited_by_ip(client, rate_limit_on):
|
|
"""暴力破解应在若干次失败后被挡住。"""
|
|
statuses = []
|
|
for _ in range(15):
|
|
r = await client.post(
|
|
"/api/auth/login", data={"username": "nobody", "password": "wrong"}
|
|
)
|
|
statuses.append(r.status_code)
|
|
|
|
assert 401 in statuses, "前几次应正常返回 401"
|
|
assert 429 in statuses, "超过阈值后应返回 429"
|
|
# 一旦触发，后续应持续被限
|
|
assert statuses[-1] == 429
|
|
|
|
|
|
async def test_ai_limit_is_per_user_not_per_ip(client, auth_headers, rate_limit_on):
|
|
"""同一 IP 下不同用户的 AI 额度互相独立——否则一人触发会误伤所有人。"""
|
|
ha = await auth_headers("alice")
|
|
hb = await auth_headers("bob")
|
|
|
|
q = {
|
|
"type": "single_choice",
|
|
"stem_markdown": "q",
|
|
"options": [{"key": "A", "text_markdown": "1"}],
|
|
"correct_answer": ["A"],
|
|
}
|
|
qa = (await client.post("/api/questions", json=q, headers=ha)).json()["id"]
|
|
qb = (await client.post("/api/questions", json=q, headers=hb)).json()["id"]
|
|
|
|
async def attempt(qid, headers):
|
|
r = await client.post(
|
|
f"/api/questions/{qid}/attempts", json={"user_answer": ["A"]}, headers=headers
|
|
)
|
|
return r.json()["id"]
|
|
|
|
# alice 一直判分到被限
|
|
alice_hit_limit = False
|
|
for _ in range(30):
|
|
aid = await attempt(qa, ha)
|
|
r = await client.post(f"/api/attempts/{aid}/judge", headers=ha)
|
|
if r.status_code == 429:
|
|
alice_hit_limit = True
|
|
break
|
|
assert alice_hit_limit, "alice 应触发 AI 限流"
|
|
|
|
# bob 不该受影响
|
|
bid = await attempt(qb, hb)
|
|
r = await client.post(f"/api/attempts/{bid}/judge", headers=hb)
|
|
assert r.status_code == 200, "bob 与 alice 共用 IP，但不该被牵连"
|
|
|
|
|
|
async def test_recording_attempt_is_not_ai_rate_limited(
|
|
client, auth_headers, rate_limit_on
|
|
):
|
|
"""记录作答不消耗 AI 额度，就不该受 AI 限流约束。
|
|
|
|
否则"先记下来、以后再判"这个主路径会被判分的额度拖住。
|
|
"""
|
|
h = await auth_headers()
|
|
qid = (
|
|
await client.post(
|
|
"/api/questions",
|
|
json={"type": "unclassified", "stem_markdown": "随手记的题"},
|
|
headers=h,
|
|
)
|
|
).json()["id"]
|
|
|
|
# 远超 AI 限额的次数，全部应成功
|
|
for i in range(40):
|
|
r = await client.post(
|
|
f"/api/questions/{qid}/attempts",
|
|
json={"user_answer": [f"第{i}次尝试"]},
|
|
headers=h,
|
|
)
|
|
assert r.status_code == 201, f"第 {i + 1} 次记录被限流了"
|
|
|
|
|
|
async def test_limit_counts_across_different_ids_in_path(
|
|
client, auth_headers, configure_ai, rate_limit_on
|
|
):
|
|
"""回归：限流必须按 endpoint 计数，不能按完整 URL。
|
|
|
|
/api/questions/{id}/ai/explain 这类路径带 ID，若按 URL 计数则每个 id
|
|
各占一个桶 —— AI 端点就等于没有限流，别人可以刷爆你的 API 额度。
|
|
"""
|
|
from unittest.mock import patch
|
|
|
|
from app.config import settings
|
|
from app.services.llm.client import LLMResult
|
|
|
|
h = await auth_headers()
|
|
await configure_ai(h)
|
|
|
|
async def fake_text(self, system, user, json_mode=False, model=None):
|
|
return LLMResult("讲解", "text-model", 1, 1)
|
|
|
|
# 每次都换一条不同的笔记（也就是不同的 URL）
|
|
limit_num = int(settings.rate_limit_ai.split("/")[0])
|
|
statuses = []
|
|
with patch("app.services.llm.client.LLMClient.chat_text", fake_text):
|
|
for i in range(limit_num + 3):
|
|
qid = (
|
|
await client.post(
|
|
"/api/questions",
|
|
json={"type": "unclassified", "stem_markdown": f"第{i}条"},
|
|
headers=h,
|
|
)
|
|
).json()["id"]
|
|
r = await client.post(f"/api/questions/{qid}/ai/explain", headers=h)
|
|
statuses.append(r.status_code)
|
|
|
|
assert 429 in statuses, "不同 URL 也必须共享同一个 AI 限流额度"
|
|
|
|
|
|
async def test_health_not_rate_limited(client, rate_limit_on):
|
|
"""健康检查会被探针频繁调用，不能限流。"""
|
|
for _ in range(30):
|
|
assert (await client.get("/api/health")).status_code == 200
|