backend/tests/test_rate_limit.py

"""
Unit tests for rate limiting middleware and identifier resolution (#445).

Verifies key function fallback resolutions (User ID vs IP Address), route
attribute assignments, and that the global handler intercepts limit breaches
to return a 429 status code.
"""
from types import SimpleNamespace
from uuid import uuid4
import pytest
from fastapi.testclient import TestClient
from slowapi.errors import RateLimitExceeded

from app.auth import create_access_token
from app.rate_limit import CHAT_QUERY_RATE_LIMIT, rate_limit_key_func
from app.routes.chat import ask_question, ask_question_stream
from app.main import app


class DummyRequest:
    def __init__(self, headers=None):
        self.headers = headers or {}
        self.client = SimpleNamespace(host="203.0.113.10")


# ── Your Original Key Resolution & Route Tests ───────────────────────────────

def test_rate_limit_key_prefers_authenticated_user_id():
    token = create_access_token("user-123")

    key = rate_limit_key_func(
        DummyRequest(headers={"authorization": f"Bearer {token}"})
    )

    assert key == "user:user-123"


def test_rate_limit_key_falls_back_to_client_ip():
    key = rate_limit_key_func(DummyRequest())

    assert key.startswith("ip:")


def test_chat_endpoints_use_required_rate_limit():
    assert CHAT_QUERY_RATE_LIMIT == "15/minute"
    assert ask_question.__rate_limits__ == [CHAT_QUERY_RATE_LIMIT]
    assert ask_question_stream.__rate_limits__ == [CHAT_QUERY_RATE_LIMIT]


# ── Middleware 429 Response Verification ──────────────────────────────────────

def test_rate_limit_handler_returns_429(client: TestClient):
    """
    Verify that hitting an endpoint that triggers a RateLimitExceeded exception
    correctly triggers the global handler, returning a 429 status code and the
    exact JSON error layout specified in app/main.py.
    """
    from fastapi import APIRouter
    test_router = APIRouter()
    @test_router.get("/api/v1/test-rate-limiting-trigger-429")
    def trigger_rate_limit():
        raise RateLimitExceeded("Too Many Requests")
    
    app.include_router(test_router)
    # Move the newly added route to the top so it takes precedence over the SPA catch-all
    new_route = app.router.routes.pop()
    app.router.routes.insert(0, new_route)

    response = client.get("/api/v1/test-rate-limiting-trigger-429")
    
    # Verify the 429 status code requirement
    assert response.status_code == 429
    
    # Verify the specific JSON payload structure from app/main.py
    json_data = response.json()
    assert "error" in json_data
    assert json_data["error"]["code"] == "RATE_LIMIT_EXCEEDED"
    assert "Rate limit exceeded. Please try again later." in json_data["error"]["message"]
    assert "request_id" in json_data["error"]
    assert isinstance(json_data["error"]["details"], dict)


# ── Manual /chat/ws Rate Limit Enforcement ──────────────────────────────────

def test_chat_ws_rate_limit_allows_up_to_configured_threshold():
    """check_chat_ws_rate_limit must allow exactly CHAT_QUERY_RATE_LIMIT hits."""
    from app.rate_limit import CHAT_QUERY_RATE_LIMIT, check_chat_ws_rate_limit

    limit = int(CHAT_QUERY_RATE_LIMIT.split("/")[0])
    user_id = f"ws-rate-limit-{uuid4()}"

    for _ in range(limit):
        assert check_chat_ws_rate_limit(user_id) is True


def test_chat_ws_rate_limit_blocks_after_threshold_exceeded():
    """The hit beyond CHAT_QUERY_RATE_LIMIT must be rejected — this is the
    bypass from #639, where /chat/ws had no rate limiting at all."""
    from app.rate_limit import CHAT_QUERY_RATE_LIMIT, check_chat_ws_rate_limit

    limit = int(CHAT_QUERY_RATE_LIMIT.split("/")[0])
    user_id = f"ws-rate-limit-{uuid4()}"

    for _ in range(limit):
        check_chat_ws_rate_limit(user_id)

    assert check_chat_ws_rate_limit(user_id) is False


def test_chat_ws_rate_limit_is_scoped_per_user():
    """One user's usage must not consume another user's rate-limit budget."""
    from app.rate_limit import check_chat_ws_rate_limit

    user_a = f"ws-rate-limit-{uuid4()}"
    user_b = f"ws-rate-limit-{uuid4()}"

    assert check_chat_ws_rate_limit(user_a) is True
    assert check_chat_ws_rate_limit(user_b) is True