148 lines
4.4 KiB
Python
148 lines
4.4 KiB
Python
"""
|
|
Tests for Rate Limiter.
|
|
|
|
Tests per-tenant rate limiting functionality.
|
|
"""
|
|
|
|
import pytest
|
|
import asyncio
|
|
from unittest.mock import Mock
|
|
|
|
from src.agents.rate_limiter import (
|
|
RateLimiter,
|
|
RateLimitConfig,
|
|
RateLimitResult,
|
|
TenantPriority,
|
|
get_rate_limiter
|
|
)
|
|
|
|
|
|
class TestRateLimitConfig:
|
|
"""Tests for RateLimitConfig dataclass."""
|
|
|
|
def test_default_values(self):
|
|
"""Test default configuration values."""
|
|
config = RateLimitConfig()
|
|
|
|
assert config.requests_per_minute == 10
|
|
assert config.requests_per_hour == 100
|
|
assert config.concurrent_limit == 3
|
|
assert config.priority == TenantPriority.NORMAL
|
|
|
|
def test_custom_values(self):
|
|
"""Test custom configuration."""
|
|
config = RateLimitConfig(
|
|
requests_per_minute=20,
|
|
requests_per_hour=500,
|
|
concurrent_limit=5,
|
|
priority=TenantPriority.HIGH
|
|
)
|
|
|
|
assert config.requests_per_minute == 20
|
|
assert config.priority == TenantPriority.HIGH
|
|
|
|
|
|
class TestRateLimiter:
|
|
"""Tests for RateLimiter class."""
|
|
|
|
@pytest.fixture
|
|
def limiter(self):
|
|
"""Create a fresh rate limiter for testing."""
|
|
return RateLimiter()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_first_request_allowed(self, limiter):
|
|
"""Test that first request is always allowed."""
|
|
result = await limiter.check_limit("tenant-001")
|
|
|
|
assert result.allowed is True
|
|
assert result.tenant_id == "tenant-001"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_concurrent_limit(self, limiter):
|
|
"""Test concurrent request limiting."""
|
|
# Configure low concurrent limit
|
|
limiter.configure_tenant("tenant-002", RateLimitConfig(concurrent_limit=2))
|
|
|
|
# First two requests allowed
|
|
r1 = await limiter.check_limit("tenant-002")
|
|
r2 = await limiter.check_limit("tenant-002")
|
|
|
|
assert r1.allowed is True
|
|
assert r2.allowed is True
|
|
|
|
# Third request blocked
|
|
r3 = await limiter.check_limit("tenant-002")
|
|
|
|
assert r3.allowed is False
|
|
assert "simultâneas" in r3.reason.lower()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_release_slot(self, limiter):
|
|
"""Test releasing a request slot."""
|
|
limiter.configure_tenant("tenant-003", RateLimitConfig(concurrent_limit=1))
|
|
|
|
# Use the slot
|
|
r1 = await limiter.check_limit("tenant-003")
|
|
assert r1.allowed is True
|
|
|
|
# Should be blocked
|
|
r2 = await limiter.check_limit("tenant-003")
|
|
assert r2.allowed is False
|
|
|
|
# Release and try again
|
|
await limiter.release("tenant-003")
|
|
r3 = await limiter.check_limit("tenant-003")
|
|
assert r3.allowed is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_minute_limit(self, limiter):
|
|
"""Test per-minute rate limiting."""
|
|
limiter.configure_tenant("tenant-004", RateLimitConfig(
|
|
requests_per_minute=3,
|
|
concurrent_limit=10
|
|
))
|
|
|
|
# Use up minute quota
|
|
for _ in range(3):
|
|
r = await limiter.check_limit("tenant-004")
|
|
await limiter.release("tenant-004")
|
|
assert r.allowed is True
|
|
|
|
# Next should be blocked
|
|
r = await limiter.check_limit("tenant-004")
|
|
assert r.allowed is False
|
|
assert "/min" in r.reason
|
|
|
|
def test_get_tenant_stats(self, limiter):
|
|
"""Test tenant statistics."""
|
|
stats = limiter.get_tenant_stats("nonexistent")
|
|
assert stats == {}
|
|
|
|
def test_configure_tenant(self, limiter):
|
|
"""Test tenant configuration."""
|
|
config = RateLimitConfig(
|
|
requests_per_minute=50,
|
|
priority=TenantPriority.CRITICAL
|
|
)
|
|
|
|
limiter.configure_tenant("vip-tenant", config)
|
|
|
|
result = limiter.get_config("vip-tenant")
|
|
assert result.requests_per_minute == 50
|
|
assert result.priority == TenantPriority.CRITICAL
|
|
|
|
|
|
class TestRateLimiterSingleton:
|
|
"""Tests for singleton pattern."""
|
|
|
|
def test_singleton(self):
|
|
"""Test singleton returns same instance."""
|
|
import src.agents.rate_limiter as module
|
|
module._rate_limiter = None
|
|
|
|
l1 = get_rate_limiter()
|
|
l2 = get_rate_limiter()
|
|
|
|
assert l1 is l2
|