diff --git a/google/cloud/storage/_experimental/asyncio/async_write_object_stream.py b/google/cloud/storage/_experimental/asyncio/async_write_object_stream.py index b84dea5a4..6d1fd5b31 100644 --- a/google/cloud/storage/_experimental/asyncio/async_write_object_stream.py +++ b/google/cloud/storage/_experimental/asyncio/async_write_object_stream.py @@ -160,9 +160,9 @@ async def send( The request message to send. This is typically used to specify the read offset and limit. """ - raise NotImplementedError( - "send() is not implemented yet in _AsyncWriteObjectStream" - ) + if not self._is_stream_open: + raise ValueError("Stream is not open") + await self.socket_like_rpc.send(bidi_write_object_request) async def recv(self) -> _storage_v2.BidiWriteObjectResponse: """Receives a response from the stream. @@ -174,9 +174,9 @@ async def recv(self) -> _storage_v2.BidiWriteObjectResponse: :class:`~google.cloud._storage_v2.types.BidiWriteObjectResponse`: The response message from the server. """ - raise NotImplementedError( - "recv() is not implemented yet in _AsyncWriteObjectStream" - ) + 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: diff --git a/tests/unit/asyncio/test_async_write_object_stream.py b/tests/unit/asyncio/test_async_write_object_stream.py index ad83669c3..7fa2123c5 100644 --- a/tests/unit/asyncio/test_async_write_object_stream.py +++ b/tests/unit/asyncio/test_async_write_object_stream.py @@ -49,6 +49,7 @@ async def instantiate_write_obj_stream(mock_client, mock_cls_async_bidi_rpc, ope socket_like_rpc = AsyncMock() mock_cls_async_bidi_rpc.return_value = socket_like_rpc socket_like_rpc.open = AsyncMock() + socket_like_rpc.send = AsyncMock() socket_like_rpc.close = AsyncMock() mock_response = mock.MagicMock(spec=_storage_v2.BidiWriteObjectResponse) @@ -313,11 +314,77 @@ async def test_close_without_open_should_raise_error( @mock.patch( "google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc" ) -async def test_unimplemented_methods_raise_error(mock_async_bidi_rpc, mock_client): - """Test that unimplemented methods (send, recv) raise NotImplementedError.""" - stream = _AsyncWriteObjectStream(mock_client, BUCKET, OBJECT) - with pytest.raises(NotImplementedError): - await stream.send(_storage_v2.BidiWriteObjectRequest()) +async def test_send(mock_cls_async_bidi_rpc, mock_client): + """Test that send calls the underlying rpc's send method.""" + # Arrange + write_obj_stream = await instantiate_write_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=True + ) + + # Act + bidi_write_object_request = _storage_v2.BidiWriteObjectRequest() + await write_obj_stream.send(bidi_write_object_request) + + # Assert + write_obj_stream.socket_like_rpc.send.assert_called_once_with( + bidi_write_object_request + ) + + +@pytest.mark.asyncio +@mock.patch( + "google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc" +) +async def test_send_without_open_should_raise_error( + mock_cls_async_bidi_rpc, mock_client +): + """Test that sending on a stream that is not open raises a ValueError.""" + # Arrange + write_obj_stream = await instantiate_write_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False + ) + + # Act & Assert + with pytest.raises(ValueError, match="Stream is not open"): + await write_obj_stream.send(_storage_v2.BidiWriteObjectRequest()) + + +@pytest.mark.asyncio +@mock.patch( + "google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc" +) +async def test_recv(mock_cls_async_bidi_rpc, mock_client): + """Test that recv calls the underlying rpc's recv method.""" + # Arrange + write_obj_stream = await instantiate_write_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=True + ) + bidi_write_object_response = _storage_v2.BidiWriteObjectResponse() + write_obj_stream.socket_like_rpc.recv = AsyncMock( + return_value=bidi_write_object_response + ) + + # Act + response = await write_obj_stream.recv() + + # Assert + write_obj_stream.socket_like_rpc.recv.assert_called_once() + assert response == bidi_write_object_response + - with pytest.raises(NotImplementedError): - await stream.recv() +@pytest.mark.asyncio +@mock.patch( + "google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc" +) +async def test_recv_without_open_should_raise_error( + mock_cls_async_bidi_rpc, mock_client +): + """Test that receiving on a stream that is not open raises a ValueError.""" + # Arrange + write_obj_stream = await instantiate_write_obj_stream( + mock_client, mock_cls_async_bidi_rpc, open=False + ) + + # Act & Assert + with pytest.raises(ValueError, match="Stream is not open"): + await write_obj_stream.recv()