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)


if __name__ == '__main__':
    unittest.main()