88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
import threading
|
|
import time
|
|
import unittest
|
|
|
|
from token_bucket import TokenBucket
|
|
|
|
|
|
class TestTokenBucket(unittest.TestCase):
|
|
def test_basic_acquire_and_refill(self):
|
|
b = TokenBucket(rate=10.0, capacity=5.0)
|
|
self.assertTrue(b.try_acquire(5.0))
|
|
self.assertFalse(b.try_acquire(1.0))
|
|
time.sleep(0.15)
|
|
self.assertTrue(b.try_acquire(1.0))
|
|
|
|
def test_capacity_cap(self):
|
|
b = TokenBucket(rate=100.0, capacity=3.0)
|
|
time.sleep(0.1)
|
|
self.assertLessEqual(b.tokens(), 3.0 + 1e-9)
|
|
|
|
def test_acquire_blocks_and_returns(self):
|
|
b = TokenBucket(rate=20.0, capacity=1.0)
|
|
t0 = time.monotonic()
|
|
b.try_acquire(1.0)
|
|
ok = b.acquire(1.0, timeout=0.2)
|
|
dt = time.monotonic() - t0
|
|
self.assertTrue(ok)
|
|
self.assertGreaterEqual(dt, 0.04)
|
|
|
|
def test_n_exceeds_capacity_raises(self):
|
|
b = TokenBucket(rate=1.0, capacity=1.0)
|
|
with self.assertRaises(ValueError):
|
|
b.acquire(2.0, timeout=0.01)
|
|
with self.assertRaises(ValueError):
|
|
b.try_acquire(2.0)
|
|
|
|
def test_invalid_params(self):
|
|
for kwargs in ({"rate": 0, "capacity": 1}, {"rate": 1, "capacity": 0}, {"rate": -1, "capacity": 1}):
|
|
with self.assertRaises(ValueError):
|
|
TokenBucket(**kwargs)
|
|
with self.assertRaises(ValueError):
|
|
TokenBucket(rate=1, capacity=1).try_acquire(0)
|
|
|
|
def test_concurrent_no_overconsume(self):
|
|
b = TokenBucket(rate=1e-9, capacity=10.0)
|
|
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()
|
|
self.assertEqual(successes["n"], 10)
|
|
|
|
def test_acquire_timeout_false(self):
|
|
b = TokenBucket(rate=1.0, capacity=1.0)
|
|
b.try_acquire(1.0)
|
|
t0 = time.monotonic()
|
|
ok = b.acquire(1.0, timeout=0.05)
|
|
dt = time.monotonic() - t0
|
|
self.assertFalse(ok)
|
|
self.assertLess(dt, 0.2)
|
|
|
|
def test_reset(self):
|
|
b = TokenBucket(rate=1.0, capacity=2.0)
|
|
b.try_acquire(2.0)
|
|
self.assertFalse(b.try_acquire(1.0))
|
|
b.reset()
|
|
self.assertTrue(b.try_acquire(2.0))
|
|
|
|
def test_partial_never_consumes(self):
|
|
b = TokenBucket(rate=1.0, capacity=2.0)
|
|
b.try_acquire(2.0)
|
|
self.assertFalse(b.try_acquire(1.0))
|
|
time.sleep(0.5)
|
|
self.assertFalse(b.try_acquire(1.0))
|
|
self.assertGreaterEqual(b.tokens(), 0.4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|