tests/test_rate_limiter.py

"""
Test cases for rate limiter module.
Tests rate limiting functionality and edge cases.
"""

import unittest
import time
from src.rate_limiter import RateLimiter


class TestRateLimiter(unittest.TestCase):
    """Test cases for rate limiter."""
    
    def test_rate_limiter_allows_requests(self):
        """Test that rate limiter allows requests within limit."""
        limiter = RateLimiter(max_calls=5, period=60)
        
        @limiter
        def test_func():
            return "success"
        
        # Should allow first 5 calls
        for i in range(5):
            result = test_func()
            self.assertEqual(result, "success")
    
    def test_rate_limiter_blocks_excess(self):
        """Test that rate limiter blocks excess requests."""
        limiter = RateLimiter(max_calls=2, period=60)
        
        @limiter
        def test_func():
            return "success"
        
        # First 2 should succeed
        self.assertEqual(test_func(), "success")
        self.assertEqual(test_func(), "success")
        
        # Third should be blocked
        result = test_func()
        self.assertIn("error", result)
        self.assertIn("Rate limit exceeded", result["error"])
    
    def test_rate_limiter_resets_after_period(self):
        """Test that rate limiter resets after period."""
        limiter = RateLimiter(max_calls=2, period=1)  # 1 second period
        
        @limiter
        def test_func():
            return "success"
        
        # Use up limit
        test_func()
        test_func()
        
        # Should be blocked
        result = test_func()
        self.assertIn("error", result)
        
        # Wait for period to expire
        time.sleep(1.1)
        
        # Should work again
        result = test_func()
        self.assertEqual(result, "success")


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