diff --git a/wizard-work/beam/test_token_bucket.py b/wizard-work/beam/test_token_bucket.py index 3e54748..b46ab8d 100644 --- a/wizard-work/beam/test_token_bucket.py +++ b/wizard-work/beam/test_token_bucket.py @@ -1,116 +1,87 @@ import threading import time +import unittest 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 +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_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(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_acquire_blocks_and_returns(): - b = TokenBucket(rate=20.0, capacity=1.0) - # t0 must precede the drain: try_acquire() debits the last token and sets - # the refill origin (self._last) inside the call, so a t0 taken after it - # silently subtracts whatever the scheduler spent in between. Observed at - # 7c36185 as `assert 0.03909331373870373 >= 0.04` in a loaded full-suite - # run: ~10.9 ms of scheduler gap, i.e. 0.05 s of withheld tokens minus a - # prefix the cell never measured. dt from before the drain is >= 0.05 by - # construction (acquire cannot return until the refill origin is 0.05 s - # old), so the 0.04 floor has margin instead of sitting 1 ms from the edge. - t0 = time.monotonic() - b.try_acquire(1.0) # drain; sets the refill origin - ok = b.acquire(1.0, timeout=0.2) - dt = time.monotonic() - t0 - assert ok is True - assert dt >= 0.04 # 1 token at 20/s = 0.05s withheld, measured from the origin + 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 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: + 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) - raised2 = False - except ValueError: - raised2 = True - assert raised2 + 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) -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 +if __name__ == "__main__": + unittest.main()