Flush AsyncBytesProvider on exit

This updates the __aexit__ on AsyncBytesProvider to flush all written
data before closing. This aligns better with the expectations of
using this class as a context manager, as it's rather surprising to
lose data from an automatic close.

Author: JordonPhillips

python-packages/smithy-core/smithy_core/aio/types.py

@@ -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)

python-packages/smithy-core/tests/unit/aio/test_types.py

@@ -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