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)