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 5c3e54d71..d34c844d5 100644 --- a/google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py +++ b/google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py @@ -199,10 +199,28 @@ async def append(self, data: bytes) -> None: self.offset += chunk_size bytes_to_flush += chunk_size if bytes_to_flush >= _MAX_BUFFER_SIZE_BYTES: - await self.flush() + await self.simple_flush() bytes_to_flush = 0 start_idx = end_idx + async def simple_flush(self) -> None: + """Flushes the data to the server. + Please note: Unlike `flush` it does not do `state_lookup` + + :rtype: None + + :raises ValueError: If the stream is not open (i.e., `open()` has not + been called). + """ + if not self._is_stream_open: + raise ValueError("Stream is not open. Call open() before simple_flush().") + + await self.write_obj_stream.send( + _storage_v2.BidiWriteObjectRequest( + flush=True, + ) + ) + async def flush(self) -> int: """Flushes the data to the server. diff --git a/tests/unit/asyncio/test_async_appendable_object_writer.py b/tests/unit/asyncio/test_async_appendable_object_writer.py index 902ed5f8b..a75824f8b 100644 --- a/tests/unit/asyncio/test_async_appendable_object_writer.py +++ b/tests/unit/asyncio/test_async_appendable_object_writer.py @@ -224,6 +224,38 @@ async def test_flush_without_open_raises_value_error(mock_client): await writer.flush() +@pytest.mark.asyncio +@mock.patch( + "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream" +) +async def test_simple_flush(mock_write_object_stream, mock_client): + """Test that flush sends the correct request and updates state.""" + # Arrange + writer = AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT) + writer._is_stream_open = True + mock_stream = mock_write_object_stream.return_value + mock_stream.send = mock.AsyncMock() + + # Act + await writer.simple_flush() + + # Assert + mock_stream.send.assert_awaited_once_with( + _storage_v2.BidiWriteObjectRequest(flush=True) + ) + + +@pytest.mark.asyncio +async def test_simple_flush_without_open_raises_value_error(mock_client): + """Test that flush raises an error if the stream is not open.""" + writer = AsyncAppendableObjectWriter(mock_client, BUCKET, OBJECT) + with pytest.raises( + ValueError, + match="Stream is not open. Call open\\(\\) before simple_flush\\(\\).", + ): + await writer.simple_flush() + + @pytest.mark.asyncio @mock.patch( "google.cloud.storage._experimental.asyncio.async_appendable_object_writer._AsyncWriteObjectStream" @@ -369,7 +401,7 @@ async def test_append_sends_data_in_chunks(mock_write_object_stream, mock_client writer.persisted_size = 100 mock_stream = mock_write_object_stream.return_value mock_stream.send = mock.AsyncMock() - writer.flush = mock.AsyncMock() + writer.simple_flush = mock.AsyncMock() data = b"a" * (_MAX_CHUNK_SIZE_BYTES + 1) await writer.append(data) @@ -387,7 +419,7 @@ async def test_append_sends_data_in_chunks(mock_write_object_stream, mock_client assert len(second_call[0][0].checksummed_data.content) == 1 assert writer.offset == 100 + len(data) - writer.flush.assert_not_awaited() + writer.simple_flush.assert_not_awaited() @pytest.mark.asyncio @@ -407,12 +439,12 @@ async def test_append_flushes_when_buffer_is_full( writer.persisted_size = 0 mock_stream = mock_write_object_stream.return_value mock_stream.send = mock.AsyncMock() - writer.flush = mock.AsyncMock() + writer.simple_flush = mock.AsyncMock() data = b"a" * _MAX_BUFFER_SIZE_BYTES await writer.append(data) - writer.flush.assert_awaited_once() + writer.simple_flush.assert_awaited_once() @pytest.mark.asyncio @@ -430,12 +462,12 @@ async def test_append_handles_large_data(mock_write_object_stream, mock_client): writer.persisted_size = 0 mock_stream = mock_write_object_stream.return_value mock_stream.send = mock.AsyncMock() - writer.flush = mock.AsyncMock() + writer.simple_flush = mock.AsyncMock() data = b"a" * (_MAX_BUFFER_SIZE_BYTES * 2 + 1) await writer.append(data) - assert writer.flush.await_count == 2 + assert writer.simple_flush.await_count == 2 @pytest.mark.asyncio @@ -453,7 +485,7 @@ async def test_append_data_two_times(mock_write_object_stream, mock_client): writer.persisted_size = 0 mock_stream = mock_write_object_stream.return_value mock_stream.send = mock.AsyncMock() - writer.flush = mock.AsyncMock() + writer.simple_flush = mock.AsyncMock() data1 = b"a" * (_MAX_CHUNK_SIZE_BYTES + 10) await writer.append(data1) @@ -463,4 +495,4 @@ async def test_append_data_two_times(mock_write_object_stream, mock_client): total_data_length = len(data1) + len(data2) assert writer.offset == total_data_length - assert writer.flush.await_count == 0 + assert writer.simple_flush.await_count == 0