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
 
     async def state_lookup(self) -> int:
-        """Returns the persisted_size."""
+        """Returns the persisted_size
+
+        :rtype: int
+        :returns: persisted size.
+        """
         await self.write_obj_stream.send(
             _storage_v2.BidiWriteObjectRequest(
                 state_lookup=True,
@@ -142,18 +146,44 @@ async def append(self, data: bytes):
         raise NotImplementedError("append is not implemented yet.")
 
     async def flush(self) -> int:
-        """Returns persisted_size"""
-        raise NotImplementedError("flush is not implemented yet.")
+        """Flushes the data to the server.
+
+        :rtype: int
+        :returns: The persisted size after flush.
+        """
+        await self.write_obj_stream.send(
+            _storage_v2.BidiWriteObjectRequest(
+                flush=True,
+                state_lookup=True,
+            )
+        )
+        response = await self.write_obj_stream.recv()
+        self.persisted_size = response.persisted_size
+        self.offset = self.persisted_size
+        return self.persisted_size
 
     async def close(self, finalize_on_close=False) -> int:
         """Returns persisted_size"""
-        raise NotImplementedError("close is not implemented yet.")
+        if finalize_on_close:
+            await self.finalize()
+
+        await self.write_obj_stream.close()
+        self._is_stream_open = False
+        self.offset = None
+
+    async def finalize(self) -> _storage_v2.Object:
+        """Finalizes the Appendable Object.
 
-    async def finalize(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.
         """
-        raise NotImplementedError("finalize is not implemented yet.")
+        await self.write_obj_stream.send(
+            _storage_v2.BidiWriteObjectRequest(finish_write=True)
+        )
+        response = await self.write_obj_stream.recv()
+        self.object_resource = response.resource
 
     # helper methods.
     async def append_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):
     with pytest.raises(NotImplementedError):
         await writer.append(b"data")
 
-    with pytest.raises(NotImplementedError):
-        await writer.flush()
-
-    with pytest.raises(NotImplementedError):
-        await writer.close()
-
-    with pytest.raises(NotImplementedError):
-        await writer.finalize()
-
     with pytest.raises(NotImplementedError):
         await writer.append_from_string("data")
 
@@ -188,3 +179,91 @@ async def test_unimplemented_methods_raise_error(mock_client):
 
     with pytest.raises(NotImplementedError):
         await writer.append_from_file("file.txt")
+
+
+@pytest.mark.asyncio
+@mock.patch(
+    "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+async def test_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 = await writer.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()
+    assert writer.persisted_size == 1024
+    assert writer.offset == 1024
+    assert persisted_size == 1024
+
+
+@pytest.mark.asyncio
+@mock.patch(
+    "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+async def test_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()
+
+    await writer.close(finalize_on_close=False)
+
+    writer.finalize.assert_not_awaited()
+    mock_stream.close.assert_awaited_once()
+    assert not writer._is_stream_open
+    assert writer.offset is None
+
+
+@pytest.mark.asyncio
+@mock.patch(
+    "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+async def test_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()
+
+    await writer.close(finalize_on_close=True)
+
+    writer.finalize.assert_awaited_once()
+    mock_stream.close.assert_awaited_once()
+    assert not writer._is_stream_open
+    assert writer.offset is None
+
+
+@pytest.mark.asyncio
+@mock.patch(
+    "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+async def test_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)
+    )
+
+    await writer.finalize()
+
+    mock_stream.send.assert_awaited_once_with(
+        _storage_v2.BidiWriteObjectRequest(finish_write=True)
+    )
+    mock_stream.recv.assert_awaited_once()
+    assert writer.object_resource == mock_resource