added unit tests for dynamic backoff

Author: daniel-sanche

tests/asyncio/retry/test_retry_streaming_async.py

@@ -39,6 +39,35 @@ async def test_retry_streaming_target_bad_sleep_generator():
         await retry_target_stream(None, lambda x: True, [], None).__anext__()
 
 
+@mock.patch("asyncio.sleep", autospec=True)
+@pytest.mark.asyncio
+async def test_retry_streaming_target_dynamic_backoff(sleep):
+    """
+    sleep_generator should be iterated after on_error, to support dynamic backoff
+    """
+    from functools import partial
+    from google.api_core.retry.retry_streaming_async import retry_target_stream
+
+    sleep.side_effect = RuntimeError("stop after sleep")
+    # start with empty sleep generator; values are added after exception in push_sleep_value
+    sleep_values = []
+    error_target = partial(TestAsyncStreamingRetry._generator_mock, error_on=0)
+    inserted_sleep = 99
+
+    def push_sleep_value(err):
+        sleep_values.append(inserted_sleep)
+
+    with pytest.raises(RuntimeError):
+        await retry_target_stream(
+            error_target,
+            predicate=lambda x: True,
+            sleep_generator=sleep_values,
+            on_error=push_sleep_value,
+        ).__anext__()
+    assert sleep.call_count == 1
+    sleep.assert_called_once_with(inserted_sleep)
+
+
 class TestAsyncStreamingRetry(Test_BaseRetry):
     def _make_one(self, *args, **kwargs):
         return retry_streaming_async.AsyncStreamingRetry(*args, **kwargs)
@@ -66,8 +95,8 @@ def if_exception_type(exc):
             str(retry_),
         )
 
+    @staticmethod
     async def _generator_mock(
-        self,
         num=5,
         error_on=None,
         exceptions_seen=None,
@@ -87,7 +116,7 @@ async def _generator_mock(
             for i in range(num):
                 if sleep_time:
                     await asyncio.sleep(sleep_time)
-                if error_on and i == error_on:
+                if error_on is not None and i == error_on:
                     raise ValueError("generator mock error")
                 yield i
         except (Exception, BaseException, GeneratorExit) as e:

tests/asyncio/retry/test_retry_unary_async.py

@@ -139,6 +139,33 @@ async def test_retry_target_bad_sleep_generator():
         await retry_async.retry_target(mock.sentinel.target, lambda x: True, [], None)
 
 
+@mock.patch("asyncio.sleep", autospec=True)
+@pytest.mark.asyncio
+async def test_retry_target_dynamic_backoff(sleep):
+    """
+    sleep_generator should be iterated after on_error, to support dynamic backoff
+    """
+    sleep.side_effect = RuntimeError("stop after sleep")
+    # start with empty sleep generator; values are added after exception in push_sleep_value
+    sleep_values = []
+    exception = ValueError("trigger retry")
+    error_target = mock.Mock(side_effect=exception)
+    inserted_sleep = 99
+
+    def push_sleep_value(err):
+        sleep_values.append(inserted_sleep)
+
+    with pytest.raises(RuntimeError):
+        await retry_async.retry_target(
+            error_target,
+            predicate=lambda x: True,
+            sleep_generator=sleep_values,
+            on_error=push_sleep_value,
+        )
+    assert sleep.call_count == 1
+    sleep.assert_called_once_with(inserted_sleep)
+
+
 class TestAsyncRetry(Test_BaseRetry):
     def _make_one(self, *args, **kwargs):
         return retry_async.AsyncRetry(*args, **kwargs)

tests/unit/retry/test_retry_streaming.py

@@ -36,6 +36,35 @@ def test_retry_streaming_target_bad_sleep_generator():
         next(retry_streaming.retry_target_stream(None, lambda x: True, [], None))
 
 
+@mock.patch("time.sleep", autospec=True)
+def test_retry_streaming_target_dynamic_backoff(sleep):
+    """
+    sleep_generator should be iterated after on_error, to support dynamic backoff
+    """
+    from functools import partial
+
+    sleep.side_effect = RuntimeError("stop after sleep")
+    # start with empty sleep generator; values are added after exception in push_sleep_value
+    sleep_values = []
+    error_target = partial(TestStreamingRetry._generator_mock, error_on=0)
+    inserted_sleep = 99
+
+    def push_sleep_value(err):
+        sleep_values.append(inserted_sleep)
+
+    with pytest.raises(RuntimeError):
+        next(
+            retry_streaming.retry_target_stream(
+                error_target,
+                predicate=lambda x: True,
+                sleep_generator=sleep_values,
+                on_error=push_sleep_value,
+            )
+        )
+    assert sleep.call_count == 1
+    sleep.assert_called_once_with(inserted_sleep)
+
+
 class TestStreamingRetry(Test_BaseRetry):
     def _make_one(self, *args, **kwargs):
         return retry_streaming.StreamingRetry(*args, **kwargs)
@@ -63,8 +92,8 @@ def if_exception_type(exc):
             str(retry_),
         )
 
+    @staticmethod
     def _generator_mock(
-        self,
         num=5,
         error_on=None,
         return_val=None,
@@ -82,7 +111,7 @@ def _generator_mock(
         """
         try:
             for i in range(num):
-                if error_on and i == error_on:
+                if error_on is not None and i == error_on:
                     raise ValueError("generator mock error")
                 yield i
             return return_val

tests/unit/retry/test_retry_unary.py

@@ -149,6 +149,32 @@ def test_retry_target_bad_sleep_generator():
         retry.retry_target(mock.sentinel.target, lambda x: True, [], None)
 
 
+@mock.patch("time.sleep", autospec=True)
+def test_retry_target_dynamic_backoff(sleep):
+    """
+    sleep_generator should be iterated after on_error, to support dynamic backoff
+    """
+    sleep.side_effect = RuntimeError("stop after sleep")
+    # start with empty sleep generator; values are added after exception in push_sleep_value
+    sleep_values = []
+    exception = ValueError("trigger retry")
+    error_target = mock.Mock(side_effect=exception)
+    inserted_sleep = 99
+
+    def push_sleep_value(err):
+        sleep_values.append(inserted_sleep)
+
+    with pytest.raises(RuntimeError):
+        retry.retry_target(
+            error_target,
+            predicate=lambda x: True,
+            sleep_generator=sleep_values,
+            on_error=push_sleep_value,
+        )
+    assert sleep.call_count == 1
+    sleep.assert_called_once_with(inserted_sleep)
+
+
 class TestRetry(Test_BaseRetry):
     def _make_one(self, *args, **kwargs):
         return retry.Retry(*args, **kwargs)