From 9bae78ce63c4c2acef79db3f38d3c02aef211bf4 Mon Sep 17 00:00:00 2001 From: Alex Woods Date: Sun, 9 Mar 2025 16:00:29 -0700 Subject: [PATCH 1/2] Fix AsyncBytesReader, read when bytes are already available --- packages/smithy-core/src/smithy_core/aio/types.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/packages/smithy-core/src/smithy_core/aio/types.py b/packages/smithy-core/src/smithy_core/aio/types.py index 51a63bbc4..9dff70903 100644 --- a/packages/smithy-core/src/smithy_core/aio/types.py +++ b/packages/smithy-core/src/smithy_core/aio/types.py @@ -72,10 +72,11 @@ async def _read_from_iterable( self._remainder = b"" return result - async for element in iterator: - result += element - if len(result) >= size: - break + if len(result) < size: + async for element in iterator: + result += element + if len(result) >= size: + break self._remainder = result[size:] return result[:size] From f61e68370157749db2148b9a74b37a98862e5f8d Mon Sep 17 00:00:00 2001 From: Alex Woods Date: Mon, 10 Mar 2025 08:57:51 -0700 Subject: [PATCH 2/2] Add test --- packages/smithy-core/tests/unit/aio/test_types.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/packages/smithy-core/tests/unit/aio/test_types.py b/packages/smithy-core/tests/unit/aio/test_types.py index b62f3a4b3..ce696bb79 100644 --- a/packages/smithy-core/tests/unit/aio/test_types.py +++ b/packages/smithy-core/tests/unit/aio/test_types.py @@ -3,6 +3,7 @@ import asyncio from io import BytesIO from typing import Self +from unittest.mock import Mock import pytest @@ -96,14 +97,18 @@ async def test_read_async_iterator() -> None: assert await reader.read() == b"foo" assert source.tell() == 3 - source = BytesIO(b"foo,bar,baz\n") + source = Mock(wraps=BytesIO(b"foo,bar,baz\n")) reader = AsyncBytesReader(_AsyncIteratorWrapper(source, chunk_size=6)) assert source.tell() == 0 assert await reader.read(4) == b"foo," assert source.tell() == 6 assert await reader.read(4) == b"bar," assert source.tell() == 12 + + assert source.read.call_count == 2 assert await reader.read(4) == b"baz\n" + # sufficient bytes have been read from source, assert we don't call read again + assert source.read.call_count == 2 assert source.tell() == 12