Skip to content

Commit f49422f

Browse files
annatischlmazuel
annatisch
authored and
lmazuel
committed
[azure-core] Small fixes for aiohttp (Azure#6490)
* Aiohttp read timeout * Aiohttp response status code * Aiohttp content type header * Review feedback * Fixed typo * Fixed async transport sleep coroutine * Fix for stream generators * Reverted attribute error handling * Fixed trio sleep test * Fixed identity test mocks * Fix identity mock to return content-type as string * Fix universal policy tests
1 parent c3e17d7 commit f49422f

File tree

12 files changed

+66
-59
lines changed

12 files changed

+66
-59
lines changed

‎sdk/core/azure-core/azure/core/exceptions.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@ def __init__(self, message=None, response=None, **kwargs):
114114
if response:
115115
self.reason = response.reason
116116
self.status_code = response.status_code
117+
117118
message = message or "Operation returned an invalid status '{}'".format(self.reason)
118119
try:
119120
try:

‎sdk/core/azure-core/azure/core/pipeline/policies/universal.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -350,7 +350,7 @@ def deserialize_from_http_generics(cls, response):
350350
# Try to use content-type from headers if available
351351
content_type = None
352352
if response.content_type: # type: ignore
353-
content_type = response.content_type[0].strip().lower() # type: ignore
353+
content_type = response.content_type.split(";")[0].strip().lower() # type: ignore
354354

355355
# Ouch, this server did not declare what it sent...
356356
# Let's guess it's JSON...

‎sdk/core/azure-core/azure/core/pipeline/transport/aiohttp.py

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
import aiohttp
3232

3333
from azure.core.configuration import ConnectionConfiguration
34-
from azure.core.exceptions import ServiceRequestError
34+
from azure.core.exceptions import ServiceRequestError, ServiceResponseError
3535
from azure.core.pipeline import Pipeline
3636

3737
from requests.exceptions import (
@@ -181,7 +181,8 @@ async def send(self, request: HttpRequest, **config: Any) -> Optional[AsyncHttpR
181181
await response.load_body()
182182
except aiohttp.client_exceptions.ClientConnectorError as err:
183183
error = ServiceRequestError(err, error=err)
184-
184+
except asyncio.TimeoutError as err:
185+
error = ServiceResponseError(err, error=err)
185186
if error:
186187
raise error
187188
return response
@@ -191,19 +192,16 @@ class AioHttpStreamDownloadGenerator(AsyncIterator):
191192
"""Streams the response body data.
192193
193194
:param pipeline: The pipeline object
194-
:param request: The request object
195195
:param response: The client response object.
196-
:type response: aiohttp.ClientResponse
197196
:param block_size: block size of data sent over connection.
198197
:type block_size: int
199198
"""
200-
def __init__(self, pipeline: Pipeline, request: HttpRequest,
201-
response: aiohttp.ClientResponse, block_size: int) -> None:
199+
def __init__(self, pipeline: Pipeline, response: AsyncHttpResponse) -> None:
202200
self.pipeline = pipeline
203-
self.request = request
201+
self.request = response.request
204202
self.response = response
205-
self.block_size = block_size
206-
self.content_length = int(response.headers.get('Content-Length', 0))
203+
self.block_size = response.block_size
204+
self.content_length = int(response.internal_response.headers.get('Content-Length', 0))
207205
self.downloaded = 0
208206

209207
def __len__(self):
@@ -215,13 +213,13 @@ async def __anext__(self):
215213
retry_interval = 1000
216214
while retry_active:
217215
try:
218-
chunk = await self.response.content.read(self.block_size)
216+
chunk = await self.response.internal_response.content.read(self.block_size)
219217
if not chunk:
220218
raise _ResponseStopIteration()
221219
self.downloaded += self.block_size
222220
return chunk
223221
except _ResponseStopIteration:
224-
self.response.close()
222+
self.response.internal_response.close()
225223
raise StopAsyncIteration()
226224
except (ChunkedEncodingError, ConnectionError):
227225
retry_total -= 1
@@ -233,7 +231,7 @@ async def __anext__(self):
233231
resp = self.pipeline.run(self.request, stream=True, headers=headers)
234232
if resp.status_code == 416:
235233
raise
236-
chunk = await self.response.content.read(self.block_size)
234+
chunk = await self.response.internal_response.content.read(self.block_size)
237235
if not chunk:
238236
raise StopIteration()
239237
self.downloaded += chunk
@@ -243,7 +241,7 @@ async def __anext__(self):
243241
raise
244242
except Exception as err:
245243
_LOGGER.warning("Unable to stream download: %s", err)
246-
self.response.close()
244+
self.response.internal_response.close()
247245
raise
248246

249247
class AioHttpTransportResponse(AsyncHttpResponse):
@@ -282,4 +280,4 @@ def stream_download(self, pipeline) -> AsyncIteratorType[bytes]:
282280
:param pipeline: The pipeline object
283281
:type pipeline: azure.core.pipeline
284282
"""
285-
return AioHttpStreamDownloadGenerator(pipeline, self.request, self.internal_response, self.block_size)
283+
return AioHttpStreamDownloadGenerator(pipeline, self)

‎sdk/core/azure-core/azure/core/pipeline/transport/requests_asyncio.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,9 @@ async def __aenter__(self):
7878
async def __aexit__(self, *exc_details): # pylint: disable=arguments-differ
7979
return super(AsyncioRequestsTransport, self).__exit__()
8080

81+
async def sleep(self, duration):
82+
await asyncio.sleep(duration)
83+
8184
async def send(self, request: HttpRequest, **kwargs: Any) -> AsyncHttpResponse: # type: ignore
8285
"""Send the request using this HTTP sender.
8386
@@ -135,18 +138,16 @@ class AsyncioStreamDownloadGenerator(AsyncIterator):
135138
"""Streams the response body data.
136139
137140
:param pipeline: The pipeline object
138-
:param request: The request object
139141
:param response: The response object.
140-
:param int block_size: block size of data sent over connection.
141142
:param generator iter_content_func: Iterator for response data.
142143
:param int content_length: size of body in bytes.
143144
"""
144-
def __init__(self, pipeline: Pipeline, request: HttpRequest, response: requests.Response, block_size: int) -> None:
145+
def __init__(self, pipeline: Pipeline, response: AsyncHttpResponse) -> None:
145146
self.pipeline = pipeline
146-
self.request = request
147+
self.request = response.request
147148
self.response = response
148-
self.block_size = block_size
149-
self.iter_content_func = self.response.iter_content(self.block_size)
149+
self.block_size = response.block_size
150+
self.iter_content_func = self.response.internal_response.iter_content(self.block_size)
150151
self.content_length = int(response.headers.get('Content-Length', 0))
151152
self.downloaded = 0
152153

@@ -170,7 +171,7 @@ async def __anext__(self):
170171
self.downloaded += self.block_size
171172
return chunk
172173
except _ResponseStopIteration:
173-
self.response.close()
174+
self.response.internal_response.close()
174175
raise StopAsyncIteration()
175176
except (requests.exceptions.ChunkedEncodingError,
176177
requests.exceptions.ConnectionError):
@@ -197,7 +198,7 @@ async def __anext__(self):
197198
raise
198199
except Exception as err:
199200
_LOGGER.warning("Unable to stream download: %s", err)
200-
self.response.close()
201+
self.response.internal_response.close()
201202
raise
202203

203204

@@ -206,5 +207,4 @@ class AsyncioRequestsTransportResponse(AsyncHttpResponse, RequestsTransportRespo
206207
"""
207208
def stream_download(self, pipeline) -> AsyncIteratorType[bytes]: # type: ignore
208209
"""Generator for streaming request body data."""
209-
return AsyncioStreamDownloadGenerator(pipeline, self.request,
210-
self.internal_response, self.block_size) # type: ignore
210+
return AsyncioStreamDownloadGenerator(pipeline, self) # type: ignore

‎sdk/core/azure-core/azure/core/pipeline/transport/requests_basic.py

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -65,9 +65,7 @@ def __init__(self, request, requests_response, block_size=None):
6565
self.status_code = requests_response.status_code
6666
self.headers = requests_response.headers
6767
self.reason = requests_response.reason
68-
content_type = requests_response.headers.get('content-type')
69-
if content_type:
70-
self.content_type = content_type.split(";")
68+
self.content_type = requests_response.headers.get('content-type')
7169

7270
def body(self):
7371
return self.internal_response.content
@@ -82,18 +80,16 @@ class StreamDownloadGenerator(object):
8280
"""Generator for streaming response data.
8381
8482
:param pipeline: The pipeline object
85-
:param request: The request object
8683
:param response: The response object.
87-
:param int block_size: Number of bytes to read into memory.
8884
:param generator iter_content_func: Iterator for response data.
8985
:param int content_length: size of body in bytes.
9086
"""
91-
def __init__(self, pipeline, request, response, block_size):
87+
def __init__(self, pipeline, response):
9288
self.pipeline = pipeline
93-
self.request = request
89+
self.request = response.request
9490
self.response = response
95-
self.block_size = block_size
96-
self.iter_content_func = self.response.iter_content(self.block_size)
91+
self.block_size = response.block_size
92+
self.iter_content_func = self.response.internal_response.iter_content(self.block_size)
9793
self.content_length = int(response.headers.get('Content-Length', 0))
9894
self.downloaded = 0
9995

@@ -115,7 +111,7 @@ def __next__(self):
115111
self.downloaded += self.block_size
116112
return chunk
117113
except StopIteration:
118-
self.response.close()
114+
self.response.internal_response.close()
119115
raise StopIteration()
120116
except (requests.exceptions.ChunkedEncodingError,
121117
requests.exceptions.ConnectionError):
@@ -138,7 +134,7 @@ def __next__(self):
138134
raise
139135
except Exception as err:
140136
_LOGGER.warning("Unable to stream download: %s", err)
141-
self.response.close()
137+
self.response.internal_response.close()
142138
raise
143139
next = __next__ # Python 2 compatibility.
144140

@@ -149,7 +145,7 @@ class RequestsTransportResponse(HttpResponse, _RequestsTransportResponseBase):
149145
def stream_download(self, pipeline):
150146
# type: (PipelineType) -> Iterator[bytes]
151147
"""Generator for streaming request body data."""
152-
return StreamDownloadGenerator(pipeline, self.request, self.internal_response, self.block_size)
148+
return StreamDownloadGenerator(pipeline, self)
153149

154150

155151
class RequestsTransport(HttpTransport):

‎sdk/core/azure-core/azure/core/pipeline/transport/requests_trio.py

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -53,18 +53,16 @@ class TrioStreamDownloadGenerator(AsyncIterator):
5353
"""Generator for streaming response data.
5454
5555
:param pipeline: The pipeline object
56-
:param request: The request object
5756
:param response: The response object.
58-
:param int block_size: Number of bytes to read into memory.
5957
:param generator iter_content_func: Iterator for response data.
6058
:param int content_length: size of body in bytes.
6159
"""
62-
def __init__(self, pipeline: Pipeline, request: HttpRequest, response: requests.Response, block_size: int) -> None:
60+
def __init__(self, pipeline: Pipeline, response: AsyncHttpResponse) -> None:
6361
self.pipeline = pipeline
64-
self.request = request
62+
self.request = response.request
6563
self.response = response
66-
self.block_size = block_size
67-
self.iter_content_func = self.response.iter_content(self.block_size)
64+
self.block_size = response.block_size
65+
self.iter_content_func = self.response.internal_response.iter_content(self.block_size)
6866
self.content_length = int(response.headers.get('Content-Length', 0))
6967
self.downloaded = 0
7068

@@ -85,7 +83,7 @@ async def __anext__(self):
8583
self.downloaded += self.block_size
8684
return chunk
8785
except _ResponseStopIteration:
88-
self.response.close()
86+
self.response.internal_response.close()
8987
raise StopAsyncIteration()
9088
except (requests.exceptions.ChunkedEncodingError,
9189
requests.exceptions.ConnectionError):
@@ -111,7 +109,7 @@ async def __anext__(self):
111109
raise
112110
except Exception as err:
113111
_LOGGER.warning("Unable to stream download: %s", err)
114-
self.response.close()
112+
self.response.internal_response.close()
115113
raise
116114

117115
class TrioRequestsTransportResponse(AsyncHttpResponse, RequestsTransportResponse): # type: ignore
@@ -120,8 +118,7 @@ class TrioRequestsTransportResponse(AsyncHttpResponse, RequestsTransportResponse
120118
def stream_download(self, pipeline) -> AsyncIteratorType[bytes]: # type: ignore
121119
"""Generator for streaming response data.
122120
"""
123-
return TrioStreamDownloadGenerator(pipeline, self.request,
124-
self.internal_response, self.block_size) # type: ignore
121+
return TrioStreamDownloadGenerator(pipeline, self) # type: ignore
125122

126123

127124
class TrioRequestsTransport(RequestsTransport, AsyncHttpTransport): # type: ignore

‎sdk/core/azure-core/tests/azure_core_asynctests/test_pipeline.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,23 @@ async def test_basic_async_requests():
119119

120120
assert response.http_response.status_code == 200
121121

122+
@pytest.mark.asyncio
123+
async def test_async_transport_sleep():
124+
125+
async with AsyncioRequestsTransport() as transport:
126+
await transport.sleep(1)
127+
128+
async with AioHttpTransport() as transport:
129+
await transport.sleep(1)
130+
131+
def test_async_trio_transport_sleep():
132+
133+
async def do():
134+
async with TrioRequestsTransport() as transport:
135+
await transport.sleep(1)
136+
137+
response = trio.run(do)
138+
122139
@pytest.mark.asyncio
123140
async def test_conf_async_requests():
124141

‎sdk/core/azure-core/tests/test_universal_pipeline.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -123,9 +123,7 @@ class MockResponse(HttpResponse):
123123
def __init__(self, body, content_type):
124124
super(MockResponse, self).__init__(None, None)
125125
self._body = body
126-
self.content_type = None
127-
if content_type:
128-
self.content_type = [content_type]
126+
self.content_type = content_type
129127

130128
def body(self):
131129
return self._body

‎sdk/identity/azure-identity/tests/helpers.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ def mock_response(status_code=200, headers={}, json_payload=None):
4141
if json_payload is not None:
4242
response.text = lambda: json.dumps(json_payload)
4343
response.headers["content-type"] = "application/json"
44-
response.content_type = ["application/json"]
44+
response.content_type = "application/json"
4545
return response
4646

4747

‎sdk/identity/azure-identity/tests/test_authn_client.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def test_authn_client_deserialization():
2828
scope = "scope"
2929

3030
mock_response = Mock(
31-
headers={"content-type": "application/json"}, status_code=200, content_type=["application/json"]
31+
headers={"content-type": "application/json"}, status_code=200, content_type="application/json"
3232
)
3333
mock_send = Mock(return_value=mock_response)
3434

@@ -87,7 +87,7 @@ def test_caching_when_only_expires_in_set():
8787
text=lambda: json.dumps({"access_token": access_token, "expires_in": expires_in, "token_type": "Bearer"}),
8888
headers={"content-type": "application/json"},
8989
status_code=200,
90-
content_type=["application/json"],
90+
content_type="application/json",
9191
)
9292
mock_send = Mock(return_value=mock_response)
9393

@@ -106,7 +106,7 @@ def test_expires_in_strings():
106106
expected_token = "token"
107107

108108
mock_response = Mock(
109-
headers={"content-type": "application/json"}, status_code=200, content_type=["application/json"]
109+
headers={"content-type": "application/json"}, status_code=200, content_type="application/json"
110110
)
111111
mock_send = Mock(return_value=mock_response)
112112

@@ -133,7 +133,7 @@ def test_cache_expiry():
133133
text=lambda: json.dumps(token_payload),
134134
headers={"content-type": "application/json"},
135135
status_code=200,
136-
content_type=["application/json"],
136+
content_type="application/json",
137137
)
138138
mock_send = Mock(return_value=mock_response)
139139

‎sdk/identity/azure-identity/tests/test_identity.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ def test_imds_credential_cache():
191191
text=lambda: json.dumps(token_payload),
192192
headers={"content-type": "application/json"},
193193
status_code=200,
194-
content_type=["application/json"],
194+
content_type="application/json",
195195
)
196196
mock_send = Mock(return_value=mock_response)
197197

@@ -219,7 +219,7 @@ def test_imds_credential_retries():
219219
mock_response = Mock(
220220
text=lambda: b"{}",
221221
headers={"content-type": "application/json", "Retry-After": "0"},
222-
content_type=["application/json"],
222+
content_type="application/json",
223223
)
224224
mock_send = Mock(return_value=mock_response)
225225

‎sdk/identity/azure-identity/tests/test_identity_async.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,7 @@ async def test_imds_credential_cache():
196196
text=lambda: json.dumps(token_payload),
197197
headers={"content-type": "application/json"},
198198
status_code=200,
199-
content_type=["application/json"],
199+
content_type="application/json",
200200
)
201201
mock_send = Mock(return_value=mock_response)
202202

@@ -225,7 +225,7 @@ async def test_imds_credential_retries():
225225
mock_response = Mock(
226226
text=lambda: b"{}",
227227
headers={"content-type": "application/json", "Retry-After": "0"},
228-
content_type=["application/json"],
228+
content_type="application/json",
229229
)
230230
mock_send = Mock(return_value=mock_response)
231231

0 commit comments

Comments
 (0)