beam: convert test_token_bucket to unittest so runner collects all 9

wizard-work/beam/test_token_bucket.py

@@ -1,74 +1,48 @@
 import threading
 import time
+import unittest
 
 from token_bucket import TokenBucket
 
 
-def test_basic_acquire_and_refill():
+class TestTokenBucket(unittest.TestCase):
+    def test_basic_acquire_and_refill(self):
         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
+        self.assertTrue(b.try_acquire(5.0))
+        self.assertFalse(b.try_acquire(1.0))
         time.sleep(0.15)
-    assert b.try_acquire(1.0) is True
+        self.assertTrue(b.try_acquire(1.0))
 
-
-def test_capacity_cap():
+    def test_capacity_cap(self):
         b = TokenBucket(rate=100.0, capacity=3.0)
         time.sleep(0.1)
-    assert b.tokens() <= 3.0 + 1e-9
+        self.assertLessEqual(b.tokens(), 3.0 + 1e-9)
 
-
-def test_acquire_blocks_and_returns():
+    def test_acquire_blocks_and_returns(self):
         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
+        b.try_acquire(1.0)
         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
+        self.assertTrue(ok)
+        self.assertGreaterEqual(dt, 0.04)
 
-
-def test_n_exceeds_capacity_raises():
+    def test_n_exceeds_capacity_raises(self):
         b = TokenBucket(rate=1.0, capacity=1.0)
-    try:
+        with self.assertRaises(ValueError):
             b.acquire(2.0, timeout=0.01)
-        raised = False
-    except ValueError:
-        raised = True
-    assert raised
-    try:
+        with self.assertRaises(ValueError):
             b.try_acquire(2.0)
-        raised2 = False
-    except ValueError:
-        raised2 = True
-    assert raised2
 
-
-def test_invalid_params():
+    def test_invalid_params(self):
         for kwargs in ({"rate": 0, "capacity": 1}, {"rate": 1, "capacity": 0}, {"rate": -1, "capacity": 1}):
-        try:
+            with self.assertRaises(ValueError):
                 TokenBucket(**kwargs)
-            assert False
-        except ValueError:
-            pass
-    try:
+        with self.assertRaises(ValueError):
             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
+    def test_concurrent_no_overconsume(self):
+        b = TokenBucket(rate=1e-9, capacity=10.0)
         successes = {"n": 0}
         lock = threading.Lock()
 
@@ -82,35 +56,32 @@ def test_concurrent_no_overconsume():
             t.start()
         for t in threads:
             t.join()
-    assert successes["n"] == 10
+        self.assertEqual(successes["n"], 10)
 
-
-def test_acquire_timeout_false():
+    def test_acquire_timeout_false(self):
         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)
+        b.try_acquire(1.0)
         t0 = time.monotonic()
         ok = b.acquire(1.0, timeout=0.05)
         dt = time.monotonic() - t0
-    assert ok is False
-    assert dt < 0.2
+        self.assertFalse(ok)
+        self.assertLess(dt, 0.2)
 
-
-def test_reset():
+    def test_reset(self):
         b = TokenBucket(rate=1.0, capacity=2.0)
         b.try_acquire(2.0)
-    assert b.try_acquire(1.0) is False
+        self.assertFalse(b.try_acquire(1.0))
         b.reset()
-    assert b.try_acquire(2.0) is True
+        self.assertTrue(b.try_acquire(2.0))
 
-
-def test_partial_never_consumes():
+    def test_partial_never_consumes(self):
         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
+        b.try_acquire(2.0)
+        self.assertFalse(b.try_acquire(1.0))
         time.sleep(0.5)
-    assert b.try_acquire(1.0) is False
-    # tokens should be ~0.5, not negative
-    assert b.tokens() >= 0.4
+        self.assertFalse(b.try_acquire(1.0))
+        self.assertGreaterEqual(b.tokens(), 0.4)
+
+
+if __name__ == "__main__":
+    unittest.main()