python-learning/rate_limiters/tests/test_token_bucket.py

26 lines
599 B
Python

import pytest
from time import sleep
from token_bucket import TokenBucketLimiter, RateLimitExceeded
def test_limit_respected() -> None:
limiter = TokenBucketLimiter(10, 10)
try:
for i in range(30):
with limiter:
sleep(0.1)
except RateLimitExceeded:
assert False, "Rate limit exceeded!"
def test_limit_exceeded() -> None:
limiter = TokenBucketLimiter(10, 10)
with pytest.raises(RateLimitExceeded) as e:
for i in range(20):
with limiter:
continue
assert "Rate exceeded" in str(e.value)