kite/backend/tests/test_rate_limiting.py

"""
Rate limiting tests — D-11, D-12, Assumption A1.

Plan 06-05: all 8 xfail stubs promoted to real assertions.

D-11: trusted-proxy CIDR logic in get_client_ip (deps/utils.py).
D-12: per-account rate limiter keyed by user_id (100 req/min).
A1:   slowapi key_func ordering assumption — request.state.current_user must be
      set as the FIRST line of the handler body before slowapi reads the key.
"""
from __future__ import annotations

import uuid

import pytest
import pytest_asyncio
from fastapi import FastAPI
from fastapi.testclient import TestClient
from slowapi.errors import RateLimitExceeded
from slowapi import _rate_limit_exceeded_handler
from slowapi.middleware import SlowAPIMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse


# ── Helpers: build a minimal Starlette Request for unit tests ─────────────────

def _make_request(
    client_host: str | None,
    xff: str | None = None,
) -> Request:
    """Build a minimal starlette.requests.Request with the given peer IP and XFF header."""
    headers_list: list[tuple[bytes, bytes]] = []
    if xff is not None:
        headers_list.append((b"x-forwarded-for", xff.encode()))

    scope = {
        "type": "http",
        "method": "GET",
        "path": "/",
        "query_string": b"",
        "headers": headers_list,
        "client": (client_host, 12345) if client_host else None,
    }
    return Request(scope)


# ── D-11: trusted-proxy CIDR logic in get_client_ip ──────────────────────────

def test_get_client_ip_untrusted_returns_direct_peer():
    """When request.client.host is 8.8.8.8 (not in trusted CIDRs), ignore
    X-Forwarded-For and return the direct peer IP '8.8.8.8'."""
    from deps.utils import get_client_ip

    req = _make_request("8.8.8.8", xff="1.2.3.4")
    assert get_client_ip(req) == "8.8.8.8"


def test_get_client_ip_trusted_proxy_reads_xff_leftmost():
    """When request.client.host is 127.0.0.1 (trusted), return the leftmost
    address from X-Forwarded-For: '1.2.3.4, 5.6.7.8' → '1.2.3.4'."""
    from deps.utils import get_client_ip

    req = _make_request("127.0.0.1", xff="1.2.3.4, 5.6.7.8")
    assert get_client_ip(req) == "1.2.3.4"


def test_get_client_ip_trusted_proxy_no_xff_falls_back():
    """When the direct peer is trusted (172.16.5.5) but no X-Forwarded-For
    header is present, return the direct peer IP as fallback."""
    from deps.utils import get_client_ip

    req = _make_request("172.16.5.5", xff=None)
    assert get_client_ip(req) == "172.16.5.5"


def test_get_client_ip_invalid_peer_returns_none_or_string():
    """When request.client is None, get_client_ip returns None without raising."""
    from deps.utils import get_client_ip

    req = _make_request(None, xff=None)
    result = get_client_ip(req)
    assert result is None


# ── D-12: per-account rate limiter keyed by user_id ──────────────────────────

def test_account_limiter_key_uses_user_id():
    """_account_key(request) where request.state.current_user has id=UUID(...)
    returns str(user.id), not request.client.host."""
    from services.rate_limiting import _account_key

    req = _make_request("8.8.8.8")

    class _FakeUser:
        id = uuid.uuid4()

    req.state.current_user = _FakeUser()
    result = _account_key(req)
    assert result == str(_FakeUser.id)
    assert result != "8.8.8.8"


def test_account_limiter_key_falls_back_to_ip_when_no_user():
    """When request.state.current_user is missing, the key function returns the
    direct peer IP — must not crash (Pitfall 3 guard)."""
    from services.rate_limiting import _account_key

    req = _make_request("9.9.9.9")
    # Do NOT set request.state.current_user
    result = _account_key(req)
    assert result == "9.9.9.9"


# ── A1: key_func ordering assumption — unit verification ─────────────────────

def test_account_limiter_key_ordering_assumption():
    """A1 verification: construct a FastAPI test app with one endpoint decorated
    @account_limiter.limit('100/minute') that sets request.state.current_user as
    its first line; call it 101 times with the same user; assert the 101st
    response is 429.

    This validates that slowapi reads the key_func AFTER the first line of the
    handler has already set request.state.current_user (the ordering assumption
    that makes per-account limiting work).

    Key correctness is verified separately via the returned body — the endpoint
    echoes back the key so we can confirm it is the user.id, not the peer IP.
    """
    from services.rate_limiting import _account_key

    _user_id = uuid.uuid4()

    # Build an isolated limiter so this test never pollutes the shared instance
    from slowapi import Limiter
    _isolated_limiter = Limiter(key_func=_account_key)

    test_app = FastAPI()
    test_app.state.limiter = _isolated_limiter
    test_app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
    test_app.add_middleware(SlowAPIMiddleware)

    class _FakeUser:
        id = _user_id

    @test_app.get("/limited")
    @_isolated_limiter.limit("100/minute")
    async def _limited_endpoint(request: Request):
        # FIRST line: set current_user so the key_func can read it
        request.state.current_user = _FakeUser()
        # Echo the key back so the test can assert it is the user.id
        key = _account_key(request)
        return {"ok": True, "key": key}

    with TestClient(test_app, raise_server_exceptions=False) as client:
        responses = [client.get("/limited") for _ in range(101)]

    status_codes = [r.status_code for r in responses]
    assert status_codes[-1] == 429, (
        f"Expected 429 on 101st request, got {status_codes[-1]}"
    )
    assert all(c == 200 for c in status_codes[:100]), (
        "First 100 requests should all be 200"
    )
    # Verify the key returned in the body is the user.id, not the IP
    first_body = responses[0].json()
    assert first_body.get("key") == str(_user_id), (
        f"Expected key '{_user_id}', got '{first_body.get('key')}'"
    )


# ── D-12: full integration — 429 after 100 requests/minute ───────────────────

@pytest.mark.asyncio
async def test_authenticated_endpoint_429_after_100_per_minute(
    async_client, auth_user
):
    """Full integration: GET /api/documents/ called 101 times with the same
    auth_user returns 429 on the 101st request."""
    headers = auth_user["headers"]

    responses = []
    for _ in range(101):
        r = await async_client.get("/api/documents", headers=headers)
        responses.append(r)

    status_codes = [r.status_code for r in responses]
    assert status_codes[-1] == 429, (
        f"Expected 429 on 101st request, got {status_codes[-1]}"
    )
    assert all(c == 200 for c in status_codes[:100]), (
        f"First 100 should be 200, got: {[c for c in status_codes[:100] if c != 200]}"
    )