diff --git a/botocore/response.py b/botocore/response.py index 9f7c2eb692..ba3fac9bab 100644 --- a/botocore/response.py +++ b/botocore/response.py @@ -124,6 +124,12 @@ def __next__(self): return current_chunk raise StopIteration() + def __enter__(self): + return self._raw_stream + + def __exit__(self, type, value, traceback): + self._raw_stream.close() + next = __next__ def iter_lines(self, chunk_size=_DEFAULT_CHUNK_SIZE, keepends=False): diff --git a/tests/unit/test_response.py b/tests/unit/test_response.py index 3a22e62475..8aa7a638ac 100644 --- a/tests/unit/test_response.py +++ b/tests/unit/test_response.py @@ -230,6 +230,13 @@ def test_streaming_line_empty_body(self): ) self.assert_lines(stream.iter_lines(), []) + def test_streaming_body_as_context_manager(self): + body = BytesIO(b'1234567890') + with response.StreamingBody(body, content_length=10) as stream: + self.assertEqual(stream.read(), b'1234567890') + self.assertFalse(body.closed) + self.assertTrue(body.closed) + class FakeRawResponse(BytesIO): def stream(self, amt=1024, decode_content=None):