import time import unittest from utils.rate_limiter import SimpleRateLimiter class TestSimpleRateLimiter(unittest.TestCase): def test_allows_requests_within_limit(self): limiter = SimpleRateLimiter(max_requests=5, window_seconds=10) key = 'test_key' for i in range(5): self.assertTrue(limiter.is_allowed(key), f"Request {i+1} should be allowed") def test_blocks_requests_exceeding_limit(self): limiter = SimpleRateLimiter(max_requests=3, window_seconds=10) key = 'test_key' # First 3 requests should pass for i in range(3): self.assertTrue(limiter.is_allowed(key)) # 4th request should be blocked self.assertFalse(limiter.is_allowed(key)) def test_different_keys_are_independent(self): limiter = SimpleRateLimiter(max_requests=2, window_seconds=10) self.assertTrue(limiter.is_allowed('key_a')) self.assertTrue(limiter.is_allowed('key_a')) self.assertFalse(limiter.is_allowed('key_a')) # key_b should still be allowed self.assertTrue(limiter.is_allowed('key_b')) self.assertTrue(limiter.is_allowed('key_b')) self.assertFalse(limiter.is_allowed('key_b')) def test_window_sliding(self): limiter = SimpleRateLimiter(max_requests=2, window_seconds=1) key = 'test_key' # Use up the limit self.assertTrue(limiter.is_allowed(key)) self.assertTrue(limiter.is_allowed(key)) self.assertFalse(limiter.is_allowed(key)) # Wait for window to expire time.sleep(1.1) # Should be allowed again self.assertTrue(limiter.is_allowed(key)) def test_cleanup_removes_old_entries(self): limiter = SimpleRateLimiter(max_requests=5, window_seconds=10) limiter.is_allowed('key_1') limiter.is_allowed('key_2') limiter.is_allowed('key_3') self.assertEqual(len(limiter._requests), 3) # Cleanup with very short max_age should remove nothing (requests are recent) limiter.cleanup(max_age_seconds=10) self.assertEqual(len(limiter._requests), 3) # Wait and cleanup time.sleep(0.1) limiter.cleanup(max_age_seconds=0.05) self.assertEqual(len(limiter._requests), 0) if __name__ == '__main__': unittest.main()