This repository was archived by the owner on Mar 31, 2026. It is now read-only.
Commit 7d17922
authored
File tree
- google/cloud/storage/_experimental/asyncio
- tests/unit/asyncio
Lines changed: 10 additions & 2 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -22,6 +22,7 @@ | |||
| 22 | 22 | ||
| 23 | 23 | """ | |
| 24 | 24 | from typing import Optional | |
| 25 | + from google.cloud import _storage_v2 | ||
| 25 | 26 | from google.cloud.storage._experimental.asyncio.async_grpc_client import ( | |
| 26 | 27 | AsyncGrpcClient, | |
| 27 | 28 | ) | |
@@ -112,9 +113,16 @@ def __init__( | |||
| 112 | 113 | self.offset: Optional[int] = None | |
| 113 | 114 | self.persisted_size: Optional[int] = None | |
| 114 | 115 | ||
| 115 | - async def state_lookup(self): | ||
| 116 | + async def state_lookup(self) -> int: | ||
| 116 | 117 | """Returns the persisted_size.""" | |
| 117 | - raise NotImplementedError("state_lookup is not implemented yet.") | ||
| 118 | + await self.write_obj_stream.send( | ||
| 119 | + _storage_v2.BidiWriteObjectRequest( | ||
| 120 | + state_lookup=True, | ||
| 121 | + ) | ||
| 122 | + ) | ||
| 123 | + response = await self.write_obj_stream.recv() | ||
| 124 | + self.persisted_size = response.persisted_size | ||
| 125 | + return self.persisted_size | ||
| 118 | 126 | ||
| 119 | 127 | async def open(self) -> None: | |
| 120 | 128 | """Opens the underlying bidi-gRPC stream.""" | |
Lines changed: 29 additions & 3 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -18,11 +18,14 @@ | |||
| 18 | 18 | from google.cloud.storage._experimental.asyncio.async_appendable_object_writer import ( | |
| 19 | 19 | AsyncAppendableObjectWriter, | |
| 20 | 20 | ) | |
| 21 | + from google.cloud import _storage_v2 | ||
| 22 | + | ||
| 21 | 23 | ||
| 22 | 24 | BUCKET = "test-bucket" | |
| 23 | 25 | OBJECT = "test-object" | |
| 24 | 26 | GENERATION = 123 | |
| 25 | 27 | WRITE_HANDLE = b"test-write-handle" | |
| 28 | + PERSISTED_SIZE = 456 | ||
| 26 | 29 | ||
| 27 | 30 | ||
| 28 | 31 | @pytest.fixture | |
@@ -82,14 +85,37 @@ def test_init_with_optional_args(mock_write_object_stream, mock_client): | |||
| 82 | 85 | ) | |
| 83 | 86 | ||
| 84 | 87 | ||
| 88 | + @pytest.mark.asyncio | ||
| 89 | + @mock.patch( | ||
| 90 | + "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream" | ||
| 91 | + ) | ||
| 92 | + async def test_state_lookup(mock_write_object_stream, mock_client): | ||
| 93 | + """Test state_lookup method.""" | ||
| 94 | + # Arrange | ||
| 95 | + writer = AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT) | ||
| 96 | + mock_stream = mock_write_object_stream.return_value | ||
| 97 | + mock_stream.send = mock.AsyncMock() | ||
| 98 | + mock_stream.recv = mock.AsyncMock( | ||
| 99 | + return_value=_storage_v2.BidiWriteObjectResponse(persisted_size=PERSISTED_SIZE) | ||
| 100 | + ) | ||
| 101 | + | ||
| 102 | + expected_request = _storage_v2.BidiWriteObjectRequest(state_lookup=True) | ||
| 103 | + | ||
| 104 | + # Act | ||
| 105 | + response = await writer.state_lookup() | ||
| 106 | + | ||
| 107 | + # Assert | ||
| 108 | + mock_stream.send.assert_awaited_once_with(expected_request) | ||
| 109 | + mock_stream.recv.assert_awaited_once() | ||
| 110 | + assert writer.persisted_size == PERSISTED_SIZE | ||
| 111 | + assert response == PERSISTED_SIZE | ||
| 112 | + | ||
| 113 | + | ||
| 85 | 114 | @pytest.mark.asyncio | |
| 86 | 115 | async def test_unimplemented_methods_raise_error(mock_client): | |
| 87 | 116 | """Test that all currently unimplemented methods raise NotImplementedError.""" | |
| 88 | 117 | writer = AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT) | |
| 89 | 118 | ||
| 90 | - with pytest.raises(NotImplementedError): | ||
| 91 | - await writer.state_lookup() | ||
| 92 | - | ||
| 93 | 119 | with pytest.raises(NotImplementedError): | |
| 94 | 120 | await writer.open() | |
| 95 | 121 | ||
0 commit comments