109 lines
2.7 KiB
Python
109 lines
2.7 KiB
Python
import threading
|
|
import time
|
|
|
|
from token_bucket import TokenBucket
|
|
|
|
|
|
def test_basic_acquire_and_refill():
|
|
b = TokenBucket(rate=10.0, capacity=5.0)
|
|
assert b.try_acquire(5.0) is True
|
|
assert b.try_acquire(1.0) is False
|
|
# refill: wait ~0.15s -> 1.5 tokens
|
|
time.sleep(0.15)
|
|
assert b.try_acquire(1.0) is True
|
|
|
|
|
|
def test_capacity_cap():
|
|
b = TokenBucket(rate=100.0, capacity=3.0)
|
|
time.sleep(0.1)
|
|
assert b.tokens() <= 3.0 + 1e-9
|
|
|
|
|
|
def test_acquire_blocks_and_returns():
|
|
b = TokenBucket(rate=20.0, capacity=1.0)
|
|
b.try_acquire(1.0) # drain
|
|
t0 = time.monotonic()
|
|
ok = b.acquire(1.0, timeout=0.2)
|
|
dt = time.monotonic() - t0
|
|
assert ok is True
|
|
assert dt >= 0.04 # had to wait ~0.05s for 1 token at 20/s
|
|
|
|
|
|
def test_n_exceeds_capacity_raises():
|
|
b = TokenBucket(rate=1.0, capacity=1.0)
|
|
try:
|
|
b.acquire(2.0, timeout=0.01)
|
|
raised = False
|
|
except ValueError:
|
|
raised = True
|
|
assert raised
|
|
try:
|
|
b.try_acquire(2.0)
|
|
raised2 = False
|
|
except ValueError:
|
|
raised2 = True
|
|
assert raised2
|
|
|
|
|
|
def test_invalid_params():
|
|
for kwargs in ({"rate": 0, "capacity": 1}, {"rate": 1, "capacity": 0}, {"rate": -1, "capacity": 1}):
|
|
try:
|
|
TokenBucket(**kwargs)
|
|
assert False
|
|
except ValueError:
|
|
pass
|
|
try:
|
|
TokenBucket(rate=1, capacity=1).try_acquire(0)
|
|
assert False
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
def test_concurrent_no_overconsume():
|
|
b = TokenBucket(rate=1e-9, capacity=10.0) # effectively no refill
|
|
successes = {"n": 0}
|
|
lock = threading.Lock()
|
|
|
|
def worker():
|
|
if b.try_acquire(1.0):
|
|
with lock:
|
|
successes["n"] += 1
|
|
|
|
threads = [threading.Thread(target=worker) for _ in range(50)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
assert successes["n"] == 10
|
|
|
|
|
|
def test_acquire_timeout_false():
|
|
b = TokenBucket(rate=1.0, capacity=1.0)
|
|
b.try_acquire(1.0) # drain
|
|
# request 1 token but timeout too short for refill (need 1s)
|
|
t0 = time.monotonic()
|
|
ok = b.acquire(1.0, timeout=0.05)
|
|
dt = time.monotonic() - t0
|
|
assert ok is False
|
|
assert dt < 0.2
|
|
|
|
|
|
def test_reset():
|
|
b = TokenBucket(rate=1.0, capacity=2.0)
|
|
b.try_acquire(2.0)
|
|
assert b.try_acquire(1.0) is False
|
|
b.reset()
|
|
assert b.try_acquire(2.0) is True
|
|
|
|
|
|
def test_partial_never_consumes():
|
|
b = TokenBucket(rate=1.0, capacity=2.0)
|
|
b.try_acquire(2.0) # drain
|
|
# 1 token needed, only 0 available -> no consume
|
|
assert b.try_acquire(1.0) is False
|
|
# refill 0.5s -> 0.5 tokens, still < 1
|
|
time.sleep(0.5)
|
|
assert b.try_acquire(1.0) is False
|
|
# tokens should be ~0.5, not negative
|
|
assert b.tokens() >= 0.4
|