Add more unit tests

Author: Pulkit0110

google/cloud/storage/_experimental/asyncio/retry/bidi_stream_retry_manager.py

@@ -13,11 +13,9 @@
 # limitations under the License.
 
 import asyncio
-import time
 from typing import Any, AsyncIterator, Callable
 
 from google.api_core import exceptions
-from google.api_core.retry.retry_base import exponential_sleep_generator
 from google.cloud.storage._experimental.asyncio.retry.base_strategy import (
     _BaseResumptionStrategy,
 )
@@ -45,46 +43,25 @@ async def execute(self, initial_state: Any, retry_policy):
         """
         Executes the bidi operation with the configured retry policy.
 
-        This method implements a manual retry loop that provides the necessary
-        control points to manage state between attempts.
-
         Args:
             initial_state: An object containing all state for the operation.
             retry_policy: The `google.api_core.retry.AsyncRetry` object to
                 govern the retry behavior for this specific operation.
         """
         state = initial_state
 
-        deadline = time.monotonic() + retry_policy._deadline if retry_policy._deadline else 0
-
-        sleep_generator = exponential_sleep_generator(
-            retry_policy._initial, retry_policy._maximum, retry_policy._multiplier
-        )
-
-        while True:
+        async def attempt():
+            requests = self._strategy.generate_requests(state)
+            stream = self._stream_opener(requests, state)
             try:
-                requests = self._strategy.generate_requests(state)
-                stream = self._stream_opener(requests, state)
                 async for response in stream:
                     self._strategy.update_state_from_response(response, state)
                 return
             except Exception as e:
-                # AsyncRetry may expose either 'on_error' (public) or the private
-                # '_on_error' depending on google.api_core version. Call whichever
-                # exists so the retry policy can decide to raise (non-retriable /
-                # deadline exceeded) or allow a retry.
-                on_error_callable = getattr(retry_policy, "on_error", None)
-                if on_error_callable is None:
-                    on_error_callable = getattr(retry_policy, "_on_error", None)
-
-                if on_error_callable is None:
-                    # No hook available on the policy; re-raise the error.
-                    raise
+                if retry_policy._predicate(e):
+                    await self._strategy.recover_state_on_failure(e, state)
+                raise e
 
-                # Let the retry policy handle the error (may raise RetryError).
-                await on_error_callable(e)
+        wrapped_attempt = retry_policy(attempt)
 
-                # If the retry policy did not raise, allow the strategy to recover
-                # and then sleep per policy before next attempt.
-                await self._strategy.recover_state_on_failure(e, state)
-                await retry_policy.sleep()
+        await wrapped_attempt()

tests/unit/asyncio/retry/test_bidi_stream_retry_manager.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import asyncio
-import unittest
 from unittest import mock
 
 import pytest
@@ -30,9 +29,11 @@ def _is_retriable(exc):
 DEFAULT_TEST_RETRY = AsyncRetry(predicate=_is_retriable, deadline=1)
 
 
-class TestBidiStreamRetryManager(unittest.IsolatedAsyncioTestCase):
+class TestBidiStreamRetryManager:
+    @pytest.mark.asyncio
     async def test_execute_success_on_first_try(self):
         mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
+
         async def mock_stream_opener(*args, **kwargs):
             yield "response_1"
 
@@ -41,12 +42,33 @@ async def mock_stream_opener(*args, **kwargs):
         )
         await retry_manager.execute(initial_state={}, retry_policy=DEFAULT_TEST_RETRY)
         mock_strategy.generate_requests.assert_called_once()
-        mock_strategy.update_state_from_response.assert_called_once_with("response_1", {})
+        mock_strategy.update_state_from_response.assert_called_once_with(
+            "response_1", {}
+        )
         mock_strategy.recover_state_on_failure.assert_not_called()
 
