feat(zb-experimental): implement close (#1614)

feat(zb-experimental): implement close

Author: chandra-siri

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

@@ -145,9 +145,10 @@ async def open(self) -> None:
 
 asyncdefclose(self) ->None:
 """Closes the bidi-gRPC connection."""
-raiseNotImplementedError(
-"close() is not implemented yet in _AsyncWriteObjectStream"
- )
+ifnotself._is_stream_open:
+raiseValueError("Stream is not open")
+awaitself.socket_like_rpc.close()
+self._is_stream_open=False
 
 asyncdefsend(
 self, bidi_write_object_request: _storage_v2.BidiWriteObjectRequest

tests/unit/asyncio/test_async_write_object_stream.py

@@ -15,6 +15,7 @@
 importpytest
 fromunittestimportmock
 
+fromunittest.mockimportAsyncMock
 fromgoogle.cloud.storage._experimental.asyncio.async_write_object_streamimport (
 _AsyncWriteObjectStream,
 )
@@ -43,6 +44,27 @@ def mock_client():
 returnclient
 
 
+asyncdefinstantiate_write_obj_stream(mock_client, mock_cls_async_bidi_rpc, open=True):
+"""Helper to create an instance of _AsyncWriteObjectStream and open it by default."""
+socket_like_rpc=AsyncMock()
+mock_cls_async_bidi_rpc.return_value=socket_like_rpc
+socket_like_rpc.open=AsyncMock()
+socket_like_rpc.close=AsyncMock()
+
+mock_response=mock.MagicMock(spec=_storage_v2.BidiWriteObjectResponse)
+mock_response.resource=mock.MagicMock(spec=_storage_v2.Object)
+mock_response.resource.generation=GENERATION
+mock_response.write_handle=WRITE_HANDLE
+socket_like_rpc.recv=AsyncMock(return_value=mock_response)
+
+write_obj_stream=_AsyncWriteObjectStream(mock_client, BUCKET, OBJECT)
+
+ifopen:
+awaitwrite_obj_stream.open()
+
+returnwrite_obj_stream
+
+
 deftest_async_write_object_stream_init(mock_client):
 """Test the constructor of _AsyncWriteObjectStream."""
 stream=_AsyncWriteObjectStream(mock_client, BUCKET, OBJECT)
@@ -228,7 +250,6 @@ async def test_open_raises_error_on_missing_generation(
 ValueError, match="Failed to obtain object generation after opening the stream"
  ):
 awaitstream.open()
-# assert stream.generation_number is None
 
 
 @pytest.mark.asyncio
@@ -252,13 +273,49 @@ async def test_open_raises_error_on_missing_write_handle(
 
 
 @pytest.mark.asyncio
-asyncdeftest_unimplemented_methods_raise_error(mock_client):
-"""Test that unimplemented methods raise NotImplementedError."""
-stream=_AsyncWriteObjectStream(mock_client, BUCKET, OBJECT)
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc"
+)
+asyncdeftest_close(mock_cls_async_bidi_rpc, mock_client):
+"""Test that close successfully closes the stream."""
+# Arrange
+write_obj_stream=awaitinstantiate_write_obj_stream(
+mock_client, mock_cls_async_bidi_rpc, open=True
+ )
 
-withpytest.raises(NotImplementedError):
-awaitstream.close()
+# Act
+awaitwrite_obj_stream.close()
 
+# Assert
+write_obj_stream.socket_like_rpc.close.assert_called_once()
+assertnotwrite_obj_stream.is_stream_open
+
+
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc"
+)
+asyncdeftest_close_without_open_should_raise_error(
+mock_cls_async_bidi_rpc, mock_client
+):
+"""Test that closing a stream that is not open raises a ValueError."""
+# Arrange
+write_obj_stream=awaitinstantiate_write_obj_stream(
+mock_client, mock_cls_async_bidi_rpc, open=False
+ )
+
+# Act & Assert
+withpytest.raises(ValueError, match="Stream is not open"):
+awaitwrite_obj_stream.close()
+
+
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_write_object_stream.AsyncBidiRpc"
+)
+asyncdeftest_unimplemented_methods_raise_error(mock_async_bidi_rpc, mock_client):
+"""Test that unimplemented methods (send, recv) raise NotImplementedError."""
+stream=_AsyncWriteObjectStream(mock_client, BUCKET, OBJECT)
 withpytest.raises(NotImplementedError):
 awaitstream.send(_storage_v2.BidiWriteObjectRequest())