"""Tests for unified rate-limit policies and central check_rate_limit(). Covers: - check_rate_limit() with various policies (AUTH, AI, UPLOAD, WEBHOOK) - Rate limit triggers on exceeding max attempts - Reset after window timeout (in-memory fallback) - get_client_ip() with and without trusted proxy - Forgejo error reporter uses central check_rate_limit() - AI proactive is_rate_limited() delegates to central check_rate_limit() """ from __future__ import annotations import asyncio import uuid from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException, Request from starlette.datastructures import Headers from app.core.rate_limit import ( RateLimitPolicy, check_rate_limit, check_rate_limit_policy, get_client_ip, reset_rate_limit, ) # ─── Unit Tests: RateLimitPolicy ─── @pytest.mark.asyncio async def test_rate_limit_policy_auth_limits(): """RateLimitPolicy.AUTH returns (rate_limit_auth_max, rate_limit_auth_window).""" max_attempts, window = RateLimitPolicy.AUTH.limits() assert max_attempts == 5 assert window == 300 @pytest.mark.asyncio async def test_rate_limit_policy_ai_limits(): """RateLimitPolicy.AI returns (rate_limit_ai_max, rate_limit_ai_window).""" max_attempts, window = RateLimitPolicy.AI.limits() assert max_attempts == 20 assert window == 60 @pytest.mark.asyncio async def test_rate_limit_policy_upload_limits(): """RateLimitPolicy.UPLOAD returns (rate_limit_upload_max, rate_limit_upload_window).""" max_attempts, window = RateLimitPolicy.UPLOAD.limits() assert max_attempts == 30 assert window == 60 @pytest.mark.asyncio async def test_rate_limit_policy_webhook_limits(): """RateLimitPolicy.WEBHOOK returns (rate_limit_webhook_max, rate_limit_webhook_window).""" max_attempts, window = RateLimitPolicy.WEBHOOK.limits() assert max_attempts == 100 assert window == 60 # ─── Unit Tests: check_rate_limit via in-memory fallback ─── def _force_inmemory_fallback(): """Return a context manager that patches the circuit breaker to force in-memory fallback.""" circuit = MagicMock() circuit.can_proceed = AsyncMock(return_value=False) circuit.record_success = AsyncMock() circuit.record_failure = AsyncMock() return patch("app.core.resilience.get_circuit", return_value=circuit) @pytest.mark.asyncio async def test_check_rate_limit_allows_under_max(): """check_rate_limit allows requests up to max_attempts.""" key = f"test:allow:{uuid.uuid4()}" with _force_inmemory_fallback(): for _ in range(3): await check_rate_limit(key, max_attempts=3, window_seconds=60) @pytest.mark.asyncio async def test_check_rate_limit_blocks_on_exceed(): """check_rate_limit raises 429 when max_attempts is exceeded.""" key = f"test:block:{uuid.uuid4()}" with _force_inmemory_fallback(): for _ in range(2): await check_rate_limit(key, max_attempts=2, window_seconds=60) with pytest.raises(HTTPException) as exc_info: await check_rate_limit(key, max_attempts=2, window_seconds=60) assert exc_info.value.status_code == 429 @pytest.mark.asyncio async def test_check_rate_limit_reset_after_window(): """Rate limit resets after the window timeout (in-memory fallback).""" key = f"test:reset:{uuid.uuid4()}" with _force_inmemory_fallback(): # Use up the limit with a very short window await check_rate_limit(key, max_attempts=1, window_seconds=1) with pytest.raises(HTTPException): await check_rate_limit(key, max_attempts=1, window_seconds=1) # Wait for window to expire await asyncio.sleep(1.1) # Should be allowed again await check_rate_limit(key, max_attempts=1, window_seconds=1) @pytest.mark.asyncio async def test_reset_rate_limit_clears_counter(): """reset_rate_limit clears the counter so subsequent calls succeed.""" key = f"test:reset2:{uuid.uuid4()}" with _force_inmemory_fallback(): await check_rate_limit(key, max_attempts=1, window_seconds=60) with pytest.raises(HTTPException): await check_rate_limit(key, max_attempts=1, window_seconds=60) await reset_rate_limit(key) await check_rate_limit(key, max_attempts=1, window_seconds=60) # ─── Unit Tests: check_rate_limit_policy ─── @pytest.mark.asyncio async def test_check_rate_limit_policy_auth(): """check_rate_limit_policy with AUTH policy enforces auth limits.""" key = f"test:policy:auth:{uuid.uuid4()}" with _force_inmemory_fallback(): max_attempts, _ = RateLimitPolicy.AUTH.limits() for _ in range(max_attempts): await check_rate_limit_policy(key, RateLimitPolicy.AUTH) with pytest.raises(HTTPException) as exc_info: await check_rate_limit_policy(key, RateLimitPolicy.AUTH) assert exc_info.value.status_code == 429 @pytest.mark.asyncio async def test_check_rate_limit_policy_ai(): """check_rate_limit_policy with AI policy enforces AI limits.""" key = f"test:policy:ai:{uuid.uuid4()}" with _force_inmemory_fallback(): max_attempts, _ = RateLimitPolicy.AI.limits() for _ in range(max_attempts): await check_rate_limit_policy(key, RateLimitPolicy.AI) with pytest.raises(HTTPException) as exc_info: await check_rate_limit_policy(key, RateLimitPolicy.AI) assert exc_info.value.status_code == 429 # ─── Unit Tests: get_client_ip ─── def _make_request(client_host: str = "127.0.0.1", headers: dict | None = None) -> Request: """Build a mock Request for get_client_ip testing.""" scope = { "type": "http", "method": "GET", "path": "/", "headers": Headers(headers or {}).raw, "client": (client_host, 8000), } return Request(scope) def test_get_client_ip_no_trusted_proxy(): """get_client_ip returns direct IP when no trusted proxies configured.""" with patch("app.config.get_settings") as mock_settings: settings = MagicMock() settings.trusted_proxy_cidrs = "" mock_settings.return_value = settings req = _make_request("192.168.1.100", {"x-forwarded-for": "10.0.0.1"}) ip = get_client_ip(req) assert ip == "192.168.1.100" def test_get_client_ip_with_trusted_proxy(): """get_client_ip uses X-Forwarded-For when direct client is a trusted proxy.""" with patch("app.config.get_settings") as mock_settings: settings = MagicMock() settings.trusted_proxy_cidrs = "10.0.0.0/8" mock_settings.return_value = settings req = _make_request("10.0.0.1", {"x-forwarded-for": "203.0.113.50"}) ip = get_client_ip(req) assert ip == "203.0.113.50" def test_get_client_ip_trusted_proxy_uses_x_real_ip(): """get_client_ip falls back to X-Real-IP when X-Forwarded-For is absent.""" with patch("app.config.get_settings") as mock_settings: settings = MagicMock() settings.trusted_proxy_cidrs = "10.0.0.0/8" mock_settings.return_value = settings req = _make_request("10.0.0.1", {"x-real-ip": "203.0.113.99"}) ip = get_client_ip(req) assert ip == "203.0.113.99" def test_get_client_ip_untrusted_proxy_ignored(): """get_client_ip ignores X-Forwarded-For from untrusted proxy.""" with patch("app.config.get_settings") as mock_settings: settings = MagicMock() settings.trusted_proxy_cidrs = "10.0.0.0/8" mock_settings.return_value = settings req = _make_request("192.168.1.1", {"x-forwarded-for": "203.0.113.50"}) ip = get_client_ip(req) assert ip == "192.168.1.1" def test_get_client_ip_multiple_forwarded(): """get_client_ip uses leftmost IP from comma-separated X-Forwarded-For.""" with patch("app.config.get_settings") as mock_settings: settings = MagicMock() settings.trusted_proxy_cidrs = "10.0.0.0/8" mock_settings.return_value = settings req = _make_request("10.0.0.1", {"x-forwarded-for": "203.0.113.50, 10.0.0.2, 10.0.0.3"}) ip = get_client_ip(req) assert ip == "203.0.113.50" # ─── Integration Tests: Forgejo reporter uses central rate limit ─── @pytest.mark.asyncio async def test_forgejo_reporter_uses_central_rate_limit(): """Forgejo error reporter delegates to check_rate_limit() and respects limits.""" from app.plugins.builtins.forgejo_error_reporter.service import ( _RATE_LIMIT_MAX, _RATE_LIMIT_WINDOW, report_error_to_forgejo, ) assert _RATE_LIMIT_MAX == 10 assert _RATE_LIMIT_WINDOW == 300 call_count = 0 async def mock_check_rate_limit(redis_key, max_attempts, window_seconds): nonlocal call_count call_count += 1 assert redis_key == "rate:forgejo_report:global" assert max_attempts == 10 assert window_seconds == 300 if call_count > _RATE_LIMIT_MAX: raise HTTPException(status_code=429, detail="Rate limit exceeded") with patch( "app.plugins.builtins.forgejo_error_reporter.service._is_duplicate", new_callable=AsyncMock, return_value=False, ), patch( "app.plugins.builtins.forgejo_error_reporter.service._get_settings", return_value={ "url": "https://forgejo.example.com", "token": "test-token", "owner": "test", "repo": "test", }, ), patch( "app.core.rate_limit.check_rate_limit", side_effect=mock_check_rate_limit, ): entry = {"message": "Test error", "stack": "trace"} for i in range(_RATE_LIMIT_MAX): await report_error_to_forgejo(entry) assert call_count == i + 1 result = await report_error_to_forgejo(entry) assert result is False assert call_count == _RATE_LIMIT_MAX + 1 # ─── Integration Tests: AI proactive uses central rate limit ─── @pytest.mark.asyncio async def test_ai_proactive_uses_central_rate_limit(): """AI proactive is_rate_limited() delegates to check_rate_limit().""" from app.plugins.builtins.ai_proactive.services import is_rate_limited tenant_id = uuid.uuid4() user_id = uuid.uuid4() rate_limit_seconds = 5 # First call should not be rate-limited result = await is_rate_limited(tenant_id, user_id, rate_limit_seconds) assert result is False # Second call (within window) should be rate-limited result = await is_rate_limited(tenant_id, user_id, rate_limit_seconds) assert result is True @pytest.mark.asyncio async def test_ai_proactive_rate_limit_key_format(): """AI proactive uses the correct Redis key format: rate:ai_proactive:{tenant_id}:{user_id}.""" tenant_id = uuid.uuid4() user_id = uuid.uuid4() expected_key = f"rate:ai_proactive:{tenant_id}:{user_id}" captured_key = None async def mock_check_rate_limit(redis_key, max_attempts, window_seconds): nonlocal captured_key captured_key = redis_key assert max_attempts == 1 assert window_seconds == 10 with patch("app.core.rate_limit.check_rate_limit", side_effect=mock_check_rate_limit): from app.plugins.builtins.ai_proactive.services import is_rate_limited await is_rate_limited(tenant_id, user_id, rate_limit_seconds=10) assert captured_key == expected_key @pytest.mark.asyncio async def test_ai_proactive_rate_limit_exception_allows(): """AI proactive allows request when rate limit check raises non-HTTP exception.""" with patch("app.core.rate_limit.check_rate_limit", side_effect=RuntimeError("Redis down")): from app.plugins.builtins.ai_proactive.services import is_rate_limited result = await is_rate_limited(uuid.uuid4(), uuid.uuid4(), rate_limit_seconds=10) assert result is False # ─── Integration Tests: Independent keys per user/tenant ─── @pytest.mark.asyncio async def test_independent_rate_limit_keys(): """Different keys have independent rate limits.""" key1 = f"test:indep:1:{uuid.uuid4()}" key2 = f"test:indep:2:{uuid.uuid4()}" with _force_inmemory_fallback(): await check_rate_limit(key1, max_attempts=1, window_seconds=60) with pytest.raises(HTTPException): await check_rate_limit(key1, max_attempts=1, window_seconds=60) # key2 should still be allowed await check_rate_limit(key2, max_attempts=1, window_seconds=60)