-    async def test_execute_retries_and_succeeds(self):
+    @pytest.mark.asyncio
+    async def test_execute_success_on_empty_stream(self):
+        mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
+
+        async def mock_stream_opener(*args, **kwargs):
+            if False:
+                yield
+
+        retry_manager = manager._BidiStreamRetryManager(
+            strategy=mock_strategy, stream_opener=mock_stream_opener
+        )
+        await retry_manager.execute(initial_state={}, retry_policy=DEFAULT_TEST_RETRY)
+
+        mock_strategy.generate_requests.assert_called_once()
+        mock_strategy.update_state_from_response.assert_not_called()
+        mock_strategy.recover_state_on_failure.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_execute_retries_on_initial_failure_and_succeeds(self):
         mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
         attempt_count = 0
+
         async def mock_stream_opener(*args, **kwargs):
             nonlocal attempt_count
             attempt_count += 1
@@ -59,17 +81,63 @@ async def mock_stream_opener(*args, **kwargs):
             strategy=mock_strategy, stream_opener=mock_stream_opener
         )
         retry_policy = AsyncRetry(predicate=_is_retriable, initial=0.01)
-        retry_policy.sleep = mock.AsyncMock()
 
-        await retry_manager.execute(initial_state={}, retry_policy=retry_policy)
+        with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock):
+            await retry_manager.execute(initial_state={}, retry_policy=retry_policy)
+
+        assert attempt_count == 2
+        assert mock_strategy.generate_requests.call_count == 2
+        mock_strategy.recover_state_on_failure.assert_called_once()
+        mock_strategy.update_state_from_response.assert_called_once_with(
+            "response_2", {}
+        )
+
+    @pytest.mark.asyncio
+    async def test_execute_retries_and_succeeds_mid_stream(self):
+        """Test retry logic for a stream that fails after yielding some data."""
+        mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
+        attempt_count = 0
+        # Use a list to simulate stream content for each attempt
+        stream_content = [
+            ["response_1", exceptions.ServiceUnavailable("Service is down")],
+            ["response_2"],
+        ]
+
+        async def mock_stream_opener(*args, **kwargs):
+            nonlocal attempt_count
+            content = stream_content[attempt_count]
+            attempt_count += 1
+            for item in content:
+                if isinstance(item, Exception):
+                    raise item
+                else:
+                    yield item
+
+        retry_manager = manager._BidiStreamRetryManager(
+            strategy=mock_strategy, stream_opener=mock_stream_opener
+        )
+        retry_policy = AsyncRetry(predicate=_is_retriable, initial=0.01)
+
+        with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock) as mock_sleep:
+            await retry_manager.execute(initial_state={}, retry_policy=retry_policy)
+
+        assert attempt_count == 2
+        mock_sleep.assert_called_once()
 
-        self.assertEqual(attempt_count, 2)
-        self.assertEqual(mock_strategy.generate_requests.call_count, 2)
+        assert mock_strategy.generate_requests.call_count == 2
         mock_strategy.recover_state_on_failure.assert_called_once()
-        mock_strategy.update_state_from_response.assert_called_once_with("response_2", {})
+        assert mock_strategy.update_state_from_response.call_count == 2
+        mock_strategy.update_state_from_response.assert_has_calls(
+            [
+                mock.call("response_1", {}),
+                mock.call("response_2", {}),
+            ]
+        )
 
+    @pytest.mark.asyncio
     async def test_execute_fails_after_deadline_exceeded(self):
         mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
+
         async def mock_stream_opener(*args, **kwargs):
             if False:
                 yield
@@ -79,13 +147,15 @@ async def mock_stream_opener(*args, **kwargs):
         retry_manager = manager._BidiStreamRetryManager(
             strategy=mock_strategy, stream_opener=mock_stream_opener
         )
-        with pytest.raises(exceptions.RetryError, match="Deadline of 0.01s exceeded"):
+        with pytest.raises(exceptions.RetryError, match="Timeout of 0.0s exceeded"):
             await retry_manager.execute(initial_state={}, retry_policy=fast_retry)
 
-        self.assertGreater(mock_strategy.recover_state_on_failure.call_count, 0)
+        mock_strategy.recover_state_on_failure.assert_called_once()
 
+    @pytest.mark.asyncio
     async def test_execute_fails_immediately_on_non_retriable_error(self):
         mock_strategy = mock.AsyncMock(spec=base_strategy._BaseResumptionStrategy)
+
         async def mock_stream_opener(*args, **kwargs):
             if False:
                 yield