tests/test_rate_limit.py

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