allow LRO pipelines to handle rest responses

Author: iscai-msft

sdk/core/azure-core/azure/core/polling/async_base_polling.py

@@ -116,10 +116,19 @@ async def request_status(self, status_link):  # pylint:disable=invalid-overridde
         """
         if self._path_format_arguments:
             status_link = self._client.format_url(status_link, **self._path_format_arguments)
-        request = self._client.get(status_link)
         # Re-inject 'x-ms-client-request-id' while polling
         if "request_id" not in self._operation_config:
             self._operation_config["request_id"] = self._get_request_id()
+        if hasattr(self._initial_response.http_response, "content"):
+            # if I am a azure.core.rest.HttpResponse
+            # want to keep making azure.core.rest calls
+            from azure.core.rest import HttpRequest as RestHttpRequest
+            request = RestHttpRequest("GET", status_link)
+            return await self._client.send_request(request, _return_pipeline_response=True)
+        # if I am a azure.core.pipeline.transport.HttpResponse
+        request = self._client.get(status_link)
+
+        # can't use send_request in this case, because send_request is still provisional
         return await self._client._pipeline.run(  # pylint: disable=protected-access
             request, stream=False, **self._operation_config
         )

sdk/core/azure-core/azure/core/polling/base_polling.py

@@ -121,7 +121,8 @@ def _is_empty(response):
 
     :rtype: bool
     """
-    return not bool(response.body())
+    content = response.content if hasattr(response, "content") else response.body()  # type: ignore
+    return not bool(content)
 
 
 class LongRunningOperation(ABC):
@@ -574,10 +575,19 @@ def request_status(self, status_link):
         """
         if self._path_format_arguments:
             status_link = self._client.format_url(status_link, **self._path_format_arguments)
-        request = self._client.get(status_link)
         # Re-inject 'x-ms-client-request-id' while polling
         if "request_id" not in self._operation_config:
             self._operation_config["request_id"] = self._get_request_id()
+        if hasattr(self._initial_response.http_response, "content"):
+            # if I am a azure.core.rest.HttpResponse
+            # want to keep making azure.core.rest calls
+            from azure.core.rest import HttpRequest as RestHttpRequest
+            request = RestHttpRequest("GET", status_link)
+            return self._client.send_request(request, _return_pipeline_response=True)
+        # if I am a azure.core.pipeline.transport.HttpResponse
+        request = self._client.get(status_link)
+
+        # can't use send_request in this case, because send_request is still provisional
         return self._client._pipeline.run(  # pylint: disable=protected-access
             request, stream=False, **self._operation_config
         )

sdk/core/azure-core/azure/core/rest/_aiohttp.py

@@ -52,7 +52,7 @@ def __init__(
     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
 

sdk/core/azure-core/tests/async_tests/test_rest_base_polling_async.py

@@ -0,0 +1,115 @@
+#--------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+#--------------------------------------------------------------------------
+import pytest
+from azure.core.exceptions import HttpResponseError
+from azure.core.rest import HttpRequest
+from azure.core.polling import AsyncLROPoller
+from azure.core.polling.async_base_polling import AsyncLROBasePolling
+
+@pytest.fixture
+def deserialization_callback():
+    def _callback(response):
+        return response.http_response.json()
+    return _callback
+
+@pytest.fixture
+def lro_poller(client, deserialization_callback):
+    async def _callback(request):
+        initial_response = await client.send_request(
+            request=request,
+            _return_pipeline_response=True
+        )
+        return AsyncLROPoller(
+            client._client,
+            initial_response,
+            deserialization_callback,
+            AsyncLROBasePolling(0)
+        )
+    return _callback
+
+@pytest.mark.asyncio
+async def test_post_with_location_and_operation_location_headers(lro_poller):
+    poller = await lro_poller(HttpRequest("POST", "/polling/post/location-and-operation-location"))
+    result = await poller.result()
+    assert result == {'location_result': True}
+
+@pytest.mark.asyncio
+async def test_post_with_location_and_operation_location_headers_no_body(lro_poller):
+    poller = await lro_poller(HttpRequest("POST", "/polling/post/location-and-operation-location-no-body"))
+    result = await poller.result()
+    assert result is None
+
+@pytest.mark.asyncio
+async def test_post_resource_location(lro_poller):
+    poller = await lro_poller(HttpRequest("POST", "/polling/post/resource-location"))
+    result = await poller.result()
+    assert result == {'location_result': True}
+
+@pytest.mark.asyncio
+async def test_put_no_polling(lro_poller):
+    result = await (await lro_poller(HttpRequest("PUT", "/polling/no-polling"))).result()
+    assert result['properties']['provisioningState'] == 'Succeeded'
+
+@pytest.mark.asyncio
+async def test_put_location(lro_poller):
+    result = await (await lro_poller(HttpRequest("PUT", "/polling/location"))).result()
+    assert result['location_result']
+
+@pytest.mark.asyncio
+async def test_put_initial_response_body_invalid(lro_poller):
+    # initial body is invalid
+    result = await (await lro_poller(HttpRequest("PUT", "/polling/initial-body-invalid"))).result()
+    assert result['location_result']
+
+@pytest.mark.asyncio
+async def test_put_operation_location_polling_fail(lro_poller):
+    with pytest.raises(HttpResponseError):
+        await (await lro_poller(HttpRequest("PUT", "/polling/put/bad-operation-location"))).result()
+
+@pytest.mark.asyncio
+async def test_put_location_polling_fail(lro_poller):
+    with pytest.raises(HttpResponseError):
+        await (await lro_poller(HttpRequest("PUT", "/polling/bad-location"))).result()
+
+@pytest.mark.asyncio
+async def test_patch_location(lro_poller):
+    result = await (await lro_poller(HttpRequest("PATCH", "/polling/location"))).result()
+    assert result['location_result']
+
+@pytest.mark.asyncio
+async def test_patch_operation_location_polling_fail(lro_poller):
+    with pytest.raises(HttpResponseError):
+        await (await lro_poller(HttpRequest("PUT", "/polling/bad-operation-location"))).result()
+
+@pytest.mark.asyncio
+async def test_patch_location_polling_fail(lro_poller):
+    with pytest.raises(HttpResponseError):
+        await (await lro_poller(HttpRequest("PUT", "/polling/bad-location"))).result()
+
+@pytest.mark.asyncio
+async def test_delete_operation_location(lro_poller):
+    result = await (await lro_poller(HttpRequest("DELETE", "/polling/operation-location"))).result()
+    assert result['status'] == 'Succeeded'

sdk/core/azure-core/tests/test_rest_base_polling.py

@@ -0,0 +1,103 @@
+#--------------------------------------------------------------------------
+#
+# Copyright (c) Microsoft Corporation. All rights reserved.
+#
+# The MIT License (MIT)
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the ""Software""), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+#--------------------------------------------------------------------------
+import pytest
+from azure.core.exceptions import HttpResponseError
+from azure.core.rest import HttpRequest
+from azure.core.polling import LROPoller
+from azure.core.polling.base_polling import LROBasePolling
+
+@pytest.fixture
+def deserialization_callback():
+    def _callback(response):
+        return response.http_response.json()
+    return _callback
+
+@pytest.fixture
+def lro_poller(client, deserialization_callback):
+    def _callback(request):
+        initial_response = client.send_request(
+            request=request,
+            _return_pipeline_response=True
+        )
+        return LROPoller(
+            client._client,
+            initial_response,
+            deserialization_callback,
+            LROBasePolling(0)
+        )
+    return _callback
+
+def test_post_with_location_and_operation_location_headers(lro_poller):
+    poller = lro_poller(HttpRequest("POST", "/polling/post/location-and-operation-location"))
+    result = poller.result()
+    assert result == {'location_result': True}
+
+def test_post_with_location_and_operation_location_headers_no_body(lro_poller):
+    poller = lro_poller(HttpRequest("POST", "/polling/post/location-and-operation-location-no-body"))
+    result = poller.result()
+    assert result is None
+
+def test_post_resource_location(lro_poller):
+    poller = lro_poller(HttpRequest("POST", "/polling/post/resource-location"))
+    result = poller.result()
+    assert result == {'location_result': True}
+
+def test_put_no_polling(lro_poller):
+    result = lro_poller(HttpRequest("PUT", "/polling/no-polling")).result()
+    assert result['properties']['provisioningState'] == 'Succeeded'
+
+def test_put_location(lro_poller):
+    result = lro_poller(HttpRequest("PUT", "/polling/location")).result()
+    assert result['location_result']
+
+def test_put_initial_response_body_invalid(lro_poller):
+    # initial body is invalid
+    result = lro_poller(HttpRequest("PUT", "/polling/initial-body-invalid")).result()
+    assert result['location_result']
+
+def test_put_operation_location_polling_fail(lro_poller):
+    with pytest.raises(HttpResponseError):
+        lro_poller(HttpRequest("PUT", "/polling/put/bad-operation-location")).result()
+
+def test_put_location_polling_fail(lro_poller):
+    with pytest.raises(HttpResponseError):
+        lro_poller(HttpRequest("PUT", "/polling/bad-location")).result()
+
+def test_patch_location(lro_poller):
+    result = lro_poller(HttpRequest("PATCH", "/polling/location")).result()
+    assert result['location_result']
+
+def test_patch_operation_location_polling_fail(lro_poller):
+    with pytest.raises(HttpResponseError):
+        lro_poller(HttpRequest("PUT", "/polling/bad-operation-location")).result()
+
+def test_patch_location_polling_fail(lro_poller):
+    with pytest.raises(HttpResponseError):
+        lro_poller(HttpRequest("PUT", "/polling/bad-location")).result()
+
+def test_delete_operation_location(lro_poller):
+    result = lro_poller(HttpRequest("DELETE", "/polling/operation-location")).result()
+    assert result['status'] == 'Succeeded'

sdk/core/azure-core/tests/testserver_tests/coretestserver/coretestserver/__init__.py

@@ -10,16 +10,18 @@
     basic_api,
     encoding_api,
     errors_api,
+    polling_api,
     streams_api,
     urlencoded_api,
     multipart_api,
-    xml_api
+    xml_api,
 )
 
 app = Flask(__name__)
 app.register_blueprint(basic_api, url_prefix="/basic")
 app.register_blueprint(encoding_api, url_prefix="/encoding")
 app.register_blueprint(errors_api, url_prefix="/errors")
+app.register_blueprint(polling_api, url_prefix="/polling")
 app.register_blueprint(streams_api, url_prefix="/streams")
 app.register_blueprint(urlencoded_api, url_prefix="/urlencoded")
 app.register_blueprint(multipart_api, url_prefix="/multipart")

sdk/core/azure-core/tests/testserver_tests/coretestserver/coretestserver/test_routes/__init__.py

@@ -12,12 +12,14 @@
 from .streams import streams_api
 from .urlencoded import urlencoded_api
 from .xml_route import xml_api
+from .polling import polling_api
 
 __all__ = [
     "basic_api",
     "encoding_api",
     "errors_api",
     "multipart_api",
+    "polling_api",
     "streams_api",
     "urlencoded_api",
     "xml_api",

sdk/core/azure-core/tests/testserver_tests/coretestserver/coretestserver/test_routes/helpers.py

@@ -158,6 +158,9 @@ def get_dict(*keys, **extras):
 
     return out_d
 
+def get_base_url(request):
+    return "http://" + request.host
+
 __all__ = ["assert_with_message",
            "get_dict",
            "jsonify"]