feat(zb-experimental): Add support for closing MRD (#1554)

* add `is_stream_open` and check it's value before doing any of the following operation. `open` , `close` , `send` ,`recv` * Instantiate `socket_like_rpc` inside open, so that it is a new object every time we open. * Otherwise, consider this scenario where a user is opening after closing then `initial_request` would be None because close sets it to None.

Author: chandra-siri

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

@@ -83,32 +83,41 @@ 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.
 
  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:
  self.generation_number = response.metadata.generation
 
  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

tests/unit/asyncio/test_async_read_object_stream.py

@@ -27,16 +27,42 @@
 _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"
 )
 @mock.patch(
  "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"