added set_next to TrackedBackoffGenerator

Author: daniel-sanche

google/cloud/bigtable/data/_helpers.py

@@ -264,18 +264,45 @@ def __init__(self, initial=0.01, maximum=60, multiplier=2):
         self.subgenerator = exponential_sleep_generator(
             initial=initial, maximum=maximum, multiplier=multiplier
         )
+        self._next_override: float | None = None
 
     def __iter__(self):
         return self
 
+    def set_next(self, next_value: float):
+        """
+        Set the next backoff value, instead of generating one from subgenerator.
+        After the value is yielded, it will go back to using self.subgenerator.
+
+        If set_next is called twice before the next() is called, only the latest
+        value will be used and others discarded
+
+        Args:
+            next_value: the upcomming value to yield when next() is called
+        Raises:
+            ValueError: if next_value is negative
+        """
+        if next_value < 0:
+            raise ValueError("backoff value cannot be less than 0")
+        self._next_override = next_value
+
     def __next__(self) -> float:
-        next_backoff = next(self.subgenerator)
+        if self._next_override is not None:
+            next_backoff = self._next_override
+            self._next_override = None
+        else:
+            next_backoff = next(self.subgenerator)
         self.history.append(next_backoff)
         return next_backoff
 
     def get_attempt_backoff(self, attempt_idx) -> float:
         """
         returns the backoff time for a specific attempt index, starting at 0.
+
+        Args:
+            attempt_idx: the index of the attempt to return backoff for
+        Raises:
+            IndexError: if attempt_idx is negative, or not in history
         """
         if attempt_idx < 0:
             raise IndexError("received negative attempt number")

tests/unit/data/test__helpers.py

@@ -317,3 +317,47 @@ def test_get_attempt_backoff_out_of_bounds(self):
             generator.get_attempt_backoff(2)
         with pytest.raises(IndexError):
             generator.get_attempt_backoff(-3)
+
+    def test_set_next_full_set(self):
+        """
+        try always using set_next to populate generator
+        """
+        generator = _helpers.TrackedBackoffGenerator()
+        for idx, val in enumerate(range(100, 0, -1)):
+            generator.set_next(val)
+            got = next(generator)
+            assert got == val
+            assert generator.get_attempt_backoff(idx) == val
+
+    def test_set_next_negative_value(self):
+        generator = _helpers.TrackedBackoffGenerator()
+        with pytest.raises(ValueError):
+            generator.set_next(-1)
+
+    @mock.patch("random.uniform", side_effect=lambda a, b: b)
+    def test_interleaved_set_next(self, mock_uniform):
+        import itertools
+
+        generator = _helpers.TrackedBackoffGenerator(
+            initial=1, multiplier=2, maximum=128
+        )
+        # values we expect generator to create
+        expected_values = [2**i for i in range(8)]
+        # values we will insert
+        inserted_values = [9, 61, 0, 4, 33, 12, 18, 2]
+        for idx in range(8):
+            assert next(generator) == expected_values[idx]
+            generator.set_next(inserted_values[idx])
+            assert next(generator) == inserted_values[idx]
+        # check to make sure history is as we expect
+        generator.history = itertools.chain.from_iterable(
+            zip(expected_values, inserted_values)
+        )
+
+    @mock.patch("random.uniform", side_effect=lambda a, b: b)
+    def test_set_next_replacement(self, mock_uniform):
+        generator = _helpers.TrackedBackoffGenerator(initial=1)
+        generator.set_next(99)
+        generator.set_next(88)
+        assert next(generator) == 88
+        assert next(generator) == 1