fix(experimental): implement requests_done method to signal end of requests in async streams. Gracefully close streams. (#1700)

fix(experimental): implement requests_done method to signal end of
requests in async streams. Gracefully close streams.

Author: chandra-siri

google/cloud/storage/_experimental/asyncio/async_read_object_stream.py

@@ -151,9 +151,16 @@ async def close(self) -> None:
 """Closes the bidi-gRPC connection."""
 ifnotself._is_stream_open:
 raiseValueError("Stream is not open")
+awaitself.requests_done()
 awaitself.socket_like_rpc.close()
 self._is_stream_open=False
 
+asyncdefrequests_done(self):
+"""Signals that all requests have been sent."""
+
+awaitself.socket_like_rpc.send(None)
+awaitself.socket_like_rpc.recv()
+
 asyncdefsend(
 self, bidi_read_object_request: _storage_v2.BidiReadObjectRequest
  ) ->None:

google/cloud/storage/_experimental/asyncio/async_write_object_stream.py

@@ -152,9 +152,16 @@ async def close(self) -> None:
 """Closes the bidi-gRPC connection."""
 ifnotself._is_stream_open:
 raiseValueError("Stream is not open")
+awaitself.requests_done()
 awaitself.socket_like_rpc.close()
 self._is_stream_open=False
 
+asyncdefrequests_done(self):
+"""Signals that all requests have been sent."""
+
+awaitself.socket_like_rpc.send(None)
+awaitself.socket_like_rpc.recv()
+
 asyncdefsend(
 self, bidi_write_object_request: _storage_v2.BidiWriteObjectRequest
  ) ->None:
@@ -186,3 +193,4 @@ async def recv(self) -> _storage_v2.BidiWriteObjectResponse:
 @property
 defis_stream_open(self) ->bool:
 returnself._is_stream_open
+

tests/unit/asyncio/test_async_read_object_stream.py

@@ -197,15 +197,41 @@ async def test_close(mock_client, mock_cls_async_bidi_rpc):
 read_obj_stream=awaitinstantiate_read_obj_stream(
 mock_client, mock_cls_async_bidi_rpc, open=True
  )
+read_obj_stream.requests_done=AsyncMock()
 
 # act
 awaitread_obj_stream.close()
 
 # assert
+read_obj_stream.requests_done.assert_called_once()
 read_obj_stream.socket_like_rpc.close.assert_called_once()
 assertnotread_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
+asyncdeftest_requests_done(mock_client, mock_cls_async_bidi_rpc):
+"""Test that requests_done signals the end of requests."""
+# Arrange
+read_obj_stream=awaitinstantiate_read_obj_stream(
+mock_client, mock_cls_async_bidi_rpc, open=True
+ )
+read_obj_stream.socket_like_rpc.send=AsyncMock()
+read_obj_stream.socket_like_rpc.recv=AsyncMock()
+
+# Act
+awaitread_obj_stream.requests_done()
+
+# Assert
+read_obj_stream.socket_like_rpc.send.assert_called_once_with(None)
+read_obj_stream.socket_like_rpc.recv.assert_called_once()
+
+
 @mock.patch(
 "google.cloud.storage._experimental.asyncio.async_read_object_stream.AsyncBidiRpc"
 )

tests/unit/asyncio/test_async_write_object_stream.py

@@ -289,11 +289,13 @@ async def test_close(mock_cls_async_bidi_rpc, mock_client):
 write_obj_stream=awaitinstantiate_write_obj_stream(
 mock_client, mock_cls_async_bidi_rpc, open=True
  )
+write_obj_stream.requests_done=AsyncMock()
 
 # Act
 awaitwrite_obj_stream.close()
 
 # Assert
+write_obj_stream.requests_done.assert_called_once()
 write_obj_stream.socket_like_rpc.close.assert_called_once()
 assertnotwrite_obj_stream.is_stream_open
 
@@ -394,3 +396,24 @@ async def test_recv_without_open_should_raise_error(
 # Act & Assert
 withpytest.raises(ValueError, match="Stream is not open"):
 awaitwrite_obj_stream.recv()
+
+
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc"
+)
+asyncdeftest_requests_done(mock_cls_async_bidi_rpc, mock_client):
+"""Test that requests_done signals the end of requests."""
+# Arrange
+write_obj_stream=awaitinstantiate_write_obj_stream(
+mock_client, mock_cls_async_bidi_rpc, open=True
+ )
+write_obj_stream.socket_like_rpc.send=AsyncMock()
+write_obj_stream.socket_like_rpc.recv=AsyncMock()
+
+# Act
+awaitwrite_obj_stream.requests_done()
+
+# Assert
+write_obj_stream.socket_like_rpc.send.assert_called_once_with(None)
+write_obj_stream.socket_like_rpc.recv.assert_called_once()