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

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"