From 8c8d5ce2a45a055a1e7ddaab8375654b0395cfb0 Mon Sep 17 00:00:00 2001 From: Pulkit Aggarwal Date: Mon, 6 Oct 2025 06:22:58 +0000 Subject: [PATCH 1/2] feat(experimental): refresh read handles --- .../_experimental/asyncio/async_read_object_stream.py | 7 +++++-- tests/unit/asyncio/test_async_read_object_stream.py | 5 +++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py b/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py index b14c55227..3e5961fe7 100644 --- a/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py +++ b/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py @@ -28,7 +28,6 @@ from google.cloud.storage._experimental.asyncio.async_abstract_object_stream import ( _AsyncAbstractObjectStream, ) -from google.cloud.storage._experimental.asyncio.bidi_async import AsyncBidiRpc class _AsyncReadObjectStream(_AsyncAbstractObjectStream): @@ -145,7 +144,11 @@ async def recv(self) -> _storage_v2.BidiReadObjectResponse: """ if not self._is_stream_open: raise ValueError("Stream is not open") - return await self.socket_like_rpc.recv() + response = await self.socket_like_rpc.recv() + # Update read_handle if present in response + if response and response.read_handle: + self.read_handle = response.read_handle + return response @property def is_stream_open(self) -> bool: diff --git a/tests/unit/asyncio/test_async_read_object_stream.py b/tests/unit/asyncio/test_async_read_object_stream.py index ea91b2251..e140f7b40 100644 --- a/tests/unit/asyncio/test_async_read_object_stream.py +++ b/tests/unit/asyncio/test_async_read_object_stream.py @@ -16,7 +16,9 @@ from unittest import mock from unittest.mock import AsyncMock from google.cloud import _storage_v2 +from google.cloud.storage._experimental.asyncio import bidi_async +from google.cloud.storage._experimental.asyncio import async_read_object_stream from google.cloud.storage._experimental.asyncio.async_read_object_stream import ( _AsyncReadObjectStream, ) @@ -273,3 +275,6 @@ async def test_recv_without_open_should_raise_error( # assert assert str(exc.value) == "Stream is not open" + + + From 4f964263c72033635becb19630c96e2ff237c4f9 Mon Sep 17 00:00:00 2001 From: Pulkit Aggarwal Date: Mon, 6 Oct 2025 06:43:22 +0000 Subject: [PATCH 2/2] fix unit tests --- .../asyncio/async_read_object_stream.py | 1 + .../asyncio/test_async_read_object_stream.py | 44 ++++++++++++++++++- 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py b/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py index 3e5961fe7..96e9f154a 100644 --- a/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py +++ b/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py @@ -28,6 +28,7 @@ from google.cloud.storage._experimental.asyncio.async_abstract_object_stream import ( _AsyncAbstractObjectStream, ) +from google.cloud.storage._experimental.asyncio.bidi_async import AsyncBidiRpc class _AsyncReadObjectStream(_AsyncAbstractObjectStream): diff --git a/tests/unit/asyncio/test_async_read_object_stream.py b/tests/unit/asyncio/test_async_read_object_stream.py index e140f7b40..ae5a48d76 100644 --- a/tests/unit/asyncio/test_async_read_object_stream.py +++ b/tests/unit/asyncio/test_async_read_object_stream.py @@ -16,7 +16,6 @@ from unittest import mock from unittest.mock import AsyncMock from google.cloud import _storage_v2 -from google.cloud.storage._experimental.asyncio import bidi_async from google.cloud.storage._experimental.asyncio import async_read_object_stream from google.cloud.storage._experimental.asyncio.async_read_object_stream import ( @@ -276,5 +275,48 @@ async def test_recv_without_open_should_raise_error( # assert assert str(exc.value) == "Stream is not open" +@mock.patch("google.cloud.storage._experimental.asyncio.async_read_object_stream.AsyncBidiRpc") +@mock.patch("google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client") +@pytest.mark.asyncio +async def test_recv_updates_read_handle_on_refresh(mock_client, mock_cls_async_bidi_rpc): + """ + Verify that the `recv` method correctly updates the stream's handle + when a new one is provided in a server response. + """ + # Arrange + socket_like_rpc = AsyncMock() + mock_cls_async_bidi_rpc.return_value = socket_like_rpc + socket_like_rpc.open = AsyncMock() + + initial_handle = _storage_v2.BidiReadHandle(handle=b"initial-handle-token") + response_with_initial_handle = _storage_v2.BidiReadObjectResponse(read_handle=initial_handle) + response_without_handle = _storage_v2.BidiReadObjectResponse(read_handle=None) + + refreshed_handle = _storage_v2.BidiReadHandle(handle=b"new-refreshed-handle-token") + response_with_refreshed_handle = _storage_v2.BidiReadObjectResponse(read_handle=refreshed_handle) + + socket_like_rpc.recv.side_effect = [ + response_with_initial_handle, + response_without_handle, + response_with_refreshed_handle, + ] + + starting_handle = _storage_v2.BidiReadHandle(handle=b"starting-handle-token") + stream = async_read_object_stream._AsyncReadObjectStream( + client=mock_client, + bucket_name=_TEST_BUCKET_NAME, + object_name=_TEST_OBJECT_NAME, + read_handle=starting_handle, + ) + + # Act & Assert + assert stream.read_handle == starting_handle + + await stream.open() + assert stream.read_handle == initial_handle + await stream.recv() + assert stream.read_handle == initial_handle + await stream.recv() + assert stream.read_handle == refreshed_handle