fix lro response setting

ChangeLog.md

@@ -7,6 +7,7 @@ Modelerfour version: 4.15.421
 
 **Bug Fixes**
 
+- Correctly choose schema from response with 200 status code in the case of LRO operations with multiple responses  #814
 - Fix conflict for model deserialization when operation has input param with name `models`  #819
 
 ### 2020-11-09 - 5.4.2

autorest/codegen/models/lro_operation.py

@@ -58,28 +58,30 @@ def __init__(
     def set_lro_response_type(self) -> None:
         if not self.responses:
             return
-        responses = {response.schema: response for response in self.responses if response.has_body}
-        response_types = list(responses.values())
-        response_type = None
-        if len(response_types) > 1:
+        responses_with_bodies = [r for r in self.responses if r.has_body]
+        num_response_schemas = {r.schema for r in responses_with_bodies}
+        response = None
+        if len(num_response_schemas) > 1:
             # choose the response that has a status code of 200
-            response_types_with_200_status_code = [
-                rt for rt in response_types if 200 in rt.status_codes
+            responses_with_200_status_codes = [
+                r for r in responses_with_bodies if 200 in r.status_codes
             ]
-            if not response_types_with_200_status_code:
+            try:
+                response = responses_with_200_status_codes[0]
+                schema_types = {r.schema for r in responses_with_bodies}
+                response_schema = cast(BaseSchema, response.schema).serialization_type
+                _LOGGER.warning(
+                    "Multiple schema types in responses: %s. Choosing: %s", schema_types, response_schema
+                )
+            except IndexError:
                 raise ValueError(
                     f"Your swagger is invalid because you have multiple response schemas for LRO" +
                     f" method {self.python_name} and none of them have a 200 status code."
                 )
-            response_type = response_types_with_200_status_code[0]
 
-            response_type_schema_name = cast(BaseSchema, response_type.schema).serialization_type
-            _LOGGER.warning(
-                "Multiple schema types in responses: %s. Choosing: %s", response_types, response_type_schema_name
-            )
-        elif response_types:
-            response_type = response_types[0]
-        self.lro_response = response_type
+        elif num_response_schemas:
+            response = responses_with_bodies[0]
+        self.lro_response = response
 
     @property
     def has_optional_return_type(self) -> bool:

test/azure/Expected/AcceptanceTests/Lro/lro/aio/operations/_lro_retrys_operations.py

@@ -348,14 +348,10 @@ async def begin_delete_provisioning202_accepted200_succeeded(
         kwargs.pop('content_type', None)
 
         def get_long_running_output(pipeline_response):
-            response_headers = {}
-            response = pipeline_response.http_response
-            response_headers['Location']=self._deserialize('str', response.headers.get('Location'))
-            response_headers['Retry-After']=self._deserialize('int', response.headers.get('Retry-After'))
             deserialized = self._deserialize('Product', pipeline_response)
 
             if cls:
-                return cls(pipeline_response, deserialized, response_headers)
+                return cls(pipeline_response, deserialized, {})
             return deserialized
 
         if polling is True: polling_method = AsyncARMPolling(lro_delay,  **kwargs)

test/azure/Expected/AcceptanceTests/Lro/lro/aio/operations/_lros_operations.py

@@ -2174,14 +2174,10 @@ async def begin_delete_provisioning202_accepted200_succeeded(
         kwargs.pop('content_type', None)
 
         def get_long_running_output(pipeline_response):
-            response_headers = {}
-            response = pipeline_response.http_response
-            response_headers['Location']=self._deserialize('str', response.headers.get('Location'))
-            response_headers['Retry-After']=self._deserialize('int', response.headers.get('Retry-After'))
             deserialized = self._deserialize('Product', pipeline_response)
 
             if cls:
-                return cls(pipeline_response, deserialized, response_headers)
+                return cls(pipeline_response, deserialized, {})
             return deserialized
 
         if polling is True: polling_method = AsyncARMPolling(lro_delay,  **kwargs)
@@ -2278,14 +2274,10 @@ async def begin_delete_provisioning202_deleting_failed200(
         kwargs.pop('content_type', None)
 
         def get_long_running_output(pipeline_response):
-            response_headers = {}
-            response = pipeline_response.http_response
-            response_headers['Location']=self._deserialize('str', response.headers.get('Location'))
-            response_headers['Retry-After']=self._deserialize('int', response.headers.get('Retry-After'))
             deserialized = self._deserialize('Product', pipeline_response)
 
             if cls:
-                return cls(pipeline_response, deserialized, response_headers)
+                return cls(pipeline_response, deserialized, {})
             return deserialized
 
         if polling is True: polling_method = AsyncARMPolling(lro_delay,  **kwargs)
@@ -2382,14 +2374,10 @@ async def begin_delete_provisioning202_deletingcanceled200(
         kwargs.pop('content_type', None)
 
         def get_long_running_output(pipeline_response):
-            response_headers = {}
-            response = pipeline_response.http_response
-            response_headers['Location']=self._deserialize('str', response.headers.get('Location'))
-            response_headers['Retry-After']=self._deserialize('int', response.headers.get('Retry-After'))
             deserialized = self._deserialize('Product', pipeline_response)
 
             if cls:
-                return cls(pipeline_response, deserialized, response_headers)
+                return cls(pipeline_response, deserialized, {})
             return deserialized
 
         if polling is True: polling_method = AsyncARMPolling(lro_delay,  **kwargs)

test/azure/Expected/AcceptanceTests/Lro/lro/operations/_lro_retrys_operations.py

@@ -358,14 +358,10 @@ def begin_delete_provisioning202_accepted200_succeeded(
         kwargs.pop('content_type', None)
 
         def get_long_running_output(pipeline_response):
-            response_headers = {}
-            response = pipeline_response.http_response
-            response_headers['Location']=self._deserialize('str', response.headers.get('Location'))
-            response_headers['Retry-After']=self._deserialize('int', response.headers.get('Retry-After'))
             deserialized = self._deserialize('Product', pipeline_response)
 
             if cls:
-                return cls(pipeline_response, deserialized, response_headers)
+                return cls(pipeline_response, deserialized, {})
             return deserialized
 
         if polling is True: polling_method = ARMPolling(lro_delay,  **kwargs)

test/azure/Expected/AcceptanceTests/Lro/lro/operations/_lros_operations.py

@@ -2218,14 +2218,10 @@ def begin_delete_provisioning202_accepted200_succeeded(
         kwargs.pop('content_type', None)
 
         def get_long_running_output(pipeline_response):
-            response_headers = {}
-            response = pipeline_response.http_response
-            response_headers['Location']=self._deserialize('str', response.headers.get('Location'))
-            response_headers['Retry-After']=self._deserialize('int', response.headers.get('Retry-After'))
             deserialized = self._deserialize('Product', pipeline_response)
 
             if cls:
-                return cls(pipeline_response, deserialized, response_headers)
+                return cls(pipeline_response, deserialized, {})
             return deserialized
 
         if polling is True: polling_method = ARMPolling(lro_delay,  **kwargs)
@@ -2324,14 +2320,10 @@ def begin_delete_provisioning202_deleting_failed200(
         kwargs.pop('content_type', None)
 
         def get_long_running_output(pipeline_response):
-            response_headers = {}
-            response = pipeline_response.http_response
-            response_headers['Location']=self._deserialize('str', response.headers.get('Location'))
-            response_headers['Retry-After']=self._deserialize('int', response.headers.get('Retry-After'))
             deserialized = self._deserialize('Product', pipeline_response)
 
             if cls:
-                return cls(pipeline_response, deserialized, response_headers)
+                return cls(pipeline_response, deserialized, {})
             return deserialized
 
         if polling is True: polling_method = ARMPolling(lro_delay,  **kwargs)
@@ -2430,14 +2422,10 @@ def begin_delete_provisioning202_deletingcanceled200(
         kwargs.pop('content_type', None)
 
         def get_long_running_output(pipeline_response):
-            response_headers = {}
-            response = pipeline_response.http_response
-            response_headers['Location']=self._deserialize('str', response.headers.get('Location'))
-            response_headers['Retry-After']=self._deserialize('int', response.headers.get('Retry-After'))
             deserialized = self._deserialize('Product', pipeline_response)
 
             if cls:
-                return cls(pipeline_response, deserialized, response_headers)
+                return cls(pipeline_response, deserialized, {})
             return deserialized
 
         if polling is True: polling_method = ARMPolling(lro_delay,  **kwargs)

test/azure/Expected/AcceptanceTests/Paging/paging/aio/operations/_paging_operations.py

@@ -232,7 +232,7 @@ async def get_next(next_link=None):
     def first_response_empty(
         self,
         **kwargs
-    ) -> AsyncIterable["models.ProductResultValue"]:
+    ) -> AsyncIterable["_models.ProductResultValue"]:
         """A paging operation whose first response's items list is empty, but still returns a next link.
         Second (and final) call, will give you an items list of 1.
 
@@ -241,7 +241,7 @@ def first_response_empty(
         :rtype: ~azure.core.async_paging.AsyncItemPaged[~paging.models.ProductResultValue]
         :raises: ~azure.core.exceptions.HttpResponseError
         """
-        cls = kwargs.pop('cls', None)  # type: ClsType["models.ProductResultValue"]
+        cls = kwargs.pop('cls', None)  # type: ClsType["_models.ProductResultValue"]
         error_map = {
             401: ClientAuthenticationError, 404: ResourceNotFoundError, 409: ResourceExistsError
         }

test/azure/Expected/AcceptanceTests/Paging/paging/operations/_paging_operations.py

@@ -239,7 +239,7 @@ def first_response_empty(
         self,
         **kwargs  # type: Any
     ):
-        # type: (...) -> Iterable["models.ProductResultValue"]
+        # type: (...) -> Iterable["_models.ProductResultValue"]
         """A paging operation whose first response's items list is empty, but still returns a next link.
         Second (and final) call, will give you an items list of 1.
 
@@ -248,7 +248,7 @@ def first_response_empty(
         :rtype: ~azure.core.paging.ItemPaged[~paging.models.ProductResultValue]
         :raises: ~azure.core.exceptions.HttpResponseError
         """
-        cls = kwargs.pop('cls', None)  # type: ClsType["models.ProductResultValue"]
+        cls = kwargs.pop('cls', None)  # type: ClsType["_models.ProductResultValue"]
         error_map = {
             401: ClientAuthenticationError, 404: ResourceNotFoundError, 409: ResourceExistsError
         }

test/vanilla/Expected/AcceptanceTests/XmsErrorResponse/xmserrorresponse/aio/operations/_pet_operations.py

@@ -153,3 +153,56 @@ async def do_something(
 
         return deserialized
     do_something.metadata = {'url': '/errorStatusCodes/Pets/doSomething/{whatAction}'}  # type: ignore
+
+    @distributed_trace_async
+    async def has_models_param(
+        self,
+        models: Optional[str] = "value1",
+        **kwargs
+    ) -> None:
+        """Ensure you can correctly deserialize the returned PetActionError and deserialization doesn't
+        conflict with the input param name 'models'.
+
+        :param models: Make sure model deserialization doesn't conflict with this param name, which has
+         input name 'models'. Use client default value in call.
+        :type models: str
+        :keyword callable cls: A custom type or function that will be passed the direct response
+        :return: None, or the result of cls(response)
+        :rtype: None
+        :raises: ~azure.core.exceptions.HttpResponseError
+        """
+        cls = kwargs.pop('cls', None)  # type: ClsType[None]
+        error_map = {
+            401: ClientAuthenticationError,
+            404: ResourceNotFoundError,
+            409: ResourceExistsError,
+            500: lambda response: HttpResponseError(response=response, model=self._deserialize(_models.PetActionError, response)),
+        }
+        error_map.update(kwargs.pop('error_map', {}))
+        accept = "application/json"
+
+        # Construct URL
+        url = self.has_models_param.metadata['url']  # type: ignore
+
+        # Construct parameters
+        query_parameters = {}  # type: Dict[str, Any]
+        if models is not None:
+            query_parameters['models'] = self._serialize.query("models", models, 'str')
+
+        # Construct headers
+        header_parameters = {}  # type: Dict[str, Any]
+        header_parameters['Accept'] = self._serialize.header("accept", accept, 'str')
+
+        request = self._client.post(url, query_parameters, header_parameters)
+        pipeline_response = await self._client._pipeline.run(request, stream=False, **kwargs)
+        response = pipeline_response.http_response
+
+        if response.status_code not in [200]:
+            map_error(status_code=response.status_code, response=response, error_map=error_map)
+            error = self._deserialize(_models.PetActionError, response)
+            raise HttpResponseError(response=response, model=error)
+
+        if cls:
+            return cls(pipeline_response, None, {})
+
+    has_models_param.metadata = {'url': '/errorStatusCodes/Pets/hasModelsParam'}  # type: ignore

test/vanilla/Expected/AcceptanceTests/XmsErrorResponse/xmserrorresponse/operations/_pet_operations.py

@@ -159,3 +159,57 @@ def do_something(
 
         return deserialized
     do_something.metadata = {'url': '/errorStatusCodes/Pets/doSomething/{whatAction}'}  # type: ignore
+
+    @distributed_trace
+    def has_models_param(
+        self,
+        models="value1",  # type: Optional[str]
+        **kwargs  # type: Any
+    ):
+        # type: (...) -> None
+        """Ensure you can correctly deserialize the returned PetActionError and deserialization doesn't
+        conflict with the input param name 'models'.
+
+        :param models: Make sure model deserialization doesn't conflict with this param name, which has
+         input name 'models'. Use client default value in call.
+        :type models: str
+        :keyword callable cls: A custom type or function that will be passed the direct response
+        :return: None, or the result of cls(response)
+        :rtype: None
+        :raises: ~azure.core.exceptions.HttpResponseError
+        """
+        cls = kwargs.pop('cls', None)  # type: ClsType[None]
+        error_map = {
+            401: ClientAuthenticationError,
+            404: ResourceNotFoundError,
+            409: ResourceExistsError,
+            500: lambda response: HttpResponseError(response=response, model=self._deserialize(_models.PetActionError, response)),
+        }
+        error_map.update(kwargs.pop('error_map', {}))
+        accept = "application/json"
+
+        # Construct URL
+        url = self.has_models_param.metadata['url']  # type: ignore
+
+        # Construct parameters
+        query_parameters = {}  # type: Dict[str, Any]
+        if models is not None:
+            query_parameters['models'] = self._serialize.query("models", models, 'str')
+
+        # Construct headers
+        header_parameters = {}  # type: Dict[str, Any]
+        header_parameters['Accept'] = self._serialize.header("accept", accept, 'str')
+
+        request = self._client.post(url, query_parameters, header_parameters)
+        pipeline_response = self._client._pipeline.run(request, stream=False, **kwargs)
+        response = pipeline_response.http_response
+
+        if response.status_code not in [200]:
+            map_error(status_code=response.status_code, response=response, error_map=error_map)
+            error = self._deserialize(_models.PetActionError, response)
+            raise HttpResponseError(response=response, model=error)
+
+        if cls:
+            return cls(pipeline_response, None, {})
+
+    has_models_param.metadata = {'url': '/errorStatusCodes/Pets/hasModelsParam'}  # type: ignore