336 lines
12 KiB
Python
336 lines
12 KiB
Python
|
|
"""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)
|