beam: convert test_token_bucket to unittest so runner collects all 9

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():
+class TestTokenBucket(unittest.TestCase):
-    b = TokenBucket(rate=10.0, capacity=5.0)
+    def test_basic_acquire_and_refill(self):
-    assert b.try_acquire(5.0) is True
+        b = TokenBucket(rate=10.0, capacity=5.0)
-    assert b.try_acquire(1.0) is False
+        self.assertTrue(b.try_acquire(5.0))
-    # refill: wait ~0.15s -> 1.5 tokens
+        self.assertFalse(b.try_acquire(1.0))
-    time.sleep(0.15)
+        time.sleep(0.15)
-    assert b.try_acquire(1.0) is True
+        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():
+    def test_acquire_blocks_and_returns(self):
-    b = TokenBucket(rate=100.0, capacity=3.0)
+        b = TokenBucket(rate=20.0, capacity=1.0)
-    time.sleep(0.1)
+        t0 = time.monotonic()
-    assert b.tokens() <= 3.0 + 1e-9
+        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():
+    def test_invalid_params(self):
-    b = TokenBucket(rate=20.0, capacity=1.0)
+        for kwargs in ({"rate": 0, "capacity": 1}, {"rate": 1, "capacity": 0}, {"rate": -1, "capacity": 1}):
-    # t0 must precede the drain: try_acquire() debits the last token and sets
+            with self.assertRaises(ValueError):
-    # the refill origin (self._last) inside the call, so a t0 taken after it
+                TokenBucket(**kwargs)
-    # silently subtracts whatever the scheduler spent in between.  Observed at
+        with self.assertRaises(ValueError):
-    # 7c36185 as `assert 0.03909331373870373 >= 0.04` in a loaded full-suite
+            TokenBucket(rate=1, capacity=1).try_acquire(0)
-    # 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_concurrent_no_overconsume(self):
+        b = TokenBucket(rate=1e-9, capacity=10.0)
+        successes = {"n": 0}
+        lock = threading.Lock()
 
-def test_n_exceeds_capacity_raises():
+        def worker():
-    b = TokenBucket(rate=1.0, capacity=1.0)
+            if b.try_acquire(1.0):
-    try:
+                with lock:
-        b.acquire(2.0, timeout=0.01)
+                    successes["n"] += 1
-        raised = False
+
-    except ValueError:
+        threads = [threading.Thread(target=worker) for _ in range(50)]
-        raised = True
+        for t in threads:
-    assert raised
+            t.start()
-    try:
+        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
+        self.assertFalse(b.try_acquire(1.0))
-    except ValueError:
+        b.reset()
-        raised2 = True
+        self.assertTrue(b.try_acquire(2.0))
-    assert raised2
+
+    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():
+if __name__ == "__main__":
-    for kwargs in ({"rate": 0, "capacity": 1}, {"rate": 1, "capacity": 0}, {"rate": -1, "capacity": 1}):
+    unittest.main()
-        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