97 lines
2.7 KiB
Python
97 lines
2.7 KiB
Python
import time
|
|
import threading
|
|
import pytest
|
|
from ratelimit import TokenBucket, make_bucket, RateLimitError
|
|
|
|
|
|
def test_basic_burst():
|
|
b = make_bucket(rate=10, capacity=5, clock=lambda: 0.0)
|
|
# capacity 5, all available at t=0
|
|
for _ in range(5):
|
|
assert b.acquire(1, blocking=False) is True
|
|
assert b.acquire(1, blocking=False) is False
|
|
|
|
|
|
def test_refill():
|
|
t = [0.0]
|
|
b = make_bucket(rate=2, capacity=4, clock=lambda: t[0])
|
|
assert b.acquire(4, blocking=False) is True
|
|
assert b.acquire(1, blocking=False) is False
|
|
t[0] = 1.0 # 2 tokens refilled
|
|
assert b.acquire(2, blocking=False) is True
|
|
assert b.acquire(1, blocking=False) is False
|
|
|
|
|
|
def test_acquire_zero():
|
|
b = make_bucket(rate=1, capacity=1, clock=lambda: 0.0)
|
|
assert b.acquire(0) is True
|
|
assert b.tokens() == 1.0
|
|
|
|
|
|
def test_capacity_cap():
|
|
t = [0.0]
|
|
b = make_bucket(rate=10, capacity=3, clock=lambda: t[0])
|
|
t[0] = 100.0
|
|
assert b.tokens() == 3.0 # capped, not 1003
|
|
|
|
|
|
def test_timeout():
|
|
t = [0.0]
|
|
def clock():
|
|
return t[0]
|
|
b = make_bucket(rate=1, capacity=1, clock=clock)
|
|
assert b.acquire(1, blocking=False) is True
|
|
# need 1 more token, refill 1/s; timeout 0.5 -> fail
|
|
ok = b.acquire(1, blocking=True, timeout=0.5)
|
|
assert ok is False
|
|
|
|
|
|
def test_blocking_success_with_fake_sleep(monkeypatch):
|
|
t = [0.0]
|
|
b = make_bucket(rate=2, capacity=2, clock=lambda: t[0])
|
|
assert b.acquire(2, blocking=False) is True
|
|
# monkeypatch time.sleep to advance clock
|
|
def fake_sleep(s):
|
|
t[0] += s
|
|
monkeypatch.setattr(time, "sleep", fake_sleep)
|
|
assert b.acquire(2, blocking=True, timeout=5) is True
|
|
|
|
|
|
def test_n_exceeds_capacity_nonblocking():
|
|
b = make_bucket(rate=1, capacity=2, clock=lambda: 0.0)
|
|
assert b.acquire(5, blocking=False) is False
|
|
|
|
|
|
def test_n_exceeds_capacity_blocking_raises():
|
|
b = make_bucket(rate=1, capacity=2, clock=lambda: 0.0)
|
|
with pytest.raises(RateLimitError):
|
|
b.acquire(5, blocking=True)
|
|
|
|
|
|
def test_negative_n():
|
|
b = make_bucket(rate=1, capacity=2, clock=lambda: 0.0)
|
|
with pytest.raises(ValueError):
|
|
b.acquire(-1, blocking=False)
|
|
|
|
|
|
def test_thread_safety():
|
|
b = make_bucket(rate=1000, capacity=100, clock=time.monotonic)
|
|
results = []
|
|
def worker():
|
|
for _ in range(50):
|
|
results.append(b.acquire(1, blocking=True))
|
|
threads = [threading.Thread(target=worker) for _ in range(4)]
|
|
for th in threads:
|
|
th.start()
|
|
for th in threads:
|
|
th.join()
|
|
assert all(results)
|
|
assert len(results) == 200
|
|
|
|
|
|
def test_invalid_params():
|
|
with pytest.raises(ValueError):
|
|
TokenBucket(rate=0, capacity=1)
|
|
with pytest.raises(ValueError):
|
|
TokenBucket(rate=1, capacity=0)
|