28 lines
597 B
Python
28 lines
597 B
Python
from time import sleep
|
|
|
|
import pytest
|
|
|
|
from token_bucket import RateLimitExceeded, TokenBucketLimiter
|
|
|
|
|
|
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)
|