From d015a58bdfeb1ad171a8608d16335cff7839a205 Mon Sep 17 00:00:00 2001 From: JordonPhillips Date: Mon, 25 Nov 2024 15:18:46 +0100 Subject: [PATCH] Flush AsyncBytesProvider on close by default This updates the default behavior of close on AsyncBytesProvider to flush already-written data. This should be the less surprising outcome, particularly when the class is used as a context manager. --- .../smithy-core/smithy_core/aio/types.py | 4 +-- .../smithy-core/tests/unit/aio/test_types.py | 32 +++++++++++++++++-- 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/python-packages/smithy-core/smithy_core/aio/types.py b/python-packages/smithy-core/smithy_core/aio/types.py index caedbcdb4..492425027 100644 --- a/python-packages/smithy-core/smithy_core/aio/types.py +++ b/python-packages/smithy-core/smithy_core/aio/types.py @@ -367,7 +367,7 @@ async def flush(self) -> None: # Unblock writes self._flushing = False - async def close(self, flush: bool = False) -> None: + async def close(self, flush: bool = True) -> None: """Closes the provider. Pending writing tasks queued after this will fail, so such tasks should be @@ -424,4 +424,4 @@ async def __aenter__(self) -> Self: return self async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: - await self.close() + await self.close(flush=True) diff --git a/python-packages/smithy-core/tests/unit/aio/test_types.py b/python-packages/smithy-core/tests/unit/aio/test_types.py index 330fa23e4..209e9700d 100644 --- a/python-packages/smithy-core/tests/unit/aio/test_types.py +++ b/python-packages/smithy-core/tests/unit/aio/test_types.py @@ -342,9 +342,9 @@ async def test_close_stops_writes() -> None: await provider.write(b"foo") -async def test_close_deletes_buffered_data() -> None: +async def test_close_without_flush_deletes_buffered_data() -> None: provider = AsyncBytesProvider(b"foo") - await provider.close() + await provider.close(flush=False) result: list[bytes] = [] await drain_provider(provider, result) assert result == [] @@ -400,7 +400,7 @@ async def test_close_stops_queued_writes() -> None: write_task = asyncio.create_task(provider.write(b"bar")) # Now close the provider. The write task will raise an error. - await provider.close() + await provider.close(flush=False) with pytest.raises(SmithyException): await write_task @@ -439,3 +439,29 @@ async def test_close_with_flush() -> None: assert result == [b"foo"] with pytest.raises(SmithyException): await write_task + + +async def test_aexit_flushes() -> None: + # Initialize a provider, keeping track of it in the top scope just to make + # sure it doesn't get GC'd + provider = AsyncBytesProvider() + + # Use the provider in a context manager. When this exits, it should flush + # and close the provider. + async with provider: + + # Write some data to the provider. + await provider.write(b"foo") + + # Start the task to read data from the provider and exit. Explictly do + # not await it here because the exit function should pass priority while + # it waits for the queue to drain. + result: list[bytes] = [] + drain_task = asyncio.create_task(drain_provider(provider, result)) + + # The queue should have been read by this point. + assert result == [b"foo"] + + # The draining task should be able to complete without errors. When next it + # tries to get a chunk, the provider's iterator will exit. + await drain_task