MCP Hub

"""
Test suite for Rate Limiter (Phase 7.3)

Tests the Token Bucket algorithm, rate limiting logic,
and middleware integration.

Run with: python test_rate_limiter.py
"""

import time
import unittest
from unittest.mock import patch

from core.rate_limiter import (
    ClientRateLimitState,
    RateLimitConfig,
    RateLimiter,
    TokenBucket,
    get_rate_limiter,
)


class TestTokenBucket(unittest.TestCase):
    """Test Token Bucket algorithm implementation."""

    def test_initialization(self):
        """Test bucket initializes with full capacity."""
        bucket = TokenBucket(capacity=10, refill_rate=1.0)
        self.assertEqual(bucket.capacity, 10)
        self.assertEqual(bucket.refill_rate, 1.0)
        self.assertEqual(bucket.tokens, 10.0)

    def test_consume_tokens(self):
        """Test consuming tokens from bucket."""
        bucket = TokenBucket(capacity=10, refill_rate=1.0)

        # Should successfully consume 5 tokens
        self.assertTrue(bucket.consume(5))
        self.assertEqual(bucket.get_available_tokens(), 5)

        # Should successfully consume 5 more tokens
        self.assertTrue(bucket.consume(5))
        self.assertEqual(bucket.get_available_tokens(), 0)

        # Should fail to consume when empty
        self.assertFalse(bucket.consume(1))

    def test_token_refill(self):
        """Test tokens refill over time."""
        bucket = TokenBucket(capacity=10, refill_rate=10.0)  # 10 tokens/second

        # Consume all tokens
        bucket.consume(10)
        self.assertEqual(bucket.get_available_tokens(), 0)

        # Wait 0.5 seconds, should refill 5 tokens
        time.sleep(0.5)
        available = bucket.get_available_tokens()
        # Allow some tolerance for timing variations
        self.assertGreaterEqual(available, 4)
        self.assertLessEqual(available, 6)

    def test_refill_cap(self):
        """Test refill doesn't exceed capacity."""
        bucket = TokenBucket(capacity=10, refill_rate=10.0)

        # Wait long enough to refill many times over
        time.sleep(2.0)

        # Should be capped at capacity
        self.assertEqual(bucket.get_available_tokens(), 10)

    def test_wait_time_calculation(self):
        """Test wait time calculation when tokens unavailable."""
        bucket = TokenBucket(capacity=10, refill_rate=2.0)  # 2 tokens/second

        # Consume all tokens
        bucket.consume(10)

        # Need 1 token, should take 0.5 seconds
        wait_time = bucket.get_wait_time(1)
        self.assertAlmostEqual(wait_time, 0.5, delta=0.1)

        # Need 10 tokens, should take 5 seconds
        wait_time = bucket.get_wait_time(10)
        self.assertAlmostEqual(wait_time, 5.0, delta=0.1)


class TestRateLimitConfig(unittest.TestCase):
    """Test rate limit configuration."""

    def test_default_values(self):
        """Test default configuration values."""
        config = RateLimitConfig()
        self.assertEqual(config.per_minute, 60)
        self.assertEqual(config.per_hour, 1000)
        self.assertEqual(config.per_day, 10000)

    def test_custom_values(self):
        """Test custom configuration values."""
        config = RateLimitConfig(per_minute=100, per_hour=2000, per_day=20000)
        self.assertEqual(config.per_minute, 100)
        self.assertEqual(config.per_hour, 2000)
        self.assertEqual(config.per_day, 20000)

    def test_from_env(self):
        """Test loading configuration from environment."""
        with patch.dict(
            "os.environ",
            {
                "RATE_LIMIT_PER_MINUTE": "120",
                "RATE_LIMIT_PER_HOUR": "2000",
                "RATE_LIMIT_PER_DAY": "15000",
            },
        ):
            config = RateLimitConfig.from_env()
            self.assertEqual(config.per_minute, 120)
            self.assertEqual(config.per_hour, 2000)
            self.assertEqual(config.per_day, 15000)

    def test_from_env_with_prefix(self):
        """Test loading configuration with prefix."""
        with patch.dict(
            "os.environ",
            {
                "WORDPRESS_RATE_LIMIT_PER_MINUTE": "80",
                "WORDPRESS_RATE_LIMIT_PER_HOUR": "1500",
                "WORDPRESS_RATE_LIMIT_PER_DAY": "12000",
            },
        ):
            config = RateLimitConfig.from_env(prefix="WORDPRESS")
            self.assertEqual(config.per_minute, 80)
            self.assertEqual(config.per_hour, 1500)
            self.assertEqual(config.per_day, 12000)


class TestClientRateLimitState(unittest.TestCase):
    """Test client rate limit state tracking."""

    def test_initialization(self):
        """Test client state initialization."""
        state = ClientRateLimitState(
            client_id="test_client",
            minute_bucket=TokenBucket(60, 1.0),
            hour_bucket=TokenBucket(1000, 1.0),
            day_bucket=TokenBucket(10000, 1.0),
        )
        self.assertEqual(state.client_id, "test_client")
        self.assertEqual(state.total_requests, 0)
        self.assertEqual(state.rejected_requests, 0)

    def test_successful_request(self):
        """Test successful request consumes tokens."""
        state = ClientRateLimitState(
            client_id="test_client",
            minute_bucket=TokenBucket(60, 1.0),
            hour_bucket=TokenBucket(1000, 1.0),
            day_bucket=TokenBucket(10000, 1.0),
        )

        allowed, message, retry_after = state.check_and_consume()
        self.assertTrue(allowed)
        self.assertEqual(message, "")
        self.assertEqual(retry_after, 0.0)
        self.assertEqual(state.total_requests, 1)
        self.assertEqual(state.rejected_requests, 0)

    def test_minute_limit_exceeded(self):
        """Test rejection when per-minute limit exceeded."""
        state = ClientRateLimitState(
            client_id="test_client",
            minute_bucket=TokenBucket(2, 1.0),  # Only 2 requests/minute
            hour_bucket=TokenBucket(1000, 1.0),
            day_bucket=TokenBucket(10000, 1.0),
        )

        # First 2 requests should succeed
        state.check_and_consume()
        state.check_and_consume()

        # Third request should fail
        allowed, message, retry_after = state.check_and_consume()
        self.assertFalse(allowed)
        self.assertIn("per minute", message.lower())
        self.assertGreater(retry_after, 0)
        self.assertEqual(state.rejected_requests, 1)

    def test_hour_limit_exceeded(self):
        """Test rejection when per-hour limit exceeded."""
        state = ClientRateLimitState(
            client_id="test_client",
            minute_bucket=TokenBucket(100, 100.0),  # High minute limit
            hour_bucket=TokenBucket(2, 1.0),  # Only 2 requests/hour
            day_bucket=TokenBucket(10000, 1.0),
        )

        # First 2 requests should succeed
        state.check_and_consume()
        state.check_and_consume()

        # Third request should fail at hour limit
        allowed, message, retry_after = state.check_and_consume()
        self.assertFalse(allowed)
        self.assertIn("per hour", message.lower())

    def test_day_limit_exceeded(self):
        """Test rejection when daily limit exceeded."""
        state = ClientRateLimitState(
            client_id="test_client",
            minute_bucket=TokenBucket(100, 100.0),
            hour_bucket=TokenBucket(100, 100.0),
            day_bucket=TokenBucket(2, 1.0),  # Only 2 requests/day
        )

        # First 2 requests should succeed
        state.check_and_consume()
        state.check_and_consume()

        # Third request should fail at day limit
        allowed, message, retry_after = state.check_and_consume()
        self.assertFalse(allowed)
        self.assertIn("daily", message.lower())

    def test_token_refund_on_rejection(self):
        """Test tokens are refunded when request rejected at later stage."""
        state = ClientRateLimitState(
            client_id="test_client",
            minute_bucket=TokenBucket(10, 1.0),
            hour_bucket=TokenBucket(2, 1.0),  # Will fail here
            day_bucket=TokenBucket(10, 1.0),
        )

        # Consume hour bucket
        state.check_and_consume()
        state.check_and_consume()

        # Check minute bucket before rejection
        minute_before = state.minute_bucket.get_available_tokens()

        # This should fail at hour limit
        allowed, message, retry_after = state.check_and_consume()
        self.assertFalse(allowed)

        # Minute bucket should be refunded
        minute_after = state.minute_bucket.get_available_tokens()
        self.assertEqual(minute_before, minute_after)

    def test_get_stats(self):
        """Test statistics generation."""
        state = ClientRateLimitState(
            client_id="test_client",
            minute_bucket=TokenBucket(60, 1.0),
            hour_bucket=TokenBucket(1000, 1.0),
            day_bucket=TokenBucket(10000, 1.0),
        )

        # Make some requests
        state.check_and_consume()
        state.check_and_consume()

        stats = state.get_stats()
        self.assertEqual(stats["client_id"], "test_client")
        self.assertEqual(stats["total_requests"], 2)
        self.assertEqual(stats["rejected_requests"], 0)
        self.assertEqual(stats["success_rate"], 1.0)
        self.assertIn("available_tokens", stats)
        self.assertIn("limits", stats)


class TestRateLimiter(unittest.TestCase):
    """Test RateLimiter class."""

    def setUp(self):
        """Create a fresh rate limiter for each test."""
        # Reset singleton
        import core.rate_limiter

        core.rate_limiter._rate_limiter = None
        self.limiter = RateLimiter()

    def test_initialization(self):
        """Test rate limiter initialization."""
        self.assertIsNotNone(self.limiter.default_config)
        self.assertEqual(len(self.limiter.clients), 0)
        self.assertEqual(self.limiter.global_stats["total_requests"], 0)

    def test_singleton_pattern(self):
        """Test get_rate_limiter returns singleton."""
        limiter1 = get_rate_limiter()
        limiter2 = get_rate_limiter()
        self.assertIs(limiter1, limiter2)

    def test_client_creation(self):
        """Test automatic client state creation."""
        allowed, message, retry_after = self.limiter.check_rate_limit(
            client_id="client1", tool_name="test_tool"
        )
        self.assertTrue(allowed)
        self.assertIn("client1", self.limiter.clients)

    def test_successful_rate_limit_check(self):
        """Test successful rate limit check."""
        allowed, message, retry_after = self.limiter.check_rate_limit(
            client_id="client1", tool_name="test_tool"
        )
        self.assertTrue(allowed)
        self.assertEqual(message, "")
        self.assertEqual(retry_after, 0.0)
        self.assertEqual(self.limiter.global_stats["total_requests"], 1)
        self.assertEqual(self.limiter.global_stats["total_rejected"], 0)

    def test_rate_limit_exceeded(self):
        """Test rate limit exceeded scenario."""
        # Configure very low limits
        self.limiter.configure_limits("test_plugin", per_minute=2, per_hour=2, per_day=2)

        # First 2 requests should succeed
        self.limiter.check_rate_limit("client1", "test", "test_plugin")
        self.limiter.check_rate_limit("client1", "test", "test_plugin")

        # Third should fail
        allowed, message, retry_after = self.limiter.check_rate_limit(
            "client1", "test", "test_plugin"
        )
        self.assertFalse(allowed)
        self.assertGreater(retry_after, 0)
        self.assertEqual(self.limiter.global_stats["total_rejected"], 1)

    def test_plugin_specific_limits(self):
        """Test plugin-specific rate limits."""
        # Configure different limits for wordpress
        self.limiter.configure_limits("wordpress", per_minute=5, per_hour=5, per_day=5)

        # Make 5 wordpress requests
        for _i in range(5):
            allowed, _, _ = self.limiter.check_rate_limit("client1", "wordpress_tool", "wordpress")
            self.assertTrue(allowed)

        # 6th should fail
        allowed, message, _ = self.limiter.check_rate_limit(
            "client1", "wordpress_tool", "wordpress"
        )
        self.assertFalse(allowed)

    def test_multiple_clients(self):
        """Test multiple clients tracked separately."""
        # Configure low limits
        self.limiter.configure_limits("test", per_minute=2, per_hour=2, per_day=2)

        # Client 1 makes 2 requests
        self.limiter.check_rate_limit("client1", "test", "test")
        self.limiter.check_rate_limit("client1", "test", "test")

        # Client 1's third request should fail
        allowed, _, _ = self.limiter.check_rate_limit("client1", "test", "test")
        self.assertFalse(allowed)

        # But client 2's first request should succeed
        allowed, _, _ = self.limiter.check_rate_limit("client2", "test", "test")
        self.assertTrue(allowed)

    def test_get_client_stats(self):
        """Test getting client statistics."""
        # Make some requests
        self.limiter.check_rate_limit("client1", "test")
        self.limiter.check_rate_limit("client1", "test")

        stats = self.limiter.get_client_stats("client1")
        self.assertIsNotNone(stats)
        self.assertEqual(stats["client_id"], "client1")
        self.assertEqual(stats["total_requests"], 2)

    def test_get_client_stats_not_found(self):
        """Test getting stats for non-existent client."""
        stats = self.limiter.get_client_stats("nonexistent")
        self.assertIsNone(stats)

    def test_get_all_stats(self):
        """Test getting all statistics."""
        # Make requests from multiple clients
        self.limiter.check_rate_limit("client1", "test")
        self.limiter.check_rate_limit("client2", "test")

        stats = self.limiter.get_all_stats()
        self.assertIn("global", stats)
        self.assertIn("default_limits", stats)
        self.assertIn("plugin_limits", stats)
        self.assertIn("clients", stats)
        self.assertEqual(stats["global"]["total_requests"], 2)
        self.assertEqual(len(stats["clients"]), 2)

    def test_reset_client(self):
        """Test resetting specific client."""
        # Make some requests
        self.limiter.check_rate_limit("client1", "test")
        self.assertIn("client1", self.limiter.clients)

        # Reset client
        success = self.limiter.reset_client("client1")
        self.assertTrue(success)
        self.assertNotIn("client1", self.limiter.clients)

    def test_reset_client_not_found(self):
        """Test resetting non-existent client."""
        success = self.limiter.reset_client("nonexistent")
        self.assertFalse(success)

    def test_reset_all(self):
        """Test resetting all clients."""
        # Make requests from multiple clients
        self.limiter.check_rate_limit("client1", "test")
        self.limiter.check_rate_limit("client2", "test")
        self.limiter.check_rate_limit("client3", "test")

        # Reset all
        count = self.limiter.reset_all()
        self.assertEqual(count, 3)
        self.assertEqual(len(self.limiter.clients), 0)
        self.assertEqual(self.limiter.global_stats["total_requests"], 0)

    def test_configure_limits(self):
        """Test dynamic limit configuration."""
        self.limiter.configure_limits("custom_plugin", per_minute=100, per_hour=2000, per_day=20000)

        config = self.limiter.plugin_configs["custom_plugin"]
        self.assertEqual(config.per_minute, 100)
        self.assertEqual(config.per_hour, 2000)
        self.assertEqual(config.per_day, 20000)


class TestIntegration(unittest.TestCase):
    """Integration tests simulating real-world usage."""

    def setUp(self):
        """Create fresh rate limiter."""
        import core.rate_limiter

        core.rate_limiter._rate_limiter = None
        self.limiter = RateLimiter()

    def test_burst_traffic(self):
        """Test handling of burst traffic within limits."""
        # Configure to allow 10/minute
        self.limiter.configure_limits("test", per_minute=10, per_hour=100, per_day=1000)

        # Simulate burst of 10 requests
        success_count = 0
        for _i in range(10):
            allowed, _, _ = self.limiter.check_rate_limit("client1", "test", "test")
            if allowed:
                success_count += 1

        # All 10 should succeed (burst allowed)
        self.assertEqual(success_count, 10)

        # 11th should fail
        allowed, _, _ = self.limiter.check_rate_limit("client1", "test", "test")
        self.assertFalse(allowed)

    def test_sustained_traffic(self):
        """Test handling of sustained traffic with refill."""
        # Configure high refill rate for testing
        self.limiter.configure_limits("test", per_minute=5, per_hour=100, per_day=1000)

        # Make 5 requests (consume all minute tokens)
        for _i in range(5):
            self.limiter.check_rate_limit("client1", "test", "test")

        # Should be rate limited now
        allowed, _, _ = self.limiter.check_rate_limit("client1", "test", "test")
        self.assertFalse(allowed)

        # Wait for refill (should get ~2-3 tokens in 0.5 seconds at 5/min rate)
        time.sleep(0.5)

        # Should be able to make 1-2 more requests
        allowed, _, _ = self.limiter.check_rate_limit("client1", "test", "test")
        # May or may not succeed depending on exact timing
        # Just verify no crash and reasonable behavior


def run_tests():
    """Run all tests."""
    # Create test suite
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()

    # Add all test classes
    suite.addTests(loader.loadTestsFromTestCase(TestTokenBucket))
    suite.addTests(loader.loadTestsFromTestCase(TestRateLimitConfig))
    suite.addTests(loader.loadTestsFromTestCase(TestClientRateLimitState))
    suite.addTests(loader.loadTestsFromTestCase(TestRateLimiter))
    suite.addTests(loader.loadTestsFromTestCase(TestIntegration))

    # Run tests
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)

    # Return exit code
    return 0 if result.wasSuccessful() else 1


if __name__ == "__main__":
    import sys

    sys.exit(run_tests())