[core] accept 'final-state-via' lro_options in base polling (Azure#22713)

Author: iscai-msft

sdk/core/azure-core/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- Add support for `final-state-via` LRO option in core.  #22713
+
 ### Breaking Changes
 
 ### Bugs Fixed

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

@@ -26,6 +26,7 @@
 import abc
 import base64
 import json
+from enum import Enum
 from typing import TYPE_CHECKING, Optional, Any, Union
 
 from ..exceptions import HttpResponseError, DecodeError
@@ -174,20 +175,38 @@ def get_final_get_url(self, pipeline_response):
         """
         raise NotImplementedError()
 
+class _LroOption(str, Enum):
+    """Known LRO options from Swagger."""
+
+    FINAL_STATE_VIA = "final-state-via"
+
+
+class _FinalStateViaOption(str, Enum):
+    """Possible final-state-via options."""
+
+    AZURE_ASYNC_OPERATION_FINAL_STATE = "azure-async-operation"
+    LOCATION_FINAL_STATE = "location"
+    OPERATION_LOCATION_FINAL_STATE = "operation-location"
+
 
 class OperationResourcePolling(LongRunningOperation):
     """Implements a operation resource polling, typically from Operation-Location.
 
     :param str operation_location_header: Name of the header to return operation format (default 'operation-location')
+    :keyword dict[str, any] lro_options: Additional options for LRO. For more information, see
+     https://aka.ms/azsdk/autorest/openapi/lro-options
     """
 
-    def __init__(self, operation_location_header="operation-location"):
+    def __init__(
+        self, operation_location_header="operation-location", **kwargs
+    ):
         self._operation_location_header = operation_location_header
 
         # Store the initial URLs
         self._async_url = None
         self._location_url = None
         self._request = None
+        self._lro_options = kwargs.pop("lro_options", {}) or {}
 
     def can_poll(self, pipeline_response):
         """Answer if this polling method could be used.
@@ -207,6 +226,19 @@ def get_final_get_url(self, pipeline_response):
 
         :rtype: str
         """
+        if (
+            self._lro_options.get(_LroOption.FINAL_STATE_VIA) == _FinalStateViaOption.LOCATION_FINAL_STATE
+            and self._location_url
+        ):
+            return self._location_url
+        if (
+            self._lro_options.get(_LroOption.FINAL_STATE_VIA)
+            in [
+                _FinalStateViaOption.AZURE_ASYNC_OPERATION_FINAL_STATE,
+                _FinalStateViaOption.OPERATION_LOCATION_FINAL_STATE
+            ]
+        ):
+            return None
         response = pipeline_response.http_response
         if not _is_empty(response):
             body = _as_json(response)
@@ -381,7 +413,7 @@ def __init__(
         **operation_config
     ):
         self._lro_algorithms = lro_algorithms or [
-            OperationResourcePolling(),
+            OperationResourcePolling(lro_options=lro_options),
             LocationPolling(),
             StatusCheckPolling(),
         ]

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

@@ -42,7 +42,7 @@
 
 from msrest import Deserializer
 
-from azure.core.polling import async_poller
+from azure.core.polling import async_poller, AsyncLROPoller
 from azure.core.exceptions import DecodeError, HttpResponseError
 from azure.core import AsyncPipelineClient
 from azure.core.pipeline import PipelineResponse, AsyncPipeline, PipelineContext
@@ -52,6 +52,7 @@
     AsyncLROBasePolling,
 )
 from utils import ASYNCIO_REQUESTS_TRANSPORT_RESPONSES, request_and_responses_product, create_transport_response
+from rest_client_async import AsyncTestRestClient
 
 class SimpleResource:
     """An implementation of Python 3 SimpleNamespace.
@@ -748,3 +749,125 @@ async def test_long_running_negative(http_request, http_response):
     LOCATION_BODY = json.dumps({ 'name': TEST_NAME })
     POLLING_STATUS = 200
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("http_request,http_response", request_and_responses_product(ASYNCIO_REQUESTS_TRANSPORT_RESPONSES))
+async def test_post_final_state_via(async_pipeline_client_builder, deserialization_cb, http_request, http_response):
+    # Test POST LRO with both Location and Operation-Location
+    CLIENT.http_request_type = http_request
+    CLIENT.http_response_type = http_response
+    # The initial response contains both Location and Operation-Location, a 202 and no Body
+    initial_response = TestBasePolling.mock_send(
+        http_request,
+        http_response,
+        'POST',
+        202,
+        {
+            'location': 'http://example.org/location',
+            'operation-location': 'http://example.org/async_monitor',
+        },
+        ''
+    )
+
+    async def send(request, **kwargs):
+        assert request.method == 'GET'
+
+        if request.url == 'http://example.org/location':
+            return TestBasePolling.mock_send(
+                http_request,
+                http_response,
+                'GET',
+                200,
+                body={'location_result': True}
+            ).http_response
+        elif request.url == 'http://example.org/async_monitor':
+            return TestBasePolling.mock_send(
+                http_request,
+                http_response,
+                'GET',
+                200,
+                body={'status': 'Succeeded'}
+            ).http_response
+        else:
+            pytest.fail("No other query allowed")
+
+    client = async_pipeline_client_builder(send)
+
+    # Test 1, LRO options with Location final state
+    poll = async_poller(
+        client,
+        initial_response,
+        deserialization_cb,
+        AsyncLROBasePolling(0, lro_options={"final-state-via": "location"}))
+    result = await poll
+    assert result['location_result'] == True
+
+    # Test 2, LRO options with Operation-Location final state
+    poll = async_poller(
+        client,
+        initial_response,
+        deserialization_cb,
+        AsyncLROBasePolling(0, lro_options={"final-state-via": "operation-location"}))
+    result = await poll
+    assert result['status'] == 'Succeeded'
+
+    # Test 3, "do the right thing" and use Location by default
+    poll = async_poller(
+        client,
+        initial_response,
+        deserialization_cb,
+        AsyncLROBasePolling(0))
+    result = await poll
+    assert result['location_result'] == True
+
+    # Test 4, location has no body
+
+    async def send(request, **kwargs):
+        assert request.method == 'GET'
+
+        if request.url == 'http://example.org/location':
+            return TestBasePolling.mock_send(
+                http_request,
+                http_response,
+                'GET',
+                200,
+                body=None
+            ).http_response
+        elif request.url == 'http://example.org/async_monitor':
+            return TestBasePolling.mock_send(
+                http_request,
+                http_response,
+                'GET',
+                200,
+                body={'status': 'Succeeded'}
+            ).http_response
+        else:
+            pytest.fail("No other query allowed")
+
+    client = async_pipeline_client_builder(send)
+
+    poll = async_poller(
+        client,
+        initial_response,
+        deserialization_cb,
+        AsyncLROBasePolling(0, lro_options={"final-state-via": "location"}))
+    result = await poll
+    assert result is None
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
+async def test_final_get_via_location(port, http_request, deserialization_cb):
+    client = AsyncTestRestClient(port)
+    request = http_request(
+        "PUT",
+        "http://localhost:{}/polling/polling-with-options".format(port),
+    )
+    request.set_json_body({"hello": "world!"})
+    initial_response = await client._client._pipeline.run(request)
+    poller = AsyncLROPoller(
+        client._client,
+        initial_response,
+        deserialization_cb,
+        AsyncLROBasePolling(0, lro_options={"final-state-via": "location"}),
+    )
+    result = await poller.result()
+    assert result == {"returnedFrom": "locationHeaderUrl"}

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

@@ -31,6 +31,7 @@
 import pickle
 import platform
 import six
+
 try:
     from unittest import mock
 except ImportError:
@@ -48,8 +49,9 @@
 
 from azure.core.polling.base_polling import LROBasePolling
 from azure.core.pipeline.policies._utils import _FixedOffset
-from utils import request_and_responses_product, REQUESTS_TRANSPORT_RESPONSES, create_transport_response
+from utils import request_and_responses_product, REQUESTS_TRANSPORT_RESPONSES, create_transport_response, HTTP_REQUESTS
 from azure.core.pipeline._tools import is_rest
+from rest_client import TestRestClient
 
 class SimpleResource:
     """An implementation of Python 3 SimpleNamespace.
@@ -756,3 +758,124 @@ def test_long_running_negative(self, http_request, http_response):
 
         LOCATION_BODY = json.dumps({ 'name': TEST_NAME })
         POLLING_STATUS = 200
+
+    @pytest.mark.parametrize("http_request,http_response", request_and_responses_product(REQUESTS_TRANSPORT_RESPONSES))
+    def test_post_final_state_via(self, pipeline_client_builder, deserialization_cb, http_request, http_response):
+        # Test POST LRO with both Location and Operation-Location
+        CLIENT.http_request_type = http_request
+        CLIENT.http_response_type = http_response
+        # The initial response contains both Location and Operation-Location, a 202 and no Body
+        initial_response = TestBasePolling.mock_send(
+            http_request,
+            http_response,
+            'POST',
+            202,
+            {
+                'location': 'http://example.org/location',
+                'operation-location': 'http://example.org/async_monitor',
+            },
+            ''
+        )
+
+        def send(request, **kwargs):
+            assert request.method == 'GET'
+
+            if request.url == 'http://example.org/location':
+                return TestBasePolling.mock_send(
+                    http_request,
+                    http_response,
+                    'GET',
+                    200,
+                    body={'location_result': True}
+                ).http_response
+            elif request.url == 'http://example.org/async_monitor':
+                return TestBasePolling.mock_send(
+                    http_request,
+                    http_response,
+                    'GET',
+                    200,
+                    body={'status': 'Succeeded'}
+                ).http_response
+            else:
+                pytest.fail("No other query allowed")
+
+        client = pipeline_client_builder(send)
+
+        # Test 1, LRO options with Location final state
+        poll = LROPoller(
+            client,
+            initial_response,
+            deserialization_cb,
+            LROBasePolling(0, lro_options={"final-state-via": "location"}))
+        result = poll.result()
+        assert result['location_result'] == True
+
+        # Test 2, LRO options with Operation-Location final state
+        poll = LROPoller(
+            client,
+            initial_response,
+            deserialization_cb,
+            LROBasePolling(0, lro_options={"final-state-via": "operation-location"}))
+        result = poll.result()
+        assert result['status'] == 'Succeeded'
+
+        # Test 3, "do the right thing" and use Location by default
+        poll = LROPoller(
+            client,
+            initial_response,
+            deserialization_cb,
+            LROBasePolling(0))
+        result = poll.result()
+        assert result['location_result'] == True
+
+        # Test 4, location has no body
+
+        def send(request, **kwargs):
+            assert request.method == 'GET'
+
+            if request.url == 'http://example.org/location':
+                return TestBasePolling.mock_send(
+                    http_request,
+                    http_response,
+                    'GET',
+                    200,
+                    body=None
+                ).http_response
+            elif request.url == 'http://example.org/async_monitor':
+                return TestBasePolling.mock_send(
+                    http_request,
+                    http_response,
+                    'GET',
+                    200,
+                    body={'status': 'Succeeded'}
+                ).http_response
+            else:
+                pytest.fail("No other query allowed")
+
+        client = pipeline_client_builder(send)
+
+        poll = LROPoller(
+            client,
+            initial_response,
+            deserialization_cb,
+            LROBasePolling(0, lro_options={"final-state-via": "location"}))
+        result = poll.result()
+        assert result is None
+
+@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
+def test_final_get_via_location(port, http_request, deserialization_cb):
+    client = TestRestClient(port)
+    request = http_request(
+        "PUT",
+        "http://localhost:{}/polling/polling-with-options".format(port),
+    )
+    request.set_json_body({"hello": "world!"})
+    initial_response = client._client._pipeline.run(request)
+    poller = LROPoller(
+        client._client,
+        initial_response,
+        deserialization_cb,
+        LROBasePolling(0, lro_options={"final-state-via": "location"}),
+    )
+    result = poller.result()
+    assert result == {"returnedFrom": "locationHeaderUrl"}

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

@@ -151,3 +151,22 @@ def request_id_location():
         '{"status": "Succeeded"}',
         status=200
     )
+
+@polling_api.route('/polling-with-options', methods=["PUT"])
+def polling_with_options_first():
+    base_url = get_base_url(request)
+    return Response(
+        '{"properties":{"provisioningState": "InProgress"}}',
+        headers={
+            'location': '{}/polling/final-get-with-location'.format(base_url),
+            'operation-location': '{}/polling/post/resource-location/operation-location-url'.format(base_url),
+        },
+        status=202
+    )
+
+@polling_api.route('/final-get-with-location', methods=["GET"])
+def polling_with_options_final_get_with_location():
+    return Response(
+        '{"returnedFrom": "locationHeaderUrl"}',
+        status=200
+    )