flaskpaste/tests/test_rate_limiting.py

"""Tests for IP-based rate limiting."""

import json


class TestRateLimiting:
    """Tests for rate limiting on paste creation."""

    def test_rate_limit_allows_normal_usage(self, client, sample_text):
        """Normal usage within rate limit succeeds."""
        # TestingConfig has RATE_LIMIT_MAX=100
        for i in range(5):
            response = client.post(
                "/",
                data=f"paste {i}",
                content_type="text/plain",
            )
            assert response.status_code == 201

    def test_rate_limit_exceeded_returns_429(self, client, app):
        """Exceeding rate limit returns 429."""
        # Temporarily lower rate limit for test
        original_max = app.config["RATE_LIMIT_MAX"]
        app.config["RATE_LIMIT_MAX"] = 3

        try:
            # Make requests up to limit
            for i in range(3):
                response = client.post(
                    "/",
                    data=f"paste {i}",
                    content_type="text/plain",
                )
                assert response.status_code == 201

            # Next request should be rate limited
            response = client.post(
                "/",
                data="one more",
                content_type="text/plain",
            )
            assert response.status_code == 429
            data = json.loads(response.data)
            assert "error" in data
            assert "Rate limit" in data["error"]
            assert "retry_after" in data
        finally:
            app.config["RATE_LIMIT_MAX"] = original_max

    def test_rate_limit_headers(self, client, app):
        """Rate limit response includes proper headers."""
        original_max = app.config["RATE_LIMIT_MAX"]
        app.config["RATE_LIMIT_MAX"] = 1

        try:
            # First request succeeds
            client.post("/", data="first", content_type="text/plain")

            # Second request is rate limited
            response = client.post("/", data="second", content_type="text/plain")
            assert response.status_code == 429
            assert "Retry-After" in response.headers
            assert "X-RateLimit-Remaining" in response.headers
            assert response.headers["X-RateLimit-Remaining"] == "0"
        finally:
            app.config["RATE_LIMIT_MAX"] = original_max

    def test_rate_limit_headers_on_success(self, client, app):
        """Successful responses include rate limit headers."""
        original_max = app.config["RATE_LIMIT_MAX"]
        app.config["RATE_LIMIT_MAX"] = 5

        try:
            # First request should include rate limit headers
            response = client.post("/", data="first", content_type="text/plain")
            assert response.status_code == 201

            # Check rate limit headers
            assert "X-RateLimit-Limit" in response.headers
            assert "X-RateLimit-Remaining" in response.headers
            assert "X-RateLimit-Reset" in response.headers

            # Verify values
            assert response.headers["X-RateLimit-Limit"] == "5"
            assert response.headers["X-RateLimit-Remaining"] == "4"  # 5 - 1 = 4

            # Reset timestamp should be a valid unix timestamp
            reset = int(response.headers["X-RateLimit-Reset"])
            import time

            assert reset > int(time.time())  # Should be in the future
        finally:
            app.config["RATE_LIMIT_MAX"] = original_max

    def test_rate_limit_auth_multiplier(self, client, app, auth_header):
        """Authenticated users get higher rate limits."""
        original_max = app.config["RATE_LIMIT_MAX"]
        original_mult = app.config["RATE_LIMIT_AUTH_MULTIPLIER"]
        app.config["RATE_LIMIT_MAX"] = 2
        app.config["RATE_LIMIT_AUTH_MULTIPLIER"] = 3  # 2 * 3 = 6 for auth users

        try:
            # Authenticated user can make more requests than base limit
            for i in range(5):
                response = client.post(
                    "/",
                    data=f"auth {i}",
                    content_type="text/plain",
                    headers=auth_header,
                )
                assert response.status_code == 201

            # 6th request should succeed (limit is 2*3=6)
            response = client.post(
                "/",
                data="auth 6",
                content_type="text/plain",
                headers=auth_header,
            )
            assert response.status_code == 201

            # 7th should fail
            response = client.post(
                "/",
                data="auth 7",
                content_type="text/plain",
                headers=auth_header,
            )
            assert response.status_code == 429
        finally:
            app.config["RATE_LIMIT_MAX"] = original_max
            app.config["RATE_LIMIT_AUTH_MULTIPLIER"] = original_mult

    def test_rate_limit_can_be_disabled(self, client, app):
        """Rate limiting can be disabled via config."""
        original_enabled = app.config["RATE_LIMIT_ENABLED"]
        original_max = app.config["RATE_LIMIT_MAX"]
        app.config["RATE_LIMIT_ENABLED"] = False
        app.config["RATE_LIMIT_MAX"] = 1

        try:
            # Should be able to make many requests
            for i in range(5):
                response = client.post(
                    "/",
                    data=f"paste {i}",
                    content_type="text/plain",
                )
                assert response.status_code == 201
        finally:
            app.config["RATE_LIMIT_ENABLED"] = original_enabled
            app.config["RATE_LIMIT_MAX"] = original_max

    def test_rate_limit_only_affects_paste_creation(self, client, app, sample_text):
        """Rate limiting only affects POST /, not GET endpoints."""
        original_max = app.config["RATE_LIMIT_MAX"]
        app.config["RATE_LIMIT_MAX"] = 2

        try:
            # Create paste
            create = client.post("/", data=sample_text, content_type="text/plain")
            assert create.status_code == 201
            paste_id = json.loads(create.data)["id"]

            # Use up rate limit
            client.post("/", data="second", content_type="text/plain")

            # Should be rate limited for creation
            response = client.post("/", data="third", content_type="text/plain")
            assert response.status_code == 429

            # But GET should still work
            response = client.get(f"/{paste_id}")
            assert response.status_code == 200

            response = client.get(f"/{paste_id}/raw")
            assert response.status_code == 200

            response = client.get("/health")
            assert response.status_code == 200
        finally:
            app.config["RATE_LIMIT_MAX"] = original_max


class TestRateLimitCleanup:
    """Tests for rate limit cleanup."""

    def test_rate_limit_window_expiry(self, app):
        """Rate limit cleanup works with explicit window."""
        from app.api.routes import cleanup_rate_limits

        # Should work without app context when window is explicit
        cleaned = cleanup_rate_limits(window=60)
        assert isinstance(cleaned, int)
        assert cleaned >= 0