Skip to content

Commit

Permalink
[rest] correctly pickle rest aiohttp responses (#20577)
Browse files Browse the repository at this point in the history
  • Loading branch information
iscai-msft authored Sep 8, 2021
1 parent 84bd0e9 commit 6ccb4ad
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 1 deletion.
2 changes: 1 addition & 1 deletion sdk/core/azure-core/azure/core/rest/_aiohttp.py
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ async def iter_bytes(self) -> AsyncIterator[bytes]:
def __getstate__(self):
state = self.__dict__.copy()
# Remove the unpicklable entries.
state['internal_response'] = None # aiohttp response are not pickable (see headers comments)
state['_internal_response'] = None # aiohttp response are not pickable (see headers comments)
state['headers'] = CIMultiDict(self.headers) # MultiDictProxy is not pickable
return state

Expand Down
13 changes: 13 additions & 0 deletions sdk/core/azure-core/tests/async_tests/test_rest_polling_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,3 +118,16 @@ async def test_delete_operation_location(lro_poller):
async def test_request_id(lro_poller):
result = await (await lro_poller(HttpRequest("POST", "/polling/request-id"), request_id="123456789")).result()
assert result['status'] == "Succeeded"

@pytest.mark.asyncio
async def test_continuation_token(client, lro_poller, deserialization_callback):
poller = await lro_poller(HttpRequest("POST", "/polling/post/location-and-operation-location"))
token = poller.continuation_token()
new_poller = AsyncLROPoller.from_continuation_token(
continuation_token=token,
polling_method=AsyncLROBasePolling(0),
client=client._client,
deserialization_callback=deserialization_callback,
)
result = await new_poller.result()
assert result == {'location_result': True}
12 changes: 12 additions & 0 deletions sdk/core/azure-core/tests/test_rest_polling.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,3 +104,15 @@ def test_delete_operation_location(lro_poller):

def test_request_id(lro_poller):
result = lro_poller(HttpRequest("POST", "/polling/request-id"), request_id="123456789").result()

def test_continuation_token(client, lro_poller, deserialization_callback):
poller = lro_poller(HttpRequest("POST", "/polling/post/location-and-operation-location"))
token = poller.continuation_token()
new_poller = LROPoller.from_continuation_token(
continuation_token=token,
polling_method=LROBasePolling(0),
client=client._client,
deserialization_callback=deserialization_callback,
)
result = new_poller.result()
assert result == {'location_result': True}

0 comments on commit 6ccb4ad

Please sign in to comment.