"""Tests for Rate Limiter (core/rate_limiter.py).""" import time import pytest from core.rate_limiter import ClientRateLimitState, RateLimitConfig, RateLimiter, TokenBucket class TestRateLimitConfig: """Test rate limit configuration.""" def test_defaults(self): config = RateLimitConfig() assert config.per_minute == 60 assert config.per_hour == 1000 assert config.per_day == 10000 def test_custom_values(self): config = RateLimitConfig(per_minute=10, per_hour=100, per_day=500) assert config.per_minute == 10 def test_from_env(self, monkeypatch): monkeypatch.setenv("RATE_LIMIT_PER_MINUTE", "30") monkeypatch.setenv("RATE_LIMIT_PER_HOUR", "500") monkeypatch.setenv("RATE_LIMIT_PER_DAY", "5000") config = RateLimitConfig.from_env() assert config.per_minute == 30 assert config.per_hour == 500 assert config.per_day == 5000 def test_from_env_with_prefix(self, monkeypatch): monkeypatch.setenv("WP_RATE_LIMIT_PER_MINUTE", "20") config = RateLimitConfig.from_env("WP") assert config.per_minute == 20 class TestTokenBucket: """Test token bucket algorithm.""" def test_initial_capacity(self): bucket = TokenBucket(capacity=10, refill_rate=1.0) assert bucket.get_available_tokens() == 10 def test_consume_success(self): bucket = TokenBucket(capacity=10, refill_rate=1.0) assert bucket.consume(1) is True assert bucket.get_available_tokens() == 9 def test_consume_multiple(self): bucket = TokenBucket(capacity=10, refill_rate=1.0) assert bucket.consume(5) is True assert bucket.get_available_tokens() == 5 def test_consume_exceeds_capacity(self): bucket = TokenBucket(capacity=5, refill_rate=1.0) assert bucket.consume(6) is False # Tokens should not be consumed on failure assert bucket.get_available_tokens() == 5 def test_consume_until_empty(self): bucket = TokenBucket(capacity=3, refill_rate=0.0) # no refill assert bucket.consume(1) is True assert bucket.consume(1) is True assert bucket.consume(1) is True assert bucket.consume(1) is False def test_get_wait_time_when_available(self): bucket = TokenBucket(capacity=10, refill_rate=1.0) assert bucket.get_wait_time(1) == 0.0 def test_get_wait_time_when_empty(self): bucket = TokenBucket(capacity=1, refill_rate=1.0) bucket.consume(1) wait = bucket.get_wait_time(1) assert wait > 0 class TestClientRateLimitState: """Test per-client rate limit state.""" def _make_client(self, per_minute=5, per_hour=100, per_day=1000): return ClientRateLimitState( client_id="test-client", minute_bucket=TokenBucket(capacity=per_minute, refill_rate=per_minute / 60.0), hour_bucket=TokenBucket(capacity=per_hour, refill_rate=per_hour / 3600.0), day_bucket=TokenBucket(capacity=per_day, refill_rate=per_day / 86400.0), ) def test_allowed_request(self): client = self._make_client() allowed, reason, retry_after = client.check_and_consume() assert allowed is True assert reason == "" assert retry_after == 0.0 def test_minute_limit_exceeded(self): client = self._make_client(per_minute=2) client.check_and_consume() client.check_and_consume() allowed, reason, retry_after = client.check_and_consume() assert allowed is False assert "per minute" in reason def test_stats(self): client = self._make_client() client.check_and_consume() client.check_and_consume() stats = client.get_stats() assert stats["client_id"] == "test-client" assert stats["total_requests"] == 2 class TestRateLimiter: """Test the main RateLimiter class.""" @pytest.fixture def limiter(self): return RateLimiter() def test_first_request_allowed(self, limiter): allowed, msg, retry = limiter.check_rate_limit("client1") assert allowed is True def test_different_clients_independent(self, limiter): for _ in range(50): limiter.check_rate_limit("client1") allowed, _, _ = limiter.check_rate_limit("client2") assert allowed is True def test_get_client_stats(self, limiter): limiter.check_rate_limit("client1") stats = limiter.get_client_stats("client1") assert stats is not None assert stats["client_id"] == "client1" def test_get_client_stats_unknown(self, limiter): assert limiter.get_client_stats("unknown") is None def test_reset_client(self, limiter): limiter.check_rate_limit("client1") assert limiter.reset_client("client1") is True assert limiter.get_client_stats("client1") is None def test_reset_nonexistent_client(self, limiter): assert limiter.reset_client("nonexistent") is False def test_reset_all(self, limiter): limiter.check_rate_limit("client1") limiter.check_rate_limit("client2") count = limiter.reset_all() assert count == 2 assert limiter.get_client_stats("client1") is None def test_get_all_stats(self, limiter): limiter.check_rate_limit("client1") stats = limiter.get_all_stats() assert "global" in stats assert "default_limits" in stats assert stats["global"]["total_requests"] >= 1 assert stats["global"]["active_clients"] == 1 def test_configure_limits(self, limiter): limiter.configure_limits("n8n", per_minute=10, per_hour=50) assert limiter.plugin_configs["n8n"].per_minute == 10 assert limiter.plugin_configs["n8n"].per_hour == 50