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 316a6750a..b14c55227 100644 --- a/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py +++ b/google/cloud/storage/_experimental/asyncio/async_read_object_stream.py @@ -83,15 +83,14 @@ def __init__( self.rpc = self.client._client._transport._wrapped_methods[ self.client._client._transport.bidi_read_object ] - first_bidi_read_req = _storage_v2.BidiReadObjectRequest( + self.first_bidi_read_req = _storage_v2.BidiReadObjectRequest( read_object_spec=_storage_v2.BidiReadObjectSpec( bucket=self._full_bucket_name, object=object_name ), ) self.metadata = (("x-goog-request-params", f"bucket={self._full_bucket_name}"),) - self.socket_like_rpc = AsyncBidiRpc( - self.rpc, initial_request=first_bidi_read_req, metadata=self.metadata - ) + self.socket_like_rpc: Optional[AsyncBidiRpc] = None + self._is_stream_open: bool = False async def open(self) -> None: """Opens the bidi-gRPC connection to read from the object. @@ -99,6 +98,11 @@ async def open(self) -> None: This method sends an initial request to start the stream and receives the first response containing metadata and a read handle. """ + if self._is_stream_open: + raise ValueError("Stream is already open") + self.socket_like_rpc = AsyncBidiRpc( + self.rpc, initial_request=self.first_bidi_read_req, metadata=self.metadata + ) await self.socket_like_rpc.open() # this is actually 1 send response = await self.socket_like_rpc.recv() if self.generation_number is None: @@ -106,9 +110,14 @@ async def open(self) -> None: self.read_handle = response.read_handle + self._is_stream_open = True + async def close(self) -> None: """Closes the bidi-gRPC connection.""" + if not self._is_stream_open: + raise ValueError("Stream is not open") await self.socket_like_rpc.close() + self._is_stream_open = False async def send( self, bidi_read_object_request: _storage_v2.BidiReadObjectRequest @@ -120,6 +129,8 @@ async def send( The request message to send. This is typically used to specify the read offset and limit. """ + if not self._is_stream_open: + raise ValueError("Stream is not open") await self.socket_like_rpc.send(bidi_read_object_request) async def recv(self) -> _storage_v2.BidiReadObjectResponse: @@ -132,4 +143,10 @@ async def recv(self) -> _storage_v2.BidiReadObjectResponse: :class:`~google.cloud._storage_v2.types.BidiReadObjectResponse`: The response message from the server. """ + if not self._is_stream_open: + raise ValueError("Stream is not open") return await self.socket_like_rpc.recv() + + @property + def is_stream_open(self) -> bool: + return self._is_stream_open diff --git a/tests/unit/asyncio/test_async_read_object_stream.py b/tests/unit/asyncio/test_async_read_object_stream.py index d7170c4a4..ea91b2251 100644 --- a/tests/unit/asyncio/test_async_read_object_stream.py +++ b/tests/unit/asyncio/test_async_read_object_stream.py @@ -27,6 +27,30 @@ _TEST_READ_HANDLE = b"test-read-handle" +async def instantiate_read_obj_stream(mock_client, mock_cls_async_bidi_rpc, open=True): + """Helper to create an instance of _AsyncReadObjectStream and open it by default.""" + socket_like_rpc = AsyncMock() + mock_cls_async_bidi_rpc.return_value = socket_like_rpc + socket_like_rpc.open = AsyncMock() + + recv_response = mock.MagicMock(spec=_storage_v2.BidiReadObjectResponse) + recv_response.metadata = mock.MagicMock(spec=_storage_v2.Object) + recv_response.metadata.generation = _TEST_GENERATION_NUMBER + recv_response.read_handle = _TEST_READ_HANDLE + socket_like_rpc.recv = AsyncMock(return_value=recv_response) + + read_obj_stream = _AsyncReadObjectStream( + client=mock_client, + bucket_name=_TEST_BUCKET_NAME, + object_name=_TEST_OBJECT_NAME, + ) + + if open: + await read_obj_stream.open() + + return read_obj_stream + + @mock.patch( "google.cloud.storage._experimental.asyncio.async_read_object_stream.AsyncBidiRpc" ) @@ -34,9 +58,11 @@ "google.cloud.storage._experimental.asyncio.async_grpc_client.AsyncGrpcClient.grpc_client" ) def test_init_with_bucket_object_generation(mock_client, mock_async_bidi_rpc): + # Arrange + rpc_sentinel = mock.sentinel.A mock_client._client._transport.bidi_read_object = "bidi_read_object_rpc" mock_client._client._transport._wrapped_methods = { - "bidi_read_object_rpc": mock.sentinel.A + "bidi_read_object_rpc": rpc_sentinel, } full_bucket_name = f"projects/_/buckets/{_TEST_BUCKET_NAME}" first_bidi_read_req = _storage_v2.BidiReadObjectRequest( @@ -45,39 +71,36 @@ def test_init_with_bucket_object_generation(mock_client, mock_async_bidi_rpc): ), ) + # Act read_obj_stream = _AsyncReadObjectStream( client=mock_client, bucket_name=_TEST_BUCKET_NAME, object_name=_TEST_OBJECT_NAME, generation_number=_TEST_GENERATION_NUMBER, + read_handle=_TEST_READ_HANDLE, ) - mock_async_bidi_rpc.assert_called_once_with( - mock.sentinel.A, - initial_request=first_bidi_read_req, - metadata=(("x-goog-request-params", f"bucket={full_bucket_name}"),), - ) - assert read_obj_stream.socket_like_rpc is mock_async_bidi_rpc.return_value + # Assert + assert read_obj_stream.bucket_name == _TEST_BUCKET_NAME + assert read_obj_stream.object_name == _TEST_OBJECT_NAME + assert read_obj_stream.generation_number == _TEST_GENERATION_NUMBER + assert read_obj_stream.read_handle == _TEST_READ_HANDLE + assert read_obj_stream.first_bidi_read_req == first_bidi_read_req + assert read_obj_stream.rpc == rpc_sentinel +@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_open(mock_client): +async def test_open(mock_client, mock_cls_async_bidi_rpc): # arrange - read_obj_stream = _AsyncReadObjectStream( - client=mock_client, - bucket_name=_TEST_BUCKET_NAME, - object_name=_TEST_OBJECT_NAME, + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False ) - read_obj_stream.socket_like_rpc.open = AsyncMock() - - recv_response = mock.MagicMock(spec=_storage_v2.BidiReadObjectResponse) - recv_response.metadata = mock.MagicMock(spec=_storage_v2.Object) - recv_response.metadata.generation = _TEST_GENERATION_NUMBER - recv_response.read_handle = _TEST_READ_HANDLE - read_obj_stream.socket_like_rpc.recv = AsyncMock(return_value=recv_response) # act await read_obj_stream.open() @@ -88,40 +111,88 @@ async def test_open(mock_client): assert read_obj_stream.generation_number == _TEST_GENERATION_NUMBER assert read_obj_stream.read_handle == _TEST_READ_HANDLE + assert read_obj_stream.is_stream_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_close(mock_client): +async def test_open_when_already_open_should_raise_error( + mock_client, mock_cls_async_bidi_rpc +): # arrange - read_obj_stream = _AsyncReadObjectStream( - client=mock_client, - bucket_name=_TEST_BUCKET_NAME, - object_name=_TEST_OBJECT_NAME, + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=True + ) + + # act + assert (pythonic) + with pytest.raises(ValueError) as exc: + await read_obj_stream.open() + + # assert + assert str(exc.value) == "Stream is already 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_close(mock_client, mock_cls_async_bidi_rpc): + # arrange + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=True ) - read_obj_stream.socket_like_rpc.close = AsyncMock() # act await read_obj_stream.close() # assert read_obj_stream.socket_like_rpc.close.assert_called_once() + assert not read_obj_stream.is_stream_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_send(mock_client): +async def test_close_without_open_should_raise_error( + mock_client, mock_cls_async_bidi_rpc +): # arrange - read_obj_stream = _AsyncReadObjectStream( - client=mock_client, - bucket_name=_TEST_BUCKET_NAME, - object_name=_TEST_OBJECT_NAME, + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False + ) + + # act + assert (pythonic) + with pytest.raises(ValueError) as exc: + await read_obj_stream.close() + + # 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_send(mock_client, mock_cls_async_bidi_rpc): + # arrange + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=True ) - read_obj_stream.socket_like_rpc.send = AsyncMock() # act bidi_read_object_request = _storage_v2.BidiReadObjectRequest() @@ -133,16 +204,40 @@ async def test_send(mock_client): ) +@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(mock_client): +async def test_send_without_open_should_raise_error( + mock_client, mock_cls_async_bidi_rpc +): # arrange - read_obj_stream = _AsyncReadObjectStream( - client=mock_client, - bucket_name=_TEST_BUCKET_NAME, - object_name=_TEST_OBJECT_NAME, + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False + ) + + # act + assert (pythonic) + with pytest.raises(ValueError) as exc: + await read_obj_stream.send(_storage_v2.BidiReadObjectRequest()) + + # 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(mock_client, mock_cls_async_bidi_rpc): + # arrange + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=True ) bidi_read_object_response = _storage_v2.BidiReadObjectResponse() read_obj_stream.socket_like_rpc.recv = AsyncMock( @@ -155,3 +250,26 @@ async def test_recv(mock_client): # assert read_obj_stream.socket_like_rpc.recv.assert_called_once() assert response == bidi_read_object_response + + +@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_without_open_should_raise_error( + mock_client, mock_cls_async_bidi_rpc +): + # arrange + read_obj_stream = await instantiate_read_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False + ) + + # act + assert (pythonic) + with pytest.raises(ValueError) as exc: + await read_obj_stream.recv() + + # assert + assert str(exc.value) == "Stream is not open"