mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-24 06:09:55 +02:00
feat(restruct tests):
This commit is contained in:
Alexander Myasoedov
@@ -0,0 +1,209 @@
|
||||
"""Tests for CircuitBreaker."""
|
||||
|
||||
import time
|
||||
from agentic_security.executor.circuit_breaker import CircuitBreaker
|
||||
|
||||
|
||||
class TestCircuitBreaker:
|
||||
"""Test CircuitBreaker functionality."""
|
||||
|
||||
def test_initialization(self):
|
||||
"""Test circuit breaker initialization."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.5, recovery_timeout=30)
|
||||
|
||||
assert breaker.failure_threshold == 0.5
|
||||
assert breaker.recovery_timeout == 30
|
||||
assert breaker.state == "closed"
|
||||
assert breaker.failures == 0
|
||||
assert breaker.successes == 0
|
||||
|
||||
def test_record_success(self):
|
||||
"""Test recording successful requests."""
|
||||
breaker = CircuitBreaker()
|
||||
|
||||
breaker.record_success()
|
||||
assert breaker.successes == 1
|
||||
assert breaker.failures == 0
|
||||
assert breaker.state == "closed"
|
||||
|
||||
def test_record_failure(self):
|
||||
"""Test recording failed requests."""
|
||||
breaker = CircuitBreaker()
|
||||
|
||||
breaker.record_failure()
|
||||
assert breaker.failures == 1
|
||||
assert breaker.successes == 0
|
||||
assert breaker.last_failure_time is not None
|
||||
|
||||
def test_circuit_opens_on_failure_threshold(self):
|
||||
"""Test that circuit opens when failure threshold is exceeded."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.5, recovery_timeout=30)
|
||||
|
||||
# Record 10 requests: 6 failures, 4 successes (60% failure rate)
|
||||
for _ in range(4):
|
||||
breaker.record_success()
|
||||
for _ in range(6):
|
||||
breaker.record_failure()
|
||||
|
||||
# Circuit should be open (60% > 50% threshold)
|
||||
assert breaker.state == "open"
|
||||
assert breaker.is_open() is True
|
||||
|
||||
def test_circuit_stays_closed_below_threshold(self):
|
||||
"""Test that circuit stays closed when below threshold."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.5, recovery_timeout=30)
|
||||
|
||||
# Record 10 requests: 4 failures, 6 successes (40% failure rate)
|
||||
for _ in range(6):
|
||||
breaker.record_success()
|
||||
for _ in range(4):
|
||||
breaker.record_failure()
|
||||
|
||||
# Circuit should stay closed (40% < 50% threshold)
|
||||
assert breaker.state == "closed"
|
||||
assert breaker.is_open() is False
|
||||
|
||||
def test_minimum_sample_size_required(self):
|
||||
"""Test that minimum sample size is required before opening."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.5)
|
||||
|
||||
# Only 5 failures (below minimum of 10 total requests)
|
||||
for _ in range(5):
|
||||
breaker.record_failure()
|
||||
|
||||
# Circuit should stay closed (not enough samples)
|
||||
assert breaker.state == "closed"
|
||||
assert breaker.is_open() is False
|
||||
|
||||
def test_circuit_recovery_after_timeout(self):
|
||||
"""Test that circuit enters half-open state after recovery timeout."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.5, recovery_timeout=1)
|
||||
|
||||
# Open the circuit
|
||||
for _ in range(4):
|
||||
breaker.record_success()
|
||||
for _ in range(6):
|
||||
breaker.record_failure()
|
||||
|
||||
assert breaker.state == "open"
|
||||
|
||||
# Wait for recovery timeout
|
||||
time.sleep(1.1)
|
||||
|
||||
# Check if circuit moves to half-open
|
||||
is_open = breaker.is_open()
|
||||
assert is_open is False
|
||||
assert breaker.state == "half_open"
|
||||
|
||||
def test_half_open_to_closed_on_successes(self):
|
||||
"""Test that circuit closes from half-open after enough successes."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.5, recovery_timeout=1)
|
||||
|
||||
# Open the circuit
|
||||
for _ in range(4):
|
||||
breaker.record_success()
|
||||
for _ in range(6):
|
||||
breaker.record_failure()
|
||||
|
||||
# Wait for recovery
|
||||
time.sleep(1.1)
|
||||
breaker.is_open() # Triggers transition to half-open
|
||||
|
||||
assert breaker.state == "half_open"
|
||||
|
||||
# Record 3 successes
|
||||
breaker.record_success()
|
||||
breaker.record_success()
|
||||
breaker.record_success()
|
||||
|
||||
# Should transition to closed
|
||||
assert breaker.state == "closed"
|
||||
|
||||
def test_get_state(self):
|
||||
"""Test get_state method."""
|
||||
breaker = CircuitBreaker()
|
||||
|
||||
assert breaker.get_state() == "closed"
|
||||
|
||||
# Open the circuit
|
||||
for _ in range(10):
|
||||
breaker.record_failure()
|
||||
|
||||
assert breaker.get_state() == "open"
|
||||
|
||||
def test_get_failure_rate(self):
|
||||
"""Test get_failure_rate method."""
|
||||
breaker = CircuitBreaker()
|
||||
|
||||
# No requests
|
||||
assert breaker.get_failure_rate() == 0.0
|
||||
|
||||
# 3 failures, 7 successes (30% failure rate)
|
||||
for _ in range(7):
|
||||
breaker.record_success()
|
||||
for _ in range(3):
|
||||
breaker.record_failure()
|
||||
|
||||
assert breaker.get_failure_rate() == 0.3
|
||||
|
||||
def test_reset(self):
|
||||
"""Test reset method."""
|
||||
breaker = CircuitBreaker()
|
||||
|
||||
# Record some activity
|
||||
breaker.record_success()
|
||||
breaker.record_failure()
|
||||
for _ in range(10):
|
||||
breaker.record_failure()
|
||||
|
||||
# Reset
|
||||
breaker.reset()
|
||||
|
||||
# Should be back to initial state
|
||||
assert breaker.state == "closed"
|
||||
assert breaker.failures == 0
|
||||
assert breaker.successes == 0
|
||||
assert breaker.last_failure_time is None
|
||||
|
||||
def test_exact_failure_threshold(self):
|
||||
"""Test behavior at exact failure threshold."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.5)
|
||||
|
||||
# Exactly 50% failure rate (5 failures, 5 successes)
|
||||
for _ in range(5):
|
||||
breaker.record_success()
|
||||
for _ in range(5):
|
||||
breaker.record_failure()
|
||||
|
||||
# Should be open (>= threshold)
|
||||
assert breaker.state == "open"
|
||||
|
||||
def test_high_failure_threshold(self):
|
||||
"""Test with high failure threshold."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.9)
|
||||
|
||||
# 80% failure rate (8 failures, 2 successes)
|
||||
for _ in range(2):
|
||||
breaker.record_success()
|
||||
for _ in range(8):
|
||||
breaker.record_failure()
|
||||
|
||||
# Should stay closed (80% < 90%)
|
||||
assert breaker.state == "closed"
|
||||
|
||||
def test_zero_recovery_timeout(self):
|
||||
"""Test with zero recovery timeout."""
|
||||
breaker = CircuitBreaker(failure_threshold=0.5, recovery_timeout=0)
|
||||
|
||||
# Open the circuit
|
||||
for _ in range(10):
|
||||
breaker.record_failure()
|
||||
|
||||
assert breaker.state == "open"
|
||||
|
||||
# Should immediately allow recovery attempt
|
||||
time.sleep(0.01)
|
||||
is_open = breaker.is_open()
|
||||
|
||||
assert is_open is False
|
||||
assert breaker.state == "half_open"
|
||||
@@ -0,0 +1,279 @@
|
||||
"""Tests for ConcurrentExecutor."""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
from unittest.mock import Mock, patch
|
||||
from agentic_security.executor.concurrent import ConcurrentExecutor, ExecutorMetrics
|
||||
from agentic_security.probe_actor.state import FuzzerState
|
||||
|
||||
|
||||
class TestExecutorMetrics:
|
||||
"""Test ExecutorMetrics functionality."""
|
||||
|
||||
def test_initialization(self):
|
||||
"""Test metrics initialization."""
|
||||
metrics = ExecutorMetrics()
|
||||
|
||||
assert metrics.successful_requests == 0
|
||||
assert metrics.failed_requests == 0
|
||||
assert metrics.total_latency == 0.0
|
||||
assert len(metrics.latencies) == 0
|
||||
|
||||
def test_record_success(self):
|
||||
"""Test recording successful requests."""
|
||||
metrics = ExecutorMetrics()
|
||||
|
||||
metrics.record_success(0.5)
|
||||
metrics.record_success(0.3)
|
||||
|
||||
assert metrics.successful_requests == 2
|
||||
assert metrics.total_latency == 0.8
|
||||
assert len(metrics.latencies) == 2
|
||||
|
||||
def test_record_failure(self):
|
||||
"""Test recording failed requests."""
|
||||
metrics = ExecutorMetrics()
|
||||
|
||||
metrics.record_failure()
|
||||
metrics.record_failure()
|
||||
|
||||
assert metrics.failed_requests == 2
|
||||
assert metrics.successful_requests == 0
|
||||
|
||||
def test_get_stats_no_requests(self):
|
||||
"""Test get_stats with no requests."""
|
||||
metrics = ExecutorMetrics()
|
||||
|
||||
stats = metrics.get_stats()
|
||||
|
||||
assert stats["total_requests"] == 0
|
||||
assert stats["success_rate"] == 0.0
|
||||
assert stats["avg_latency_ms"] == 0.0
|
||||
assert stats["p95_latency_ms"] == 0.0
|
||||
|
||||
def test_get_stats_with_requests(self):
|
||||
"""Test get_stats with recorded requests."""
|
||||
metrics = ExecutorMetrics()
|
||||
|
||||
# Record some requests
|
||||
metrics.record_success(0.1) # 100ms
|
||||
metrics.record_success(0.2) # 200ms
|
||||
metrics.record_success(0.3) # 300ms
|
||||
metrics.record_failure()
|
||||
|
||||
stats = metrics.get_stats()
|
||||
|
||||
assert stats["total_requests"] == 4
|
||||
assert stats["successful_requests"] == 3
|
||||
assert stats["failed_requests"] == 1
|
||||
assert stats["success_rate"] == 0.75
|
||||
assert stats["avg_latency_ms"] == pytest.approx(200.0, rel=0.01)
|
||||
|
||||
def test_get_stats_p95_latency(self):
|
||||
"""Test p95 latency calculation."""
|
||||
metrics = ExecutorMetrics()
|
||||
|
||||
# Add 100 requests with varying latencies
|
||||
for i in range(100):
|
||||
metrics.record_success(i * 0.001) # 0ms to 99ms
|
||||
|
||||
stats = metrics.get_stats()
|
||||
|
||||
# p95 should be around 95ms
|
||||
assert stats["p95_latency_ms"] >= 90.0
|
||||
assert stats["p95_latency_ms"] <= 100.0
|
||||
|
||||
|
||||
class TestConcurrentExecutor:
|
||||
"""Test ConcurrentExecutor functionality."""
|
||||
|
||||
def test_initialization(self):
|
||||
"""Test executor initialization."""
|
||||
executor = ConcurrentExecutor(
|
||||
max_concurrent=20,
|
||||
rate_limit=10,
|
||||
burst=5,
|
||||
failure_threshold=0.5,
|
||||
recovery_timeout=30,
|
||||
)
|
||||
|
||||
assert executor.semaphore._value == 20
|
||||
assert executor.rate_limiter.rate == 10
|
||||
assert executor.rate_limiter.burst == 5
|
||||
assert executor.circuit_breaker.failure_threshold == 0.5
|
||||
assert executor.circuit_breaker.recovery_timeout == 30
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_batch_success(self):
|
||||
"""Test successful batch execution."""
|
||||
executor = ConcurrentExecutor(max_concurrent=10, rate_limit=100, burst=10)
|
||||
fuzzer_state = FuzzerState()
|
||||
|
||||
# Mock request factory
|
||||
request_factory = Mock()
|
||||
|
||||
# Mock process_prompt to return success
|
||||
async def mock_process_prompt(rf, prompt, tokens, module, state):
|
||||
return (10, False) # 10 tokens, not refused
|
||||
|
||||
with patch(
|
||||
"agentic_security.probe_actor.fuzzer.process_prompt",
|
||||
side_effect=mock_process_prompt,
|
||||
):
|
||||
prompts = ["prompt1", "prompt2", "prompt3"]
|
||||
tokens, failures = await executor.execute_batch(
|
||||
request_factory, prompts, "test_module", fuzzer_state
|
||||
)
|
||||
|
||||
assert tokens == 30 # 3 prompts * 10 tokens
|
||||
assert failures == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_batch_with_failures(self):
|
||||
"""Test batch execution with some failures."""
|
||||
executor = ConcurrentExecutor(max_concurrent=10, rate_limit=100, burst=10)
|
||||
fuzzer_state = FuzzerState()
|
||||
|
||||
request_factory = Mock()
|
||||
|
||||
# Mock process_prompt to alternate success/failure
|
||||
call_count = [0]
|
||||
|
||||
async def mock_process_prompt(rf, prompt, tokens, module, state):
|
||||
call_count[0] += 1
|
||||
if call_count[0] % 2 == 0:
|
||||
return (10, True) # Refused
|
||||
return (10, False) # Success
|
||||
|
||||
with patch(
|
||||
"agentic_security.probe_actor.fuzzer.process_prompt",
|
||||
side_effect=mock_process_prompt,
|
||||
):
|
||||
prompts = ["p1", "p2", "p3", "p4"]
|
||||
tokens, failures = await executor.execute_batch(
|
||||
request_factory, prompts, "test_module", fuzzer_state
|
||||
)
|
||||
|
||||
assert tokens == 40 # 4 prompts * 10 tokens
|
||||
assert failures == 2 # 2 refused
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_batch_respects_concurrency_limit(self):
|
||||
"""Test that concurrency limit is respected."""
|
||||
executor = ConcurrentExecutor(max_concurrent=2, rate_limit=100, burst=10)
|
||||
fuzzer_state = FuzzerState()
|
||||
|
||||
request_factory = Mock()
|
||||
|
||||
# Track concurrent executions
|
||||
concurrent_count = [0]
|
||||
max_concurrent = [0]
|
||||
|
||||
async def mock_process_prompt(rf, prompt, tokens, module, state):
|
||||
concurrent_count[0] += 1
|
||||
max_concurrent[0] = max(max_concurrent[0], concurrent_count[0])
|
||||
await asyncio.sleep(0.01) # Simulate work
|
||||
concurrent_count[0] -= 1
|
||||
return (10, False)
|
||||
|
||||
with patch(
|
||||
"agentic_security.probe_actor.fuzzer.process_prompt",
|
||||
side_effect=mock_process_prompt,
|
||||
):
|
||||
prompts = ["p1", "p2", "p3", "p4", "p5"]
|
||||
await executor.execute_batch(
|
||||
request_factory, prompts, "test_module", fuzzer_state
|
||||
)
|
||||
|
||||
# Max concurrent should not exceed limit
|
||||
assert max_concurrent[0] <= 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_circuit_breaker_integration(self):
|
||||
"""Test that circuit breaker opens on failures."""
|
||||
executor = ConcurrentExecutor(
|
||||
max_concurrent=10,
|
||||
rate_limit=100,
|
||||
burst=20,
|
||||
failure_threshold=0.5,
|
||||
recovery_timeout=1,
|
||||
)
|
||||
fuzzer_state = FuzzerState()
|
||||
request_factory = Mock()
|
||||
|
||||
# Mock process_prompt to always fail
|
||||
async def mock_process_prompt_fail(rf, prompt, tokens, module, state):
|
||||
raise Exception("Request failed")
|
||||
|
||||
# First batch - all failures
|
||||
with patch(
|
||||
"agentic_security.probe_actor.fuzzer.process_prompt",
|
||||
side_effect=mock_process_prompt_fail,
|
||||
):
|
||||
prompts = ["p1", "p2", "p3", "p4", "p5", "p6", "p7", "p8", "p9", "p10"]
|
||||
tokens, failures = await executor.execute_batch(
|
||||
request_factory, prompts, "test_module", fuzzer_state
|
||||
)
|
||||
|
||||
# All should have failed
|
||||
assert failures == 10
|
||||
|
||||
# Circuit should be open now
|
||||
assert executor.circuit_breaker.state == "open"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_metrics(self):
|
||||
"""Test getting executor metrics."""
|
||||
executor = ConcurrentExecutor(max_concurrent=10, rate_limit=100, burst=10)
|
||||
fuzzer_state = FuzzerState()
|
||||
request_factory = Mock()
|
||||
|
||||
async def mock_process_prompt(rf, prompt, tokens, module, state):
|
||||
return (10, False)
|
||||
|
||||
with patch(
|
||||
"agentic_security.probe_actor.fuzzer.process_prompt",
|
||||
side_effect=mock_process_prompt,
|
||||
):
|
||||
await executor.execute_batch(
|
||||
request_factory, ["p1", "p2"], "test_module", fuzzer_state
|
||||
)
|
||||
|
||||
metrics = executor.get_metrics()
|
||||
|
||||
assert "total_requests" in metrics
|
||||
assert "success_rate" in metrics
|
||||
assert "circuit_breaker_state" in metrics
|
||||
assert "available_tokens" in metrics
|
||||
assert metrics["total_requests"] == 2
|
||||
assert metrics["circuit_breaker_state"] == "closed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limiting_applied(self):
|
||||
"""Test that rate limiting is applied."""
|
||||
executor = ConcurrentExecutor(max_concurrent=10, rate_limit=5, burst=2)
|
||||
fuzzer_state = FuzzerState()
|
||||
request_factory = Mock()
|
||||
|
||||
async def mock_process_prompt(rf, prompt, tokens, module, state):
|
||||
return (10, False)
|
||||
|
||||
import time
|
||||
|
||||
with patch(
|
||||
"agentic_security.probe_actor.fuzzer.process_prompt",
|
||||
side_effect=mock_process_prompt,
|
||||
):
|
||||
start = time.monotonic()
|
||||
# 5 requests with rate=5/s and burst=2
|
||||
# First 2 immediate, next 3 should take ~0.6s total
|
||||
await executor.execute_batch(
|
||||
request_factory,
|
||||
["p1", "p2", "p3", "p4", "p5"],
|
||||
"test_module",
|
||||
fuzzer_state,
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
|
||||
# Should take at least 0.5s (3 requests / 5 per second)
|
||||
assert elapsed >= 0.4
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user