Spaces:
Runtime error
Runtime error
| import pytest | |
| from uuid import uuid4 | |
| from unittest.mock import patch | |
| from app.rate_limit import _hits | |
| def test_query_rate_limit(client, db_session): | |
| # Clear the in-memory rate limit dictionary between tests | |
| _hits.clear() | |
| client.post("/api/v1/auth/signup", json={"email": "ratelimit@test.com", "password": "Password123"}) | |
| token = client.post("/api/v1/auth/login", json={"email": "ratelimit@test.com", "password": "Password123"}).json()["access_token"] | |
| headers = {"Authorization": f"Bearer {token}"} | |
| hub_res = client.post("/api/v1/hubs", headers={**headers, "X-Idempotency-Key": str(uuid4())}, json={"name": "Rate Limit Hub"}) | |
| hub_id = hub_res.json()["id"] | |
| # We mock the AI to be fast so we can make 10 rapid requests | |
| with patch("app.routers.hubs.expand_query", return_value=[]), \ | |
| patch("app.routers.hubs.embed_texts", return_value=[[0.1]*768]), \ | |
| patch("app.routers.hubs.answer_question", return_value={"answer": "Test", "evidence": None, "grounded": False}): | |
| # Make 10 successful queries | |
| for _ in range(10): | |
| res = client.post(f"/api/v1/hubs/{hub_id}/query", headers=headers, json={"question": "Test query?"}) | |
| assert res.status_code == 200 | |
| # The 11th query should be rate limited | |
| res_11 = client.post(f"/api/v1/hubs/{hub_id}/query", headers=headers, json={"question": "Test query?"}) | |
| assert res_11.status_code == 429 | |
| assert res_11.json()["error"]["code"] == "RATE_LIMITED" | |
| assert "Retry-After" in res_11.headers | |
| # A different user should not be rate limited (per-user isolation) | |
| client.post("/api/v1/auth/signup", json={"email": "ratelimit2@test.com", "password": "Password123"}) | |
| token2 = client.post("/api/v1/auth/login", json={"email": "ratelimit2@test.com", "password": "Password123"}).json()["access_token"] | |
| with patch("app.routers.hubs.expand_query", return_value=[]), \ | |
| patch("app.routers.hubs.embed_texts", return_value=[[0.1]*768]), \ | |
| patch("app.routers.hubs.answer_question", return_value={"answer": "Test", "evidence": None, "grounded": False}): | |
| res_other = client.post(f"/api/v1/hubs/{hub_id}/query", headers={"Authorization": f"Bearer {token2}"}, json={"question": "Test query?"}) | |
| assert res_other.status_code == 404 # Because User 2 doesn't own Hub 1 |