From 2a2e1e7f1d49f7ed372c2a80f08fb67cf1ec953d Mon Sep 17 00:00:00 2001 From: iscai-msft Date: Tue, 7 Sep 2021 17:45:40 -0400 Subject: [PATCH 1/2] add tests --- .../tests/async_tests/test_rest_polling_async.py | 13 +++++++++++++ sdk/core/azure-core/tests/test_rest_polling.py | 12 ++++++++++++ 2 files changed, 25 insertions(+) diff --git a/sdk/core/azure-core/tests/async_tests/test_rest_polling_async.py b/sdk/core/azure-core/tests/async_tests/test_rest_polling_async.py index 4f97b9d1cee4..b80ac0596768 100644 --- a/sdk/core/azure-core/tests/async_tests/test_rest_polling_async.py +++ b/sdk/core/azure-core/tests/async_tests/test_rest_polling_async.py @@ -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} diff --git a/sdk/core/azure-core/tests/test_rest_polling.py b/sdk/core/azure-core/tests/test_rest_polling.py index 3f7caa96f04c..3dfebef958db 100644 --- a/sdk/core/azure-core/tests/test_rest_polling.py +++ b/sdk/core/azure-core/tests/test_rest_polling.py @@ -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} From 8a5eb53560ed5e8b5a848b841c91bde029a74aa0 Mon Sep 17 00:00:00 2001 From: iscai-msft Date: Tue, 7 Sep 2021 17:45:50 -0400 Subject: [PATCH 2/2] correctly pickle rest aiohttp --- sdk/core/azure-core/azure/core/rest/_aiohttp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/core/azure-core/azure/core/rest/_aiohttp.py b/sdk/core/azure-core/azure/core/rest/_aiohttp.py index 44e45174f6eb..055e569b8166 100644 --- a/sdk/core/azure-core/azure/core/rest/_aiohttp.py +++ b/sdk/core/azure-core/azure/core/rest/_aiohttp.py @@ -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