67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
|
"""Tests for async rate limiter."""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import pytest
|
||
|
|
import time
|
||
|
|
from aucourt_ingest.utils.rate_limiter import RateLimiter, RateLimiterRegistry
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_single_request_passes_immediately():
|
||
|
|
rl = RateLimiter(rps=1.0, concurrent=1)
|
||
|
|
t0 = time.monotonic()
|
||
|
|
async with rl:
|
||
|
|
pass
|
||
|
|
elapsed = time.monotonic() - t0
|
||
|
|
assert elapsed < 0.5 # should be near-instant
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_rate_limit_delays_rapid_requests():
|
||
|
|
rl = RateLimiter(rps=10.0, concurrent=1) # 100ms between requests
|
||
|
|
t0 = time.monotonic()
|
||
|
|
async with rl:
|
||
|
|
pass
|
||
|
|
async with rl:
|
||
|
|
pass
|
||
|
|
elapsed = time.monotonic() - t0
|
||
|
|
# Two requests at 10 rps should take ~100ms minimum
|
||
|
|
assert elapsed >= 0.05
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_concurrency_limit():
|
||
|
|
rl = RateLimiter(rps=1000.0, concurrent=1) # high rps, low concurrency
|
||
|
|
entered = 0
|
||
|
|
max_concurrent = 0
|
||
|
|
|
||
|
|
async def worker():
|
||
|
|
nonlocal entered, max_concurrent
|
||
|
|
async with rl:
|
||
|
|
entered += 1
|
||
|
|
max_concurrent = max(max_concurrent, entered)
|
||
|
|
await asyncio.sleep(0.05)
|
||
|
|
entered -= 1
|
||
|
|
|
||
|
|
await asyncio.gather(worker(), worker())
|
||
|
|
assert max_concurrent == 1 # only one at a time
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_registry_isolation():
|
||
|
|
reg = RateLimiterRegistry()
|
||
|
|
a = reg.get("source_a", rps=1.0)
|
||
|
|
b = reg.get("source_b", rps=1.0)
|
||
|
|
assert a is not b
|
||
|
|
# same key returns same instance
|
||
|
|
assert reg.get("source_a") is a
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_registry_register():
|
||
|
|
reg = RateLimiterRegistry()
|
||
|
|
reg.register("fast", rps=100.0, concurrent=5)
|
||
|
|
rl = reg.get("fast")
|
||
|
|
assert rl.rps == 100.0
|
||
|
|
assert rl.concurrent == 5
|