feat(restruct tests):

tests/unit/executor/test_rate_limiter.py

@@ -0,0 +1,145 @@
+"""Tests for TokenBucketRateLimiter."""
+
+import asyncio
+import pytest
+import time
+from agentic_security.executor.rate_limiter import TokenBucketRateLimiter
+
+
+class TestTokenBucketRateLimiter:
+    """Test TokenBucketRateLimiter functionality."""
+
+    @pytest.mark.asyncio
+    async def test_initialization(self):
+        """Test rate limiter initialization."""
+        limiter = TokenBucketRateLimiter(rate=10, burst=20)
+
+        assert limiter.rate == 10
+        assert limiter.burst == 20
+        assert limiter.tokens == 20  # Starts full
+
+    @pytest.mark.asyncio
+    async def test_acquire_with_available_tokens(self):
+        """Test acquiring tokens when they're available."""
+        limiter = TokenBucketRateLimiter(rate=10, burst=5)
+
+        start = time.monotonic()
+        await limiter.acquire()
+        elapsed = time.monotonic() - start
+
+        # Should return immediately
+        assert elapsed < 0.1
+        assert limiter.tokens < 5  # One token consumed
+
+    @pytest.mark.asyncio
+    async def test_acquire_waits_when_no_tokens(self):
+        """Test that acquire waits when no tokens available."""
+        limiter = TokenBucketRateLimiter(rate=10, burst=1)
+
+        # Consume the initial token
+        await limiter.acquire()
+
+        # Next acquire should wait
+        start = time.monotonic()
+        await limiter.acquire()
+        elapsed = time.monotonic() - start
+
+        # Should wait approximately 1/rate seconds (0.1s for rate=10)
+        assert elapsed >= 0.08  # Allow some tolerance
+
+    @pytest.mark.asyncio
+    async def test_rate_limiting(self):
+        """Test that rate limiting actually limits request rate."""
+        limiter = TokenBucketRateLimiter(rate=10, burst=2)
+
+        # Make 5 requests
+        start = time.monotonic()
+        for _ in range(5):
+            await limiter.acquire()
+        elapsed = time.monotonic() - start
+
+        # With rate=10/s and burst=2:
+        # - First 2 requests are immediate (burst)
+        # - Next 3 requests require waiting: 3 * (1/10) = 0.3s
+        # Total should be around 0.3s
+        assert elapsed >= 0.25  # Allow some tolerance
+        assert elapsed < 0.5
+
+    @pytest.mark.asyncio
+    async def test_burst_capacity(self):
+        """Test that burst capacity allows immediate requests."""
+        limiter = TokenBucketRateLimiter(rate=5, burst=10)
+
+        # Make burst number of requests immediately
+        start = time.monotonic()
+        for _ in range(10):
+            await limiter.acquire()
+        elapsed = time.monotonic() - start
+
+        # All 10 requests should be nearly immediate (using burst capacity)
+        assert elapsed < 0.2
+
+    @pytest.mark.asyncio
+    async def test_token_replenishment(self):
+        """Test that tokens are replenished over time."""
+        limiter = TokenBucketRateLimiter(rate=10, burst=5)
+
+        # Consume all tokens
+        for _ in range(5):
+            await limiter.acquire()
+
+        assert limiter.tokens < 1
+
+        # Wait for tokens to replenish
+        await asyncio.sleep(0.3)  # Should add 3 tokens at rate=10
+
+        # Should have tokens again (approximately 3)
+        available = limiter.get_available_tokens()
+        assert available >= 2.5
+        assert available <= 3.5
+
+    @pytest.mark.asyncio
+    async def test_get_available_tokens(self):
+        """Test get_available_tokens method."""
+        limiter = TokenBucketRateLimiter(rate=10, burst=5)
+
+        # Initially full
+        assert limiter.get_available_tokens() == 5
+
+        # After consuming one
+        await limiter.acquire()
+        assert limiter.get_available_tokens() < 5
+
+    @pytest.mark.asyncio
+    async def test_concurrent_requests(self):
+        """Test rate limiter with concurrent requests."""
+        limiter = TokenBucketRateLimiter(rate=10, burst=3)
+
+        async def make_request(limiter):
+            await limiter.acquire()
+            return time.monotonic()
+
+        # Make 5 concurrent requests
+        start = time.monotonic()
+        tasks = [make_request(limiter) for _ in range(5)]
+        timestamps = await asyncio.gather(*tasks)
+        total_elapsed = time.monotonic() - start
+
+        # First 3 should be immediate (burst=3)
+        # Next 2 should wait
+        # Total time should be around 0.2s (2 * 1/10)
+        assert total_elapsed >= 0.15
+        assert total_elapsed < 0.4
+
+    @pytest.mark.asyncio
+    async def test_max_burst_capacity(self):
+        """Test that tokens don't exceed burst capacity."""
+        limiter = TokenBucketRateLimiter(rate=100, burst=5)
+
+        # Wait longer than needed to fill
+        await asyncio.sleep(0.2)  # Would add 20 tokens, but capped at 5
+
+        # Check tokens don't exceed burst
+        available = limiter.get_available_tokens()
+        assert available <= 5
+        assert available >= 4.5  # Close to full