hold-slayer/tests/test_rate_limit.py

"""
Rate-limiting tests.

The limiter guards the unauthenticated `/auth/*` edge — the only routes that
must answer before an identity exists. Everything else is owner-gated, so a
limit there would mostly throttle the single legitimate operator.

The properties worth pinning: the cap actually blocks, windows expire, clients
and routes don't share a bucket, the store can't grow without bound, and
identity comes from the socket peer rather than a spoofable header.
"""

import httpx
import pytest
from fastapi import Depends, FastAPI

import main
from core.rate_limit import RateLimiter, client_key, get_limiter, rate_limit


@pytest.fixture(autouse=True)
def _clean_limiter():
    get_limiter().reset()
    yield
    get_limiter().reset()


class TestRateLimiter:
    def test_allows_up_to_the_limit(self):
        rl = RateLimiter(limit=3, window_seconds=60)
        assert [rl.check("k", now=100.0)[0] for _ in range(3)] == [True, True, True]

    def test_blocks_past_the_limit(self):
        rl = RateLimiter(limit=3, window_seconds=60)
        for _ in range(3):
            rl.check("k", now=100.0)
        allowed, retry_after = rl.check("k", now=100.0)
        assert allowed is False
        assert retry_after > 0

    def test_window_expiry_resets_the_count(self):
        rl = RateLimiter(limit=2, window_seconds=60)
        rl.check("k", now=100.0)
        rl.check("k", now=100.0)
        assert rl.check("k", now=100.0)[0] is False
        # A full window later, the caller is welcome again.
        assert rl.check("k", now=161.0)[0] is True

    def test_retry_after_shrinks_as_the_window_drains(self):
        rl = RateLimiter(limit=1, window_seconds=60)
        rl.check("k", now=100.0)
        early = rl.check("k", now=110.0)[1]
        late = rl.check("k", now=150.0)[1]
        assert early > late >= 1

    def test_clients_do_not_share_a_bucket(self):
        rl = RateLimiter(limit=1, window_seconds=60)
        assert rl.check("auth:me:1.1.1.1", now=100.0)[0] is True
        # A different client must be unaffected by the first one's usage.
        assert rl.check("auth:me:2.2.2.2", now=100.0)[0] is True

    def test_routes_do_not_share_a_bucket(self):
        rl = RateLimiter(limit=1, window_seconds=60)
        assert rl.check("auth:me:1.1.1.1", now=100.0)[0] is True
        assert rl.check("auth:callback:1.1.1.1", now=100.0)[0] is True

    def test_per_call_limit_overrides_the_default(self):
        rl = RateLimiter(limit=100, window_seconds=60)
        rl.check("k", limit=1, now=100.0)
        assert rl.check("k", limit=1, now=100.0)[0] is False

    def test_bucket_store_is_bounded(self):
        # Otherwise a spray of source addresses is itself a memory exhaustion
        # vector — the thing the limiter exists to prevent.
        rl = RateLimiter(limit=5, window_seconds=60, max_clients=10)
        for i in range(50):
            rl.check(f"client-{i}", now=100.0 + i)
        assert len(rl._buckets) <= 10

    def test_eviction_drops_oldest_first(self):
        rl = RateLimiter(limit=5, window_seconds=600, max_clients=3)
        for i in range(4):
            rl.check(f"client-{i}", now=100.0 + i)
        assert "client-0" not in rl._buckets
        assert "client-3" in rl._buckets


class TestClientKey:
    def _request(self, peer: str | None, headers: dict | None = None):
        scope = {
            "type": "http",
            "method": "GET",
            "path": "/auth/me",
            "headers": [(k.lower().encode(), v.encode()) for k, v in (headers or {}).items()],
            "client": (peer, 12345) if peer else None,
        }
        from starlette.requests import Request

        return Request(scope)

    def test_uses_the_socket_peer(self):
        assert client_key(self._request("10.0.0.5"), "auth:me") == "auth:me:10.0.0.5"

    def test_ignores_x_forwarded_for(self):
        # Trusting a spoofable header would let one client present as
        # thousands, making the limiter worse than useless.
        key = client_key(
            self._request("10.0.0.5", {"X-Forwarded-For": "1.2.3.4"}), "auth:me"
        )
        assert key == "auth:me:10.0.0.5"
        assert "1.2.3.4" not in key

    def test_missing_peer_does_not_crash(self):
        assert client_key(self._request(None), "auth:me") == "auth:me:unknown"


class TestDependency:
    """The FastAPI integration: a 429 with a Retry-After header."""

    def _app(self, limit=2):
        app = FastAPI()

        @app.get("/limited", dependencies=[Depends(rate_limit("test", limit=limit))])
        async def limited():
            return {"ok": True}

        @app.get("/unlimited")
        async def unlimited():
            return {"ok": True}

        return app

    async def _get(self, app, path, n=1):
        transport = httpx.ASGITransport(app=app)
        async with httpx.AsyncClient(transport=transport, base_url="http://test") as c:
            return [await c.get(path) for _ in range(n)]

    async def test_returns_429_past_the_limit(self):
        responses = await self._get(self._app(limit=2), "/limited", n=3)
        assert [r.status_code for r in responses] == [200, 200, 429]

    async def test_429_carries_retry_after(self):
        responses = await self._get(self._app(limit=1), "/limited", n=2)
        blocked = responses[-1]
        assert blocked.status_code == 429
        assert int(blocked.headers["retry-after"]) >= 1

    async def test_unlimited_routes_are_untouched(self):
        responses = await self._get(self._app(limit=1), "/unlimited", n=10)
        assert {r.status_code for r in responses} == {200}


class TestAuthRoutesAreLimited:
    """The wiring: the unauthenticated edge is covered, the rest is not."""

    def _is_limited(self, path: str) -> bool:
        """True if the route carries a dependency built by `rate_limit`.

        Identified by the closure's qualname rather than a string search, so
        this can't pass on an unrelated dependency that happens to stringify
        similarly.
        """
        for route in main.app.routes:
            if getattr(route, "path", None) != path:
                continue
            return any(
                getattr(d.dependency, "__qualname__", "").startswith("rate_limit")
                for d in getattr(route, "dependencies", [])
            )
        raise AssertionError(f"route {path} not found")

    @pytest.mark.parametrize(
        "path", ["/auth/login", "/auth/callback", "/auth/me", "/auth/refresh-callback"]
    )
    def test_unauthenticated_auth_routes_are_limited(self, path):
        assert self._is_limited(path)

    def test_owner_gated_routes_are_not_limited(self):
        # They're already behind is_owner; limiting them would throttle the
        # single legitimate operator.
        assert not self._is_limited("/api/v1/calls/active")

    def test_logout_is_not_limited(self):
        # Pure redirect builder — no I/O, nothing to exhaust.
        assert not self._is_limited("/auth/logout")