From 0304a0ca35acb1dc836d423e1b2f4d9bf848a9cf Mon Sep 17 00:00:00 2001 From: iscai-msft Date: Mon, 30 Aug 2021 16:47:41 -0400 Subject: [PATCH] add set_read_checks --- sdk/core/azure-core/azure/core/_pipeline_client_async.py | 2 +- .../azure-core/azure/core/rest/_http_response_impl_async.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/sdk/core/azure-core/azure/core/_pipeline_client_async.py b/sdk/core/azure-core/azure/core/_pipeline_client_async.py index b64e7e47daa4..5523782ee1e4 100644 --- a/sdk/core/azure-core/azure/core/_pipeline_client_async.py +++ b/sdk/core/azure-core/azure/core/_pipeline_client_async.py @@ -209,7 +209,7 @@ async def _make_pipeline_call(self, request, **kwargs): # the body is loaded. instead of doing response.read(), going to set the body # to the internal content rest_response._content = response.body() # pylint: disable=protected-access - await rest_response.read() + await rest_response._set_read_checks() # pylint: disable=protected-access await rest_response.close() except Exception as exc: await rest_response.close() diff --git a/sdk/core/azure-core/azure/core/rest/_http_response_impl_async.py b/sdk/core/azure-core/azure/core/rest/_http_response_impl_async.py index 1f3e85ffec99..8ea5a1d03d25 100644 --- a/sdk/core/azure-core/azure/core/rest/_http_response_impl_async.py +++ b/sdk/core/azure-core/azure/core/rest/_http_response_impl_async.py @@ -46,6 +46,10 @@ class AsyncHttpResponseImpl(_HttpResponseBaseImpl, _AsyncHttpResponse): :keyword Callable stream_download_generator: The stream download generator that we use to stream the response. """ + async def _set_read_checks(self): + self._is_stream_consumed = True + await self.close() + async def read(self) -> bytes: """Read the response's bytes into memory. @@ -57,7 +61,7 @@ async def read(self) -> bytes: async for part in self.iter_bytes(): parts.append(part) self._content = b"".join(parts) - self._is_stream_consumed = True + await self._set_read_checks() return self._content async def iter_raw(self) -> AsyncIterator[bytes]: