beam: add injectable clock to TokenBucket for deterministic testing

wizard-work/beam/token_bucket.py

@@ -14,18 +14,16 @@ from typing import Optional
 
 
 class TokenBucket:
-    def __init__(self, rate: float, capacity: float):
+    def __init__(self, rate: float, capacity: float, clock=None):
         if rate <= 0:
             raise ValueError("rate must be > 0")
         if capacity <= 0:
             raise ValueError("capacity must be > 0")
-        if rate > capacity:
-            # allow; but typically rate <= capacity. No restriction.
-            pass
         self._rate = float(rate)
         self._capacity = float(capacity)
+        self._clock = clock if clock is not None else time.monotonic
         self._tokens = float(capacity)
-        self._last = time.monotonic()
+        self._last = self._clock()
         self._lock = threading.Lock()
         self._cond = threading.Condition(self._lock)
 
@@ -46,7 +44,7 @@ class TokenBucket:
 
     def tokens(self) -> float:
         with self._lock:
-            self._refill(time.monotonic())
+            self._refill(self._clock())
             return self._tokens
 
     def try_acquire(self, n: float = 1.0) -> bool:
@@ -55,7 +53,7 @@ class TokenBucket:
         if n > self._capacity:
             raise ValueError("n exceeds capacity")
         with self._lock:
-            now = time.monotonic()
+            now = self._clock()
             self._refill(now)
             if self._tokens >= n - 1e-12:
                 self._tokens -= n
@@ -67,10 +65,10 @@ class TokenBucket:
             raise ValueError("n must be > 0")
         if n > self._capacity:
             raise ValueError("n exceeds capacity")
-        deadline = None if timeout is None else time.monotonic() + timeout
+        deadline = None if timeout is None else self._clock() + timeout
         with self._lock:
             while True:
-                now = time.monotonic()
+                now = self._clock()
                 self._refill(now)
                 if self._tokens >= n - 1e-12:
                     self._tokens -= n
@@ -99,5 +97,5 @@ class TokenBucket:
     def reset(self, tokens: Optional[float] = None) -> None:
         with self._lock:
             self._tokens = self._capacity if tokens is None else float(tokens)
-            self._last = time.monotonic()
+            self._last = self._clock()
             self._cond.notify_all()