async def test_response_streaming_error_behavior():
# Test to reproduce https://github.com/Azure/azure-sdk-for-python/issues/16723
block_size = 103
total_response_size = 500
req_response = requests.Response()
req_request = requests.Request()
class FakeStreamWithConnectionError:
# fake object for urllib3.response.HTTPResponse
def __init__(self):
self.total_response_size = 500
def stream(self, chunk_size, decode_content=False):
assert chunk_size == block_size
left = total_response_size
while left > 0:
if left <= block_size:
raise requests.exceptions.ConnectionError()
data = b"X" * min(chunk_size, left)
left -= len(data)
yield data
def read(self, chunk_size, decode_content=False):
assert chunk_size == block_size
if self.total_response_size > 0:
if self.total_response_size <= block_size:
raise requests.exceptions.ConnectionError()
data = b"X" * min(chunk_size, self.total_response_size)
self.total_response_size -= len(data)
return data
def close(self):
pass
req_response.raw = FakeStreamWithConnectionError()
response = AsyncioRequestsTransportResponse(
req_request,
req_response,
block_size,
)
async def mock_run(self, *args, **kwargs):
return PipelineResponse(
None,
requests.Response(),
None,
)
transport = AioHttpTransport()
pipeline = AsyncPipeline(transport)
pipeline.run = mock_run
downloader = response.stream_download(pipeline)
with pytest.raises(requests.exceptions.ConnectionError):
while True:
await downloader.__anext__()
def mock_send(method, status, headers=None, body=RESPONSE_BODY):
if headers is None:
headers = {}
response = Response()
response._content_consumed = True
response._content = json.dumps(body).encode('ascii') if body is not None else None
response.request = mock.create_autospec(Request)
response.request.method = method
response.request.url = RESOURCE_URL
response.request.headers = {
'x-ms-client-request-id': '67f4dd4e-6262-45e1-8bed-5c45cf23b6d9'
}
response.status_code = status
response.headers = headers
response.headers.update({"content-type": "application/json; charset=utf8"})
response.reason = "OK"
request = CLIENT._request(
response.request.method,
response.request.url,
None, # params
response.request.headers,
body,
None, # form_content
None # stream_content
)
return PipelineResponse(
request,
AsyncioRequestsTransportResponse(
request,
response,
),
None # context
)
def mock_update(url, headers=None):
response = Response()
response._content_consumed = True
response.request = mock.create_autospec(Request)
response.request.method = 'GET'
response.headers = headers or {}
response.headers.update(
{"content-type": "application/json; charset=utf8"})
response.reason = "OK"
if url == ASYNC_URL:
response.request.url = url
response.status_code = POLLING_STATUS
response._content = ASYNC_BODY.encode('ascii')
response.randomFieldFromPollAsyncOpHeader = None
elif url == LOCATION_URL:
response.request.url = url
response.status_code = POLLING_STATUS
response._content = LOCATION_BODY.encode('ascii')
response.randomFieldFromPollLocationHeader = None
elif url == ERROR:
raise BadEndpointError("boom")
elif url == RESOURCE_URL:
response.request.url = url
response.status_code = POLLING_STATUS
response._content = RESOURCE_BODY.encode('ascii')
else:
raise Exception('URL does not match')
request = CLIENT._request(
response.request.method,
response.request.url,
None, # params
{}, # request has no headers
None, # Request has no body
None, # form_content
None # stream_content
)
return PipelineResponse(
request,
AsyncioRequestsTransportResponse(
request,
response,
),
None # context
)
def polling_response():
polling = AsyncLROBasePolling()
headers = {}
response = Response()
response.headers = headers
response.status_code = 200
polling._pipeline_response = PipelineResponse(
None, AsyncioRequestsTransportResponse(
None,
response,
), PipelineContext(None))
polling._initial_response = polling._pipeline_response
return polling, headers