python-learning/rate_limiters/tests/test_token_bucket.py

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)