44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
"""Tests for rate limiting middleware."""
|
|
|
|
import pytest
|
|
|
|
from web.api.middleware import TokenBucket
|
|
|
|
|
|
class TestTokenBucket:
|
|
"""Tests for TokenBucket rate limiter."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_initial_burst(self) -> None:
|
|
"""Test that initial burst tokens are available."""
|
|
bucket = TokenBucket(rate=10.0, burst=5)
|
|
|
|
# Should be able to acquire burst tokens
|
|
for _ in range(5):
|
|
allowed, _ = await bucket.acquire()
|
|
assert allowed is True
|
|
|
|
# Next one should fail
|
|
allowed, _ = await bucket.acquire()
|
|
assert allowed is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_returns_retry_after_when_empty(self) -> None:
|
|
"""Test that acquire returns positive retry_after when empty."""
|
|
bucket = TokenBucket(rate=1.0, burst=1)
|
|
|
|
# Drain the bucket
|
|
await bucket.acquire()
|
|
|
|
allowed, retry_after = await bucket.acquire()
|
|
assert allowed is False
|
|
assert retry_after > 0.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acquire_returns_zero_retry_when_available(self) -> None:
|
|
"""Test that acquire returns 0 retry_after when tokens available."""
|
|
bucket = TokenBucket(rate=10.0, burst=5)
|
|
|
|
allowed, retry_after = await bucket.acquire()
|
|
assert allowed is True
|
|
assert retry_after == 0.0
|