feat(zb-experimental): implement state_lookup (#1617)

feat(zb-experimental): implement state_lookup

Author: chandra-siri

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

@@ -22,6 +22,7 @@
 
 """
 from typing import Optional
+from google.cloud import _storage_v2
 from google.cloud.storage._experimental.asyncio.async_grpc_client import (
     AsyncGrpcClient,
 )
@@ -112,9 +113,16 @@ def __init__(
         self.offset: Optional[int] = None
         self.persisted_size: Optional[int] = None
 
-    async def state_lookup(self):
+    async def state_lookup(self) -> int:
         """Returns the persisted_size."""
-        raise NotImplementedError("state_lookup is not implemented yet.")
+        await self.write_obj_stream.send(
+            _storage_v2.BidiWriteObjectRequest(
+                state_lookup=True,
+            )
+        )
+        response = await self.write_obj_stream.recv()
+        self.persisted_size = response.persisted_size
+        return self.persisted_size
 
     async def open(self) -> None:
         """Opens the underlying bidi-gRPC stream."""

tests/unit/asyncio/test_async_appendable_object_writer.py

@@ -18,11 +18,14 @@
 from google.cloud.storage._experimental.asyncio.async_appendable_object_writer import (
     AsyncAppendableObjectWriter,
 )
+from google.cloud import _storage_v2
+
 
 BUCKET = "test-bucket"
 OBJECT = "test-object"
 GENERATION = 123
 WRITE_HANDLE = b"test-write-handle"
+PERSISTED_SIZE = 456
 
 
 @pytest.fixture
@@ -82,14 +85,37 @@ def test_init_with_optional_args(mock_write_object_stream, mock_client):
     )
 
 
+@pytest.mark.asyncio
+@mock.patch(
+    "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream"
+)
+async def test_state_lookup(mock_write_object_stream, mock_client):
+    """Test state_lookup method."""
+    # Arrange
+    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=PERSISTED_SIZE)
+    )
+
+    expected_request = _storage_v2.BidiWriteObjectRequest(state_lookup=True)
+
+    # Act
+    response = await writer.state_lookup()
+
+    # Assert
+    mock_stream.send.assert_awaited_once_with(expected_request)
+    mock_stream.recv.assert_awaited_once()
+    assert writer.persisted_size == PERSISTED_SIZE
+    assert response == PERSISTED_SIZE
+
+
 @pytest.mark.asyncio
 async def test_unimplemented_methods_raise_error(mock_client):
     """Test that all currently unimplemented methods raise NotImplementedError."""
     writer = AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT)
 
-    with pytest.raises(NotImplementedError):
-        await writer.state_lookup()
-
     with pytest.raises(NotImplementedError):
         await writer.open()