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) def test_in_flight_slot_is_reusable_after_release(self): limiter = SimpleRateLimiter( max_requests=0, window_seconds=60, max_in_flight=2, ) key = 'GWS_A:query_order_detail' self.assertTrue(limiter.try_acquire(key)) self.assertTrue(limiter.try_acquire(key)) self.assertFalse(limiter.try_acquire(key)) limiter.release(key) self.assertTrue(limiter.try_acquire(key)) limiter.release(key) limiter.release(key) self.assertNotIn(key, limiter._in_flight) def test_zero_window_limit_disables_request_count_limit(self): limiter = SimpleRateLimiter( max_requests=0, window_seconds=60, max_in_flight=1, ) for _ in range(100): self.assertTrue(limiter.is_allowed('GWS_A:query_order_detail')) def test_zero_in_flight_limit_is_unbounded_and_unknown_release_is_safe(self): limiter = SimpleRateLimiter( max_requests=0, window_seconds=60, max_in_flight=0, ) self.assertTrue(limiter.try_acquire('GWS_A:query_order_detail')) limiter.release('GWS_A:query_order_detail') limiter.release('unknown') if __name__ == '__main__': unittest.main()