mcp/tests/test_rate_limiter.py

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()