From a4007115c8284d3e660d8bdf76777327de3ea309 Mon Sep 17 00:00:00 2001 From: JordonPhillips Date: Wed, 12 Mar 2025 22:13:31 +0100 Subject: [PATCH] Close crt bridge stream on end_stream if empty This is a followup to #432. I think the issue was that end_stream was getting called when the buffer was already empty, which would put it in an incomplete state. Unfortunately that fix also left it in an incomplete (but now working) state. This make sure it doesn't get into that incomplete state and adds a failing test for that case. --- packages/smithy-http/src/smithy_http/aio/crt.py | 7 ++++++- packages/smithy-http/tests/unit/aio/test_crt.py | 6 ++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/packages/smithy-http/src/smithy_http/aio/crt.py b/packages/smithy-http/src/smithy_http/aio/crt.py index ea72e822d..480b71fba 100644 --- a/packages/smithy-http/src/smithy_http/aio/crt.py +++ b/packages/smithy-http/src/smithy_http/aio/crt.py @@ -369,8 +369,10 @@ def read(self, size: int | None = -1) -> bytes: if len(self._chunks) == 0: if self._done: + self.close() return b"" else: + # When the CRT recieves this, it'll try again raise BlockingIOError("read") # We could compile all the chunks here instead of just returning @@ -429,4 +431,7 @@ def close(self) -> None: def end_stream(self) -> None: """End the stream, letting any remaining chunks be read before it is closed.""" - self._done = True + if len(self._chunks) == 0: + self.close() + else: + self._done = True diff --git a/packages/smithy-http/tests/unit/aio/test_crt.py b/packages/smithy-http/tests/unit/aio/test_crt.py index f0264f1e5..4694b6e3a 100644 --- a/packages/smithy-http/tests/unit/aio/test_crt.py +++ b/packages/smithy-http/tests/unit/aio/test_crt.py @@ -95,6 +95,12 @@ def test_done_stream_read() -> None: assert stream.read() == b"" +def test_end_empty_stream() -> None: + stream = BufferableByteStream() + stream.end_stream() + assert stream.read() == b"" + + def test_stream_read1() -> None: stream = BufferableByteStream() stream.write(b"foo")