feat(zb-experimental): implement flush, close and finalize (#1619)

feat(zb-experimental): implement flush, close and finalize

Author: chandra-siri

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

@@ -114,7 +114,11 @@ def __init__(
 self.persisted_size: Optional[int] =None
 
 asyncdefstate_lookup(self) ->int:
-"""Returns the persisted_size."""
+"""Returns the persisted_size
+
+ :rtype: int
+ :returns: persisted size.
+ """
 awaitself.write_obj_stream.send(
 _storage_v2.BidiWriteObjectRequest(
 state_lookup=True,
@@ -142,18 +146,44 @@ async def append(self, data: bytes):
 raiseNotImplementedError("append is not implemented yet.")
 
 asyncdefflush(self) ->int:
-"""Returns persisted_size"""
-raiseNotImplementedError("flush is not implemented yet.")
+"""Flushes the data to the server.
+
+ :rtype: int
+ :returns: The persisted size after flush.
+ """
+awaitself.write_obj_stream.send(
+_storage_v2.BidiWriteObjectRequest(
+flush=True,
+state_lookup=True,
+ )
+ )
+response=awaitself.write_obj_stream.recv()
+self.persisted_size=response.persisted_size
+self.offset=self.persisted_size
+returnself.persisted_size
 
 asyncdefclose(self, finalize_on_close=False) ->int:
 """Returns persisted_size"""
-raiseNotImplementedError("close is not implemented yet.")
+iffinalize_on_close:
+awaitself.finalize()
+
+awaitself.write_obj_stream.close()
+self._is_stream_open=False
+self.offset=None
+
+asyncdeffinalize(self) ->_storage_v2.Object:
+"""Finalizes the Appendable Object.
 
-asyncdeffinalize(self) ->int:
-"""Returns persisted_size
  Note: Once finalized no more data can be appended.
+
+ rtype: google.cloud.storage_v2.types.Object
+ returns: The finalized object resource.
  """
-raiseNotImplementedError("finalize is not implemented yet.")
+awaitself.write_obj_stream.send(
+_storage_v2.BidiWriteObjectRequest(finish_write=True)
+ )
+response=awaitself.write_obj_stream.recv()
+self.object_resource=response.resource
 
 # helper methods.
 asyncdefappend_from_string(self, data: str):

tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -171,15 +171,6 @@ async def test_unimplemented_methods_raise_error(mock_client):
 withpytest.raises(NotImplementedError):
 awaitwriter.append(b"data")
 
-withpytest.raises(NotImplementedError):
-awaitwriter.flush()
-
-withpytest.raises(NotImplementedError):
-awaitwriter.close()
-
-withpytest.raises(NotImplementedError):
-awaitwriter.finalize()
-
 withpytest.raises(NotImplementedError):
 awaitwriter.append_from_string("data")
 
@@ -188,3 +179,91 @@ async def test_unimplemented_methods_raise_error(mock_client):
 
 withpytest.raises(NotImplementedError):
 awaitwriter.append_from_file("file.txt")
+
+
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+asyncdeftest_flush(mock_write_object_stream, mock_client):
+"""Test that flush sends the correct request and updates state."""
+writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
+mock_stream=mock_write_object_stream.return_value
+mock_stream.send=mock.AsyncMock()
+mock_stream.recv=mock.AsyncMock(
+return_value=_storage_v2.BidiWriteObjectResponse(persisted_size=1024)
+ )
+
+persisted_size=awaitwriter.flush()
+
+expected_request=_storage_v2.BidiWriteObjectRequest(flush=True, state_lookup=True)
+mock_stream.send.assert_awaited_once_with(expected_request)
+mock_stream.recv.assert_awaited_once()
+assertwriter.persisted_size==1024
+assertwriter.offset==1024
+assertpersisted_size==1024
+
+
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+asyncdeftest_close_without_finalize(mock_write_object_stream, mock_client):
+"""Test close without finalizing."""
+writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
+writer._is_stream_open=True
+writer.offset=1024
+mock_stream=mock_write_object_stream.return_value
+mock_stream.close=mock.AsyncMock()
+writer.finalize=mock.AsyncMock()
+
+awaitwriter.close(finalize_on_close=False)
+
+writer.finalize.assert_not_awaited()
+mock_stream.close.assert_awaited_once()
+assertnotwriter._is_stream_open
+assertwriter.offsetisNone
+
+
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+asyncdeftest_close_with_finalize(mock_write_object_stream, mock_client):
+"""Test close with finalizing."""
+writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
+writer._is_stream_open=True
+writer.offset=1024
+mock_stream=mock_write_object_stream.return_value
+mock_stream.close=mock.AsyncMock()
+writer.finalize=mock.AsyncMock()
+
+awaitwriter.close(finalize_on_close=True)
+
+writer.finalize.assert_awaited_once()
+mock_stream.close.assert_awaited_once()
+assertnotwriter._is_stream_open
+assertwriter.offsetisNone
+
+
+@pytest.mark.asyncio
+@mock.patch(
+"google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+asyncdeftest_finalize(mock_write_object_stream, mock_client):
+"""Test that finalize sends the correct request and updates state."""
+writer=AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
+mock_stream=mock_write_object_stream.return_value
+mock_stream.send=mock.AsyncMock()
+mock_resource=_storage_v2.Object(name=OBJECT, bucket=BUCKET)
+mock_stream.recv=mock.AsyncMock(
+return_value=_storage_v2.BidiWriteObjectResponse(resource=mock_resource)
+ )
+
+awaitwriter.finalize()
+
+mock_stream.send.assert_awaited_once_with(
+_storage_v2.BidiWriteObjectRequest(finish_write=True)
+ )
+mock_stream.recv.assert_awaited_once()
+assertwriter.object_resource==mock_resource