diff --git a/google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py b/google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py index ca03c2b6d..be992b641 100644 --- a/google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py +++ b/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.""" diff --git a/tests/unit/asyncio/test_async_appendable_object_writer.py b/tests/unit/asyncio/test_async_appendable_object_writer.py index cf838ac05..d3a6d3830 100644 --- a/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/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()