From b1fe424c0509b2f95f30453d03bc48af36ecf0b5 Mon Sep 17 00:00:00 2001 From: Swarup Chanda Date: Thu, 16 Jul 2026 23:48:31 +0530 Subject: [PATCH 1/5] [cognitiveservices] Fix compute create: accept 202, poll resource, tolerate read-after-write 404 Override begin_create_or_update (sync + async) to accept the service's HTTP 202 and track the LRO by polling the compute resource (BodyContentPolling) instead of the computeOperations status endpoint that needs computeOperations/read many callers lack. Tolerates the ~15-20s read-after-write 404 window, blocks until terminal, and surfaces the compute's own provisioning error. Adds 14 unit tests; bumps to 15.0.0b4. --- .../azure-mgmt-cognitiveservices/CHANGELOG.md | 21 ++ .../azure/mgmt/cognitiveservices/_version.py | 2 +- .../aio/operations/_patch.py | 210 +++++++++++++++- .../cognitiveservices/operations/_patch.py | 238 +++++++++++++++++- .../tests/test_computes_operations_patch.py | 210 ++++++++++++++++ .../test_computes_operations_patch_async.py | 210 ++++++++++++++++ 6 files changed, 888 insertions(+), 3 deletions(-) create mode 100644 sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py create mode 100644 sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/CHANGELOG.md b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/CHANGELOG.md index 119efa07cee0..f18649325342 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/CHANGELOG.md +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/CHANGELOG.md @@ -1,5 +1,26 @@ # Release History +## 15.0.0b4 (Unreleased) + +### Bugs Fixed + + - `ComputesOperations.begin_create_or_update` now works for asynchronous compute creation that returns + HTTP 202 (Accepted) and no longer polls the compute operation-status endpoint. Previously the create + failed in two ways: the generated code rejected the 202 response with + `Operation returned an invalid status 'Accepted'`, and (when it did poll) the operation-status endpoint + (`.../locations/{location}/computeOperations/{operationId}`) required the + `Microsoft.CognitiveServices/locations/computeOperations/read` permission, so callers allowed to create a + compute but not granted that read permission were shown a misleading `AuthorizationFailed` error even + though the create succeeded. The 202 is now accepted and the long-running operation is tracked by polling + the compute resource itself (a `GET` on the resource URL, watching `provisioningState`) instead of the + operation-status endpoint, so it only needs the `computes/read` permission that callers already have. The + poller still blocks until the operation reaches a terminal state and raises on a genuine provisioning + failure; non-2xx create failures still propagate. Callers can opt out of blocking with `polling=False`. + The poller also tolerates the short read-after-write window right after the create where the compute is + briefly not yet queryable (a `GET` returns `404 "Cluster not found"`), so a successful create is not + wrongly reported as failed, and when provisioning does fail it surfaces the compute's own error detail + (e.g. a quota message) instead of a generic `Operation returned an invalid status 'OK'`. + ## 15.0.0b3 (2026-06-26) ### Features Added diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/_version.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/_version.py index 8cc664488da0..6fce30ba3921 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/_version.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/_version.py @@ -6,4 +6,4 @@ # Changes may cause incorrect behavior and will be lost if the code is regenerated. # -------------------------------------------------------------------------- -VERSION = "15.0.0b3" +VERSION = "15.0.0b4" diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py index 87676c65a8f0..d73ae4060aa3 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py @@ -7,9 +7,217 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ +import json +import time +from collections.abc import MutableMapping +from io import IOBase +from typing import IO, Any, Optional, Union, cast +from azure.core.exceptions import ( + ClientAuthenticationError, + HttpResponseError, + ResourceExistsError, + ResourceNotFoundError, + ResourceNotModifiedError, + map_error, +) +from azure.core.polling import AsyncLROPoller, AsyncNoPolling, AsyncPollingMethod +from azure.core.polling.base_polling import _raise_if_bad_http_status_and_method +from azure.core.tracing.decorator_async import distributed_trace_async +from azure.core.utils import case_insensitive_dict +from azure.mgmt.core.exceptions import ARMErrorFormat +from azure.mgmt.core.polling.arm_polling import BodyContentPolling, StatusCheckPolling +from azure.mgmt.core.polling.async_arm_polling import AsyncARMPolling -__all__: list[str] = [] # Add all objects you want publicly available to users at this package level +from ... import models as _models +from ..._utils.model_base import SdkJSONEncoder, _deserialize, _failsafe_deserialize +from ...operations._operations import build_computes_create_or_update_request +from ...operations._patch import _TERMINAL_FAILED_STATES, _compute_provisioning_error +from ._operations import ComputesOperations as _ComputesOperationsGenerated + +JSON = MutableMapping[str, Any] + + +class _AsyncComputeResourcePolling(AsyncARMPolling): + """``AsyncARMPolling`` that tolerates the compute-create read-after-write window. + + Right after the service accepts the create (HTTP 202), the new compute is briefly not queryable: a + ``GET`` on the resource returns ``404 "Cluster not found"`` for ~15-20s. Plain resource polling + (:class:`BodyContentPolling`) would treat that 404 as a failure and wrongly fail a create that + actually succeeded. This subclass treats a 404 as "still in progress" and keeps polling until the + resource is queryable and reaches a terminal ``provisioningState``. A generous grace period bounds + the tolerance so a genuinely missing resource still surfaces instead of hanging forever. When the + compute provisions to a failed state it raises the resource's own error detail rather than the + generic ``Operation returned an invalid status 'OK'``. + """ + + _NOT_FOUND_GRACE_SECONDS = 300 + + async def update_status(self) -> None: + self._pipeline_response = await self.request_status(self._operation.get_polling_url()) + response = self._pipeline_response.http_response + if response.status_code == 404: + first_seen = getattr(self, "_not_found_since", None) + if first_seen is None: + self._not_found_since = time.monotonic() + elif time.monotonic() - first_seen > self._NOT_FOUND_GRACE_SECONDS: + _raise_if_bad_http_status_and_method(response) # grace exhausted: surface the 404 + self._status = "InProgress" + return + self._not_found_since = None + _raise_if_bad_http_status_and_method(response) + self._status = self._operation.get_status(self._pipeline_response) + if str(self._status).lower() in _TERMINAL_FAILED_STATES: + raise _compute_provisioning_error(response) + + +class ComputesOperations(_ComputesOperationsGenerated): + """Customized ``ComputesOperations`` that polls the compute resource instead of its + operation-status endpoint. + + The generated :meth:`begin_create_or_update` has two problems against the current service: + + * it rejects the ``202 Accepted`` the service returns for the async create with + "Operation returned an invalid status 'Accepted'"; and + * its poller follows the ``Azure-AsyncOperation`` header to + ``.../locations/{location}/computeOperations/{operationId}``, which requires the + ``Microsoft.CognitiveServices/locations/computeOperations/read`` permission. Many callers who + are allowed to create a compute lack that permission, so the poll fails with + ``AuthorizationFailed`` even though the create itself succeeded. + + This override accepts the ``202`` and polls the compute resource itself (a ``GET`` on the resource + URL, watching ``provisioningState``) via :class:`BodyContentPolling`, deliberately excluding the + algorithms that would follow the ``Azure-AsyncOperation``/``Location`` headers. As a result the + poller only needs ``computes/read`` (which callers already have), still blocks until the operation + reaches a terminal state, and surfaces a genuine provisioning failure instead of masking it. + """ + + @distributed_trace_async + async def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: Union[_models.Compute, JSON, IO[bytes]], + **kwargs: Any, + ) -> AsyncLROPoller[_models.Compute]: + """Creates or updates a compute associated with the Cognitive Services account. + + This override accepts the service's ``202 Accepted`` and polls the compute resource itself + (a ``GET`` on the resource URL, watching ``provisioningState``) instead of the operation-status + endpoint that requires ``computeOperations/read``. It still blocks until the operation reaches + a terminal state and raises on a genuine provisioning failure, so callers get correct results + without needing the operation-status read permission. + + :param resource_group_name: The name of the resource group. The name is case insensitive. + Required. + :type resource_group_name: str + :param account_name: The name of Cognitive Services account. Required. + :type account_name: str + :param compute_name: The name of the compute associated with the Cognitive Services Account. + Required. + :type compute_name: str + :param resource: The compute properties. Is one of the following types: Compute, JSON, + IO[bytes] Required. + :type resource: ~azure.mgmt.cognitiveservices.models.Compute or JSON or IO[bytes] + :return: An instance of AsyncLROPoller that returns Compute. The Compute is compatible with + MutableMapping + :rtype: ~azure.core.polling.AsyncLROPoller[~azure.mgmt.cognitiveservices.models.Compute] + :raises ~azure.core.exceptions.HttpResponseError: + """ + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + _params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = ( + kwargs.pop("content_type", _headers.pop("Content-Type", None)) or "application/json" + ) + cls = kwargs.pop("cls", None) + polling: Union[bool, AsyncPollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + kwargs.pop("continuation_token", None) + + error_map: dict = { + 401: ClientAuthenticationError, + 404: ResourceNotFoundError, + 409: ResourceExistsError, + 304: ResourceNotModifiedError, + } + error_map.update(kwargs.pop("error_map", {}) or {}) + + if isinstance(resource, (IOBase, bytes)): + _content = resource + else: + _content = json.dumps(resource, cls=SdkJSONEncoder, exclude_readonly=True) # type: ignore + + _request = build_computes_create_or_update_request( + resource_group_name=resource_group_name, + account_name=account_name, + compute_name=compute_name, + subscription_id=self._config.subscription_id, + content_type=content_type, + api_version=self._config.api_version, + content=_content, + headers=_headers, + params=_params, + ) + _request.url = self._client.format_url( + _request.url, + endpoint=self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), + ) + + pipeline_response = await self._client._pipeline.run( # pylint: disable=protected-access + _request, stream=False, **kwargs + ) + response = pipeline_response.http_response + # Accept 202 (async "Accepted") in addition to 200/201. The generated create rejects 202 with + # "Operation returned an invalid status 'Accepted'", failing a create the service accepted. + if response.status_code not in [200, 201, 202]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + error = _failsafe_deserialize(_models.ErrorResponse, response) + raise HttpResponseError(response=response, model=error, error_format=ARMErrorFormat) + await response.read() # load the body so the poller/deserializer can read it + + def get_long_running_output(pipeline_response): + response = pipeline_response.http_response + deserialized = _deserialize(_models.Compute, response.json()) + if cls: + return cls(pipeline_response, deserialized, {}) + return deserialized + + path_format_arguments = { + "endpoint": self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), + } + + # Poll the compute resource itself (a GET on the resource URL, watching ``provisioningState``) + # instead of the operation-status endpoint. ``BodyContentPolling`` polls the resource URL, so + # the poller never calls ``.../computeOperations/{id}`` (which requires ``computeOperations/read`` + # that many callers lack). Excluding ``AzureAsyncOperationPolling`` and ``LocationPolling`` + # guarantees the ``Azure-AsyncOperation``/``Location`` headers the service returns are ignored, + # while the poller still blocks until terminal and raises on a failed provisioning. + # ``_AsyncComputeResourcePolling`` additionally tolerates the ~20s read-after-write window where + # the resource GET returns 404 "Cluster not found" right after the create is accepted. + if polling is True: + polling_method: AsyncPollingMethod = cast( + AsyncPollingMethod, + _AsyncComputeResourcePolling( + lro_delay, + lro_algorithms=[BodyContentPolling(), StatusCheckPolling()], + path_format_arguments=path_format_arguments, + **kwargs, + ), + ) + elif polling is False: + polling_method = cast(AsyncPollingMethod, AsyncNoPolling()) + else: + polling_method = polling + return AsyncLROPoller[_models.Compute]( + self._client, pipeline_response, get_long_running_output, polling_method # type: ignore + ) + + +__all__: list[str] = [ + "ComputesOperations", +] # Add all objects you want publicly available to users at this package level def patch_sdk(): diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py index 87676c65a8f0..acc4a2b830e0 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py @@ -7,9 +7,245 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ +import json +import time +from collections.abc import MutableMapping +from io import IOBase +from typing import IO, Any, Optional, Union, cast +from azure.core.exceptions import ( + ClientAuthenticationError, + HttpResponseError, + ResourceExistsError, + ResourceNotFoundError, + ResourceNotModifiedError, + map_error, +) +from azure.core.polling import LROPoller, NoPolling, PollingMethod +from azure.core.polling.base_polling import _raise_if_bad_http_status_and_method +from azure.core.tracing.decorator import distributed_trace +from azure.core.utils import case_insensitive_dict +from azure.mgmt.core.exceptions import ARMErrorFormat +from azure.mgmt.core.polling.arm_polling import ARMPolling, BodyContentPolling, StatusCheckPolling -__all__: list[str] = [] # Add all objects you want publicly available to users at this package level +from .. import models as _models +from .._utils.model_base import SdkJSONEncoder, _deserialize, _failsafe_deserialize +from ._operations import ComputesOperations as _ComputesOperationsGenerated +from ._operations import build_computes_create_or_update_request + +JSON = MutableMapping[str, Any] + +_TERMINAL_FAILED_STATES = frozenset({"failed", "canceled", "cancelled"}) + + +def _compute_provisioning_error(response: Any) -> HttpResponseError: + """Build a descriptive error from a compute resource whose provisioning failed. + + Because the poller watches the compute resource (not the operation-status endpoint), the failure + detail lives in the compute's ``properties.errors``. Surface that instead of azure-core's generic + ``Operation returned an invalid status 'OK'`` (which happens because the terminal poll is a plain + HTTP 200 with no OData error body). + """ + state = None + detail = "" + try: + body = response.json() + props = body.get("properties", {}) if isinstance(body, dict) else {} + state = props.get("provisioningState") + parts = [] + for err in props.get("errors") or []: + if isinstance(err, dict): + code, message = err.get("code"), err.get("message") + parts.append(f"{code}: {message}" if code and message else str(message or code)) + detail = "; ".join(p for p in parts if p) + except Exception: # pylint: disable=broad-except + pass + text = f"Compute provisioning {state or 'failed'}." + if detail: + text = f"{text} {detail}" + return HttpResponseError(response=response, message=text) + + +class _ComputeResourcePolling(ARMPolling): + """``ARMPolling`` that tolerates the compute-create read-after-write window. + + Right after the service accepts the create (HTTP 202), the new compute is briefly not queryable: a + ``GET`` on the resource returns ``404 "Cluster not found"`` for ~15-20s. Plain resource polling + (:class:`BodyContentPolling`) would treat that 404 as a failure and wrongly fail a create that + actually succeeded. This subclass treats a 404 as "still in progress" and keeps polling until the + resource is queryable and reaches a terminal ``provisioningState``. A generous grace period bounds + the tolerance so a genuinely missing resource still surfaces instead of hanging forever. When the + compute provisions to a failed state it raises the resource's own error detail rather than the + generic ``Operation returned an invalid status 'OK'``. + """ + + _NOT_FOUND_GRACE_SECONDS = 300 + + def update_status(self) -> None: + self._pipeline_response = self.request_status(self._operation.get_polling_url()) + response = self._pipeline_response.http_response + if response.status_code == 404: + first_seen = getattr(self, "_not_found_since", None) + if first_seen is None: + self._not_found_since = time.monotonic() + elif time.monotonic() - first_seen > self._NOT_FOUND_GRACE_SECONDS: + _raise_if_bad_http_status_and_method(response) # grace exhausted: surface the 404 + self._status = "InProgress" + return + self._not_found_since = None + _raise_if_bad_http_status_and_method(response) + self._status = self._operation.get_status(self._pipeline_response) + if str(self._status).lower() in _TERMINAL_FAILED_STATES: + raise _compute_provisioning_error(response) + + +class ComputesOperations(_ComputesOperationsGenerated): + """Customized ``ComputesOperations`` that polls the compute resource instead of its + operation-status endpoint. + + The generated :meth:`begin_create_or_update` has two problems against the current service: + + * it rejects the ``202 Accepted`` the service returns for the async create with + "Operation returned an invalid status 'Accepted'"; and + * its poller follows the ``Azure-AsyncOperation`` header to + ``.../locations/{location}/computeOperations/{operationId}``, which requires the + ``Microsoft.CognitiveServices/locations/computeOperations/read`` permission. Many callers who + are allowed to create a compute lack that permission, so the poll fails with + ``AuthorizationFailed`` even though the create itself succeeded. + + This override accepts the ``202`` and polls the compute resource itself (a ``GET`` on the resource + URL, watching ``provisioningState``) via :class:`BodyContentPolling`, deliberately excluding the + algorithms that would follow the ``Azure-AsyncOperation``/``Location`` headers. As a result the + poller only needs ``computes/read`` (which callers already have), still blocks until the operation + reaches a terminal state, and surfaces a genuine provisioning failure instead of masking it. + """ + + @distributed_trace + def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: Union[_models.Compute, JSON, IO[bytes]], + **kwargs: Any, + ) -> LROPoller[_models.Compute]: + """Creates or updates a compute associated with the Cognitive Services account. + + This override accepts the service's ``202 Accepted`` and polls the compute resource itself + (a ``GET`` on the resource URL, watching ``provisioningState``) instead of the operation-status + endpoint that requires ``computeOperations/read``. It still blocks until the operation reaches + a terminal state and raises on a genuine provisioning failure, so callers get correct results + without needing the operation-status read permission. + + :param resource_group_name: The name of the resource group. The name is case insensitive. + Required. + :type resource_group_name: str + :param account_name: The name of Cognitive Services account. Required. + :type account_name: str + :param compute_name: The name of the compute associated with the Cognitive Services Account. + Required. + :type compute_name: str + :param resource: The compute properties. Is one of the following types: Compute, JSON, + IO[bytes] Required. + :type resource: ~azure.mgmt.cognitiveservices.models.Compute or JSON or IO[bytes] + :return: An instance of LROPoller that returns Compute. The Compute is compatible with + MutableMapping + :rtype: ~azure.core.polling.LROPoller[~azure.mgmt.cognitiveservices.models.Compute] + :raises ~azure.core.exceptions.HttpResponseError: + """ + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + _params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = ( + kwargs.pop("content_type", _headers.pop("Content-Type", None)) or "application/json" + ) + cls = kwargs.pop("cls", None) + polling: Union[bool, PollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + kwargs.pop("continuation_token", None) + + error_map: dict = { + 401: ClientAuthenticationError, + 404: ResourceNotFoundError, + 409: ResourceExistsError, + 304: ResourceNotModifiedError, + } + error_map.update(kwargs.pop("error_map", {}) or {}) + + if isinstance(resource, (IOBase, bytes)): + _content = resource + else: + _content = json.dumps(resource, cls=SdkJSONEncoder, exclude_readonly=True) # type: ignore + + _request = build_computes_create_or_update_request( + resource_group_name=resource_group_name, + account_name=account_name, + compute_name=compute_name, + subscription_id=self._config.subscription_id, + content_type=content_type, + api_version=self._config.api_version, + content=_content, + headers=_headers, + params=_params, + ) + _request.url = self._client.format_url( + _request.url, + endpoint=self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), + ) + + pipeline_response = self._client._pipeline.run( + _request, stream=False, **kwargs + ) # pylint: disable=protected-access + response = pipeline_response.http_response + # Accept 202 (async "Accepted") in addition to 200/201. The generated create rejects 202 with + # "Operation returned an invalid status 'Accepted'", failing a create the service accepted. + if response.status_code not in [200, 201, 202]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + error = _failsafe_deserialize(_models.ErrorResponse, response) + raise HttpResponseError(response=response, model=error, error_format=ARMErrorFormat) + response.read() # load the body so the poller/deserializer can read it + + def get_long_running_output(pipeline_response): + response = pipeline_response.http_response + deserialized = _deserialize(_models.Compute, response.json()) + if cls: + return cls(pipeline_response, deserialized, {}) + return deserialized + + path_format_arguments = { + "endpoint": self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), + } + + # Poll the compute resource itself (a GET on the resource URL, watching ``provisioningState``) + # instead of the operation-status endpoint. ``BodyContentPolling`` polls the resource URL, so + # the poller never calls ``.../computeOperations/{id}`` (which requires ``computeOperations/read`` + # that many callers lack). Excluding ``AzureAsyncOperationPolling`` and ``LocationPolling`` + # guarantees the ``Azure-AsyncOperation``/``Location`` headers the service returns are ignored, + # while the poller still blocks until terminal and raises on a failed provisioning. + # ``_ComputeResourcePolling`` additionally tolerates the ~20s read-after-write window where the + # resource GET returns 404 "Cluster not found" right after the create is accepted. + if polling is True: + polling_method: PollingMethod = cast( + PollingMethod, + _ComputeResourcePolling( + lro_delay, + lro_algorithms=[BodyContentPolling(), StatusCheckPolling()], + path_format_arguments=path_format_arguments, + **kwargs, + ), + ) + elif polling is False: + polling_method = cast(PollingMethod, NoPolling()) + else: + polling_method = polling + return LROPoller[_models.Compute]( + self._client, pipeline_response, get_long_running_output, polling_method # type: ignore + ) + + +__all__: list[str] = [ + "ComputesOperations", +] # Add all objects you want publicly available to users at this package level def patch_sdk(): diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py new file mode 100644 index 000000000000..f93beb795089 --- /dev/null +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py @@ -0,0 +1,210 @@ +# coding=utf-8 +# -------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------- +"""Unit tests for the customized (patched) ComputesOperations.begin_create_or_update. + +These verify the behavior added to work around the async compute-create handling: + - the create accepts a 202 (async "Accepted") response (the generated code rejected it), + - the poller polls the compute *resource* (GET on the resource URL, watching + ``provisioningState``) via ``BodyContentPolling`` and never the operation-status endpoint + (``.../computeOperations/{id}``, which requires ``computeOperations/read``), + - the poller still blocks until a terminal state and raises on a genuine provisioning failure, and + - genuine non-2xx create errors still propagate. +""" +import json as _json +import time +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from azure.core.credentials import AccessToken +from azure.core.exceptions import HttpResponseError +from azure.core.polling import NoPolling +from azure.mgmt.core.polling.arm_polling import ( + ARMPolling, + AzureAsyncOperationPolling, + BodyContentPolling, + LocationPolling, + StatusCheckPolling, +) +from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient +from azure.mgmt.cognitiveservices.operations._patch import _ComputeResourcePolling + +RESOURCE_URL = ( + "https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000" + "/resourceGroups/rg/providers/Microsoft.CognitiveServices/accounts/acct" + "/computes/test-compute?api-version=2026-05-15-preview" +) + +ACCEPTED_BODY = {"name": "test-compute", "properties": {"provisioningState": "Accepted"}} +SUCCEEDED_BODY = { + "name": "test-compute", + "type": "Microsoft.CognitiveServices/accounts/computes", + "properties": {"provisioningState": "Succeeded"}, +} +NOT_FOUND_BODY = {"error": {"code": "NotFound", "message": "Cluster not found."}} + + +class _FakeHttpResponse: + """Minimal stand-in for an azure.core.rest HttpResponse used by the polling machinery.""" + + def __init__(self, status_code, body, method="PUT", url=RESOURCE_URL, headers=None): + self.status_code = status_code + self._body = body + self.request = SimpleNamespace(method=method, url=url, headers={"x-ms-client-request-id": "fake-request-id"}) + self.headers = headers or {} + self.reason = "reason" + self.content_type = "application/json" + + @property + def content(self): # presence of ``content`` makes azure-core treat this as a "rest" response + return _json.dumps(self._body).encode("utf-8") if self._body is not None else b"" + + def text(self, *args, **kwargs): + return _json.dumps(self._body) if self._body is not None else "" + + def json(self): + return self._body + + def read(self, *args, **kwargs): + return self.content + + +def _pipeline_response(status_code, body, method="PUT", url=RESOURCE_URL): + return SimpleNamespace(http_response=_FakeHttpResponse(status_code, body, method, url), context={}) + + +class _FakeCredential: + def get_token(self, *scopes, **kwargs): # pylint: disable=unused-argument + return AccessToken("fake-token", int(time.time()) + 3600) + + +def _make_computes(): + client = CognitiveServicesManagementClient( + credential=_FakeCredential(), subscription_id="00000000-0000-0000-0000-000000000000" + ) + return client.computes + + +def test_create_uses_resource_polling_not_computeoperations(): + """The poller is ARMPolling restricted to resource-based algorithms; the operation-status + algorithms (which would hit ``computeOperations/read``) are excluded.""" + computes = _make_computes() + # A terminal 200 lets the poller finish during initialize, so no background poll is needed here. + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(200, SUCCEEDED_BODY)) + + poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + + assert isinstance(poller._polling_method, ARMPolling) + assert isinstance(poller._polling_method, _ComputeResourcePolling) + algorithms = [type(a) for a in poller._polling_method._lro_algorithms] + assert BodyContentPolling in algorithms + assert StatusCheckPolling in algorithms + # The whole point of the fix: never follow Azure-AsyncOperation / Location to computeOperations. + assert AzureAsyncOperationPolling not in algorithms + assert LocationPolling not in algorithms + + result = poller.result() + assert result.name == "test-compute" + + +def test_create_accepts_202_and_polls_resource_until_succeeded(): + """A 202 create is accepted and the poller blocks, polling the resource URL until Succeeded.""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = MagicMock(return_value=_pipeline_response(200, SUCCEEDED_BODY, method="GET")) + + poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + result = poller.result() + + assert result.name == "test-compute" + assert result.properties.provisioning_state == "Succeeded" + # The status was polled from the resource itself, not the operation-status endpoint. + computes._client.send_request.assert_called() + polled_url = computes._client.send_request.call_args[0][0].url + assert "computeOperations" not in polled_url + assert "/computes/test-compute" in polled_url + + +def test_create_surfaces_provisioning_failure(): + """A create that provisions to Failed must raise with the resource's own error detail (not the + generic 'Operation returned an invalid status OK').""" + computes = _make_computes() + failed_body = { + "name": "test-compute", + "properties": { + "provisioningState": "Failed", + "errors": [{"code": "QuotaExceeded", "message": "exceeding subscription quota limits."}], + }, + } + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = MagicMock(return_value=_pipeline_response(200, failed_body, method="GET")) + + poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + with pytest.raises(HttpResponseError) as exc_info: + poller.result() + message = str(exc_info.value) + assert "QuotaExceeded" in message # the real reason is surfaced + assert "invalid status" not in message # not azure-core's generic fallback message + + +def test_create_propagates_non_2xx_error(): + """A genuine non-2xx create failure must still raise, without any polling.""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock( + return_value=_pipeline_response(400, {"error": {"code": "Bad", "message": "bad"}}) + ) + computes._client.send_request = MagicMock() + + with pytest.raises(HttpResponseError): + computes.begin_create_or_update("rg", "acct", "test-compute", b"{}") + computes._client.send_request.assert_not_called() + + +def test_create_polling_false_uses_no_polling_escape_hatch(): + """Callers can still opt out of blocking with ``polling=False`` (returns the accepted resource).""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = MagicMock() + + poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling=False) + + assert isinstance(poller._polling_method, NoPolling) + result = poller.result() + assert result.name == "test-compute" + computes._client.send_request.assert_not_called() + + +def test_create_tolerates_read_after_write_404_then_succeeds(): + """Right after the 202, a GET on the compute can return 404 'Cluster not found' for ~20s. The + poller must tolerate that window and keep polling instead of failing the successful create.""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = MagicMock( + side_effect=[ + _pipeline_response(404, NOT_FOUND_BODY, method="GET"), + _pipeline_response(200, SUCCEEDED_BODY, method="GET"), + ] + ) + + poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + result = poller.result() + + assert result.name == "test-compute" + assert computes._client.send_request.call_count == 2 # it kept polling past the 404 + + +def test_create_surfaces_persistent_not_found(monkeypatch): + """If the resource never becomes queryable, the bounded grace expires and the 404 surfaces, so a + genuinely missing resource does not hang forever.""" + monkeypatch.setattr(_ComputeResourcePolling, "_NOT_FOUND_GRACE_SECONDS", -1) + computes = _make_computes() + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = MagicMock(return_value=_pipeline_response(404, NOT_FOUND_BODY, method="GET")) + + poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + with pytest.raises(HttpResponseError): + poller.result() diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py new file mode 100644 index 000000000000..cdff70cbca92 --- /dev/null +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py @@ -0,0 +1,210 @@ +# coding=utf-8 +# -------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------- +"""Async unit tests for the customized (patched) ComputesOperations.begin_create_or_update. + +Mirror of the sync tests: the create accepts a 202, the poller polls the compute *resource* +(never ``computeOperations``), blocks until terminal, surfaces provisioning failures, and still +propagates genuine non-2xx create errors. +""" +import json as _json +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from azure.core.credentials import AccessToken +from azure.core.exceptions import HttpResponseError +from azure.core.polling import AsyncNoPolling +from azure.mgmt.core.polling.arm_polling import ( + AzureAsyncOperationPolling, + BodyContentPolling, + LocationPolling, + StatusCheckPolling, +) +from azure.mgmt.core.polling.async_arm_polling import AsyncARMPolling +from azure.mgmt.cognitiveservices.aio import CognitiveServicesManagementClient +from azure.mgmt.cognitiveservices.aio.operations._patch import _AsyncComputeResourcePolling + +RESOURCE_URL = ( + "https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000" + "/resourceGroups/rg/providers/Microsoft.CognitiveServices/accounts/acct" + "/computes/test-compute?api-version=2026-05-15-preview" +) + +ACCEPTED_BODY = {"name": "test-compute", "properties": {"provisioningState": "Accepted"}} +SUCCEEDED_BODY = { + "name": "test-compute", + "type": "Microsoft.CognitiveServices/accounts/computes", + "properties": {"provisioningState": "Succeeded"}, +} +NOT_FOUND_BODY = {"error": {"code": "NotFound", "message": "Cluster not found."}} + + +class _FakeAsyncHttpResponse: + """Minimal stand-in for an azure.core.rest AsyncHttpResponse used by the polling machinery.""" + + def __init__(self, status_code, body, method="PUT", url=RESOURCE_URL, headers=None): + self.status_code = status_code + self._body = body + self.request = SimpleNamespace(method=method, url=url, headers={"x-ms-client-request-id": "fake-request-id"}) + self.headers = headers or {} + self.reason = "reason" + self.content_type = "application/json" + + @property + def content(self): # presence of ``content`` makes azure-core treat this as a "rest" response + return _json.dumps(self._body).encode("utf-8") if self._body is not None else b"" + + def text(self, *args, **kwargs): + return _json.dumps(self._body) if self._body is not None else "" + + def json(self): + return self._body + + async def read(self, *args, **kwargs): + return self.content + + +def _pipeline_response(status_code, body, method="PUT", url=RESOURCE_URL): + return SimpleNamespace(http_response=_FakeAsyncHttpResponse(status_code, body, method, url), context={}) + + +class _FakeAsyncCredential: + async def get_token(self, *scopes, **kwargs): # pylint: disable=unused-argument + return AccessToken("fake-token", int(time.time()) + 3600) + + +def _make_computes(): + client = CognitiveServicesManagementClient( + credential=_FakeAsyncCredential(), subscription_id="00000000-0000-0000-0000-000000000000" + ) + return client.computes + + +@pytest.mark.asyncio +async def test_create_uses_resource_polling_not_computeoperations(): + """The poller is AsyncARMPolling restricted to resource-based algorithms; the operation-status + algorithms (which would hit ``computeOperations/read``) are excluded.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(200, SUCCEEDED_BODY)) + + poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + + assert isinstance(poller._polling_method, AsyncARMPolling) + assert isinstance(poller._polling_method, _AsyncComputeResourcePolling) + algorithms = [type(a) for a in poller._polling_method._lro_algorithms] + assert BodyContentPolling in algorithms + assert StatusCheckPolling in algorithms + assert AzureAsyncOperationPolling not in algorithms + assert LocationPolling not in algorithms + + result = await poller.result() + assert result.name == "test-compute" + + +@pytest.mark.asyncio +async def test_create_accepts_202_and_polls_resource_until_succeeded(): + """A 202 create is accepted and the poller blocks, polling the resource URL until Succeeded.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = AsyncMock(return_value=_pipeline_response(200, SUCCEEDED_BODY, method="GET")) + + poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + result = await poller.result() + + assert result.name == "test-compute" + assert result.properties.provisioning_state == "Succeeded" + computes._client.send_request.assert_called() + polled_url = computes._client.send_request.call_args[0][0].url + assert "computeOperations" not in polled_url + assert "/computes/test-compute" in polled_url + + +@pytest.mark.asyncio +async def test_create_surfaces_provisioning_failure(): + """A create that provisions to Failed must raise with the resource's own error detail (not the + generic 'Operation returned an invalid status OK').""" + computes = _make_computes() + failed_body = { + "name": "test-compute", + "properties": { + "provisioningState": "Failed", + "errors": [{"code": "QuotaExceeded", "message": "exceeding subscription quota limits."}], + }, + } + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = AsyncMock(return_value=_pipeline_response(200, failed_body, method="GET")) + + poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + with pytest.raises(HttpResponseError) as exc_info: + await poller.result() + message = str(exc_info.value) + assert "QuotaExceeded" in message # the real reason is surfaced + assert "invalid status" not in message # not azure-core's generic fallback message + + +@pytest.mark.asyncio +async def test_create_propagates_non_2xx_error(): + """A genuine non-2xx create failure must still raise, without any polling.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock( + return_value=_pipeline_response(400, {"error": {"code": "Bad", "message": "bad"}}) + ) + computes._client.send_request = AsyncMock() + + with pytest.raises(HttpResponseError): + await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}") + computes._client.send_request.assert_not_called() + + +@pytest.mark.asyncio +async def test_create_polling_false_uses_no_polling_escape_hatch(): + """Callers can still opt out of blocking with ``polling=False`` (returns the accepted resource).""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = AsyncMock() + + poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling=False) + + assert isinstance(poller._polling_method, AsyncNoPolling) + result = await poller.result() + assert result.name == "test-compute" + computes._client.send_request.assert_not_called() + + +@pytest.mark.asyncio +async def test_create_tolerates_read_after_write_404_then_succeeds(): + """Right after the 202, a GET on the compute can return 404 'Cluster not found' for ~20s. The + poller must tolerate that window and keep polling instead of failing the successful create.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = AsyncMock( + side_effect=[ + _pipeline_response(404, NOT_FOUND_BODY, method="GET"), + _pipeline_response(200, SUCCEEDED_BODY, method="GET"), + ] + ) + + poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + result = await poller.result() + + assert result.name == "test-compute" + assert computes._client.send_request.call_count == 2 # it kept polling past the 404 + + +@pytest.mark.asyncio +async def test_create_surfaces_persistent_not_found(monkeypatch): + """If the resource never becomes queryable, the bounded grace expires and the 404 surfaces, so a + genuinely missing resource does not hang forever.""" + monkeypatch.setattr(_AsyncComputeResourcePolling, "_NOT_FOUND_GRACE_SECONDS", -1) + computes = _make_computes() + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes._client.send_request = AsyncMock(return_value=_pipeline_response(404, NOT_FOUND_BODY, method="GET")) + + poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + with pytest.raises(HttpResponseError): + await poller.result() From 5baf28f604bc8511fddbe77ea852ac1d127d468e Mon Sep 17 00:00:00 2001 From: Swarup Chanda Date: Fri, 17 Jul 2026 17:10:47 +0530 Subject: [PATCH 2/5] [cognitiveservices] Track compute creation via the list API (blocking, no computeOperations) Poll provisioningState from the list API instead of a resource GET: a GET on a just-created compute returns 404 'Cluster not found' for an extended period while it provisions, whereas list reflects state from the moment the create is accepted. Keeps the 202 accept, blocks until terminal via a custom _ComputeListPolling, surfaces the compute's real error on failure, and preserves the public overloads, api-version validation and continuation_token resume path. Updates 22 unit tests and the CHANGELOG. --- .../azure-mgmt-cognitiveservices/CHANGELOG.md | 17 +- .../aio/operations/_patch.py | 390 +++++++++++----- .../cognitiveservices/operations/_patch.py | 442 +++++++++++++----- .../tests/test_computes_operations_patch.py | 180 ++++--- .../test_computes_operations_patch_async.py | 197 +++++--- 5 files changed, 844 insertions(+), 382 deletions(-) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/CHANGELOG.md b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/CHANGELOG.md index f18649325342..1b48c80da9d8 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/CHANGELOG.md +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/CHANGELOG.md @@ -11,15 +11,14 @@ (`.../locations/{location}/computeOperations/{operationId}`) required the `Microsoft.CognitiveServices/locations/computeOperations/read` permission, so callers allowed to create a compute but not granted that read permission were shown a misleading `AuthorizationFailed` error even - though the create succeeded. The 202 is now accepted and the long-running operation is tracked by polling - the compute resource itself (a `GET` on the resource URL, watching `provisioningState`) instead of the - operation-status endpoint, so it only needs the `computes/read` permission that callers already have. The - poller still blocks until the operation reaches a terminal state and raises on a genuine provisioning - failure; non-2xx create failures still propagate. Callers can opt out of blocking with `polling=False`. - The poller also tolerates the short read-after-write window right after the create where the compute is - briefly not yet queryable (a `GET` returns `404 "Cluster not found"`), so a successful create is not - wrongly reported as failed, and when provisioning does fail it surfaces the compute's own error detail - (e.g. a quota message) instead of a generic `Operation returned an invalid status 'OK'`. + though the create succeeded. The 202 is now accepted and the long-running operation is tracked by reading + `provisioningState` from the `list` API, so it only needs the `computes/read` permission that callers + already have. The `list` API is used rather than a `get` on the resource because a `GET` on a just-created + compute returns `404 "Cluster not found"` for an extended period while it provisions, whereas `list` + reflects the compute's state from the moment the create is accepted. The poller still blocks until the + operation reaches a terminal state and surfaces the compute's own error detail (e.g. a quota message) on a + genuine provisioning failure; non-2xx create failures still propagate, and callers can opt out of blocking + with `polling=False`. ## 15.0.0b3 (2026-06-26) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py index d73ae4060aa3..2fc764d672c4 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py @@ -7,11 +7,12 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ +import asyncio import json import time from collections.abc import MutableMapping from io import IOBase -from typing import IO, Any, Optional, Union, cast +from typing import IO, Any, AsyncIterator, Optional, Union, cast, overload from azure.core.exceptions import ( ClientAuthenticationError, @@ -19,124 +20,165 @@ ResourceExistsError, ResourceNotFoundError, ResourceNotModifiedError, + StreamClosedError, + StreamConsumedError, map_error, ) +from azure.core.pipeline import PipelineResponse from azure.core.polling import AsyncLROPoller, AsyncNoPolling, AsyncPollingMethod -from azure.core.polling.base_polling import _raise_if_bad_http_status_and_method from azure.core.tracing.decorator_async import distributed_trace_async from azure.core.utils import case_insensitive_dict from azure.mgmt.core.exceptions import ARMErrorFormat -from azure.mgmt.core.polling.arm_polling import BodyContentPolling, StatusCheckPolling -from azure.mgmt.core.polling.async_arm_polling import AsyncARMPolling from ... import models as _models from ..._utils.model_base import SdkJSONEncoder, _deserialize, _failsafe_deserialize +from ..._validation import api_version_validation +from ...operations._patch import _TERMINAL_FAILED_STATES, _TERMINAL_SUCCESS_STATES, _provisioning_error, _state_of from ...operations._operations import build_computes_create_or_update_request -from ...operations._patch import _TERMINAL_FAILED_STATES, _compute_provisioning_error from ._operations import ComputesOperations as _ComputesOperationsGenerated JSON = MutableMapping[str, Any] -class _AsyncComputeResourcePolling(AsyncARMPolling): - """``AsyncARMPolling`` that tolerates the compute-create read-after-write window. +class _AsyncComputeListPolling(AsyncPollingMethod): + """Async blocking LRO poller that reads compute status from the ``list`` API instead of ``get``. - Right after the service accepts the create (HTTP 202), the new compute is briefly not queryable: a - ``GET`` on the resource returns ``404 "Cluster not found"`` for ~15-20s. Plain resource polling - (:class:`BodyContentPolling`) would treat that 404 as a failure and wrongly fail a create that - actually succeeded. This subclass treats a 404 as "still in progress" and keeps polling until the - resource is queryable and reaches a terminal ``provisioningState``. A generous grace period bounds - the tolerance so a genuinely missing resource still surfaces instead of hanging forever. When the - compute provisions to a failed state it raises the resource's own error detail rather than the - generic ``Operation returned an invalid status 'OK'``. + The compute create is asynchronous (the service returns HTTP 202) and the new compute is **not** + reliably queryable via ``get`` during creation: a ``GET`` on the resource returns + ``404 "Cluster not found"`` for an extended, variable period while the cluster backend + materializes. The compute does, however, appear in the ``list`` response with its + ``provisioningState`` from the moment the create is accepted. This poller therefore polls ``list`` + (matching the resource by name) until the compute reaches a terminal ``provisioningState``. + + It blocks like a normal long-running-operation poller, never calls the permission-gated + ``.../computeOperations/{id}`` endpoint (so it does not need ``computeOperations/read``), and + surfaces the compute's own error detail on failure. A bounded grace period bounds how long the + compute may be absent from ``list`` before the poller gives up. """ _NOT_FOUND_GRACE_SECONDS = 300 - async def update_status(self) -> None: - self._pipeline_response = await self.request_status(self._operation.get_polling_url()) - response = self._pipeline_response.http_response - if response.status_code == 404: - first_seen = getattr(self, "_not_found_since", None) - if first_seen is None: - self._not_found_since = time.monotonic() - elif time.monotonic() - first_seen > self._NOT_FOUND_GRACE_SECONDS: - _raise_if_bad_http_status_and_method(response) # grace exhausted: surface the 404 - self._status = "InProgress" - return - self._not_found_since = None - _raise_if_bad_http_status_and_method(response) - self._status = self._operation.get_status(self._pipeline_response) - if str(self._status).lower() in _TERMINAL_FAILED_STATES: - raise _compute_provisioning_error(response) + def __init__( + self, + operations: "ComputesOperations", + resource_group_name: str, + account_name: str, + compute_name: str, + interval: float, + ) -> None: + self._operations = operations + self._resource_group_name = resource_group_name + self._account_name = account_name + self._compute_name = compute_name + self._interval = interval + self._status = "InProgress" + self._resource: Optional[_models.Compute] = None + self._first_missing_at: Optional[float] = None + + def initialize(self, client: Any, initial_response: Any, deserialization_callback: Any) -> None: + self._initial_response = initial_response + + async def _current_compute(self) -> Optional[_models.Compute]: + async for compute in self._operations.list(self._resource_group_name, self._account_name): + if getattr(compute, "name", None) == self._compute_name: + return compute + return None + + async def run(self) -> None: + while not self.finished(): + compute = await self._current_compute() + if compute is None: + # The create was accepted, so the compute should appear in ``list`` shortly. Tolerate a + # brief absence, but give up after a bounded grace so a create that never materializes + # does not hang forever. + if self._first_missing_at is None: + self._first_missing_at = time.monotonic() + elif time.monotonic() - self._first_missing_at > self._NOT_FOUND_GRACE_SECONDS: + raise HttpResponseError( + message=f"Compute '{self._compute_name}' did not appear in the account's compute " + "list after it was created." + ) + else: + self._first_missing_at = None + self._resource = compute + state = _state_of(compute) + if state in _TERMINAL_SUCCESS_STATES: + self._status = "Succeeded" + elif state in _TERMINAL_FAILED_STATES: + self._status = "Failed" + raise _provisioning_error(compute) + if not self.finished(): + await asyncio.sleep(self._interval) + + def status(self) -> str: + return self._status + + def finished(self) -> bool: + return self._status in ("Succeeded", "Failed", "Canceled") + + def resource(self) -> Optional[_models.Compute]: + return self._resource + + def get_continuation_token(self) -> str: + return self._compute_name + + @classmethod + def from_continuation_token(cls, continuation_token: str, **kwargs: Any): + client = kwargs["client"] + deserialization_callback = kwargs["deserialization_callback"] + return client, None, deserialization_callback class ComputesOperations(_ComputesOperationsGenerated): - """Customized ``ComputesOperations`` that polls the compute resource instead of its - operation-status endpoint. + """Customized ``ComputesOperations`` that tracks compute creation via the ``list`` API. The generated :meth:`begin_create_or_update` has two problems against the current service: - * it rejects the ``202 Accepted`` the service returns for the async create with - "Operation returned an invalid status 'Accepted'"; and + * :meth:`_create_or_update_initial` rejects the ``202 Accepted`` the service returns for the async + create with "Operation returned an invalid status 'Accepted'"; and * its poller follows the ``Azure-AsyncOperation`` header to ``.../locations/{location}/computeOperations/{operationId}``, which requires the ``Microsoft.CognitiveServices/locations/computeOperations/read`` permission. Many callers who are allowed to create a compute lack that permission, so the poll fails with ``AuthorizationFailed`` even though the create itself succeeded. - This override accepts the ``202`` and polls the compute resource itself (a ``GET`` on the resource - URL, watching ``provisioningState``) via :class:`BodyContentPolling`, deliberately excluding the - algorithms that would follow the ``Azure-AsyncOperation``/``Location`` headers. As a result the - poller only needs ``computes/read`` (which callers already have), still blocks until the operation - reaches a terminal state, and surfaces a genuine provisioning failure instead of masking it. + This override keeps the generated method's public shape (overloads, api-version validation and + ``continuation_token`` resume path) unchanged. It (1) accepts the ``202`` in + :meth:`_create_or_update_initial` and (2) blocks on an :class:`_AsyncComputeListPolling` poller + that reads ``provisioningState`` from the ``list`` API. ``list`` is used rather than ``get`` + because a ``GET`` on a just-created compute returns ``404 "Cluster not found"`` for an extended + period while it provisions, whereas ``list`` reflects the compute's state from the moment the + create is accepted. The poller only needs ``computes/read``, still blocks until a terminal state, + and surfaces a genuine provisioning failure instead of masking it. """ - @distributed_trace_async - async def begin_create_or_update( + @api_version_validation( + method_added_on="2026-03-15-preview", + params_added_on={ + "2026-03-15-preview": [ + "api_version", + "subscription_id", + "resource_group_name", + "account_name", + "compute_name", + "content_type", + "accept", + ] + }, + api_versions_list=["2026-03-15-preview", "2026-05-15-preview"], + ) + async def _create_or_update_initial( self, resource_group_name: str, account_name: str, compute_name: str, resource: Union[_models.Compute, JSON, IO[bytes]], **kwargs: Any, - ) -> AsyncLROPoller[_models.Compute]: - """Creates or updates a compute associated with the Cognitive Services account. - - This override accepts the service's ``202 Accepted`` and polls the compute resource itself - (a ``GET`` on the resource URL, watching ``provisioningState``) instead of the operation-status - endpoint that requires ``computeOperations/read``. It still blocks until the operation reaches - a terminal state and raises on a genuine provisioning failure, so callers get correct results - without needing the operation-status read permission. - - :param resource_group_name: The name of the resource group. The name is case insensitive. - Required. - :type resource_group_name: str - :param account_name: The name of Cognitive Services account. Required. - :type account_name: str - :param compute_name: The name of the compute associated with the Cognitive Services Account. - Required. - :type compute_name: str - :param resource: The compute properties. Is one of the following types: Compute, JSON, - IO[bytes] Required. - :type resource: ~azure.mgmt.cognitiveservices.models.Compute or JSON or IO[bytes] - :return: An instance of AsyncLROPoller that returns Compute. The Compute is compatible with - MutableMapping - :rtype: ~azure.core.polling.AsyncLROPoller[~azure.mgmt.cognitiveservices.models.Compute] - :raises ~azure.core.exceptions.HttpResponseError: - """ - _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) - _params = kwargs.pop("params", {}) or {} - - content_type: Optional[str] = ( - kwargs.pop("content_type", _headers.pop("Content-Type", None)) or "application/json" - ) - cls = kwargs.pop("cls", None) - polling: Union[bool, AsyncPollingMethod] = kwargs.pop("polling", True) - lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) - kwargs.pop("continuation_token", None) - - error_map: dict = { + ) -> AsyncIterator[bytes]: + # Mirrors the generated ``_create_or_update_initial`` exactly, except it accepts ``202``. The + # generated code only allows 200/201 and rejects the service's async 202 with + # "Operation returned an invalid status 'Accepted'", failing a create the service accepted. + error_map: MutableMapping = { 401: ClientAuthenticationError, 404: ResourceNotFoundError, 409: ResourceExistsError, @@ -144,6 +186,14 @@ async def begin_create_or_update( } error_map.update(kwargs.pop("error_map", {}) or {}) + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + _params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None)) + cls: Any = kwargs.pop("cls", None) + + content_type = content_type or "application/json" + _content = None if isinstance(resource, (IOBase, bytes)): _content = resource else: @@ -160,58 +210,182 @@ async def begin_create_or_update( headers=_headers, params=_params, ) - _request.url = self._client.format_url( - _request.url, - endpoint=self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), - ) + path_format_arguments = { + "endpoint": self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), + } + _request.url = self._client.format_url(_request.url, **path_format_arguments) - pipeline_response = await self._client._pipeline.run( # pylint: disable=protected-access - _request, stream=False, **kwargs + _decompress = kwargs.pop("decompress", True) + _stream = True + pipeline_response: PipelineResponse = await self._client._pipeline.run( # pylint: disable=protected-access + _request, stream=_stream, **kwargs ) + response = pipeline_response.http_response - # Accept 202 (async "Accepted") in addition to 200/201. The generated create rejects 202 with - # "Operation returned an invalid status 'Accepted'", failing a create the service accepted. + if response.status_code not in [200, 201, 202]: + try: + await response.read() # Load the body in memory and close the socket + except (StreamConsumedError, StreamClosedError): + pass map_error(status_code=response.status_code, response=response, error_map=error_map) error = _failsafe_deserialize(_models.ErrorResponse, response) raise HttpResponseError(response=response, model=error, error_format=ARMErrorFormat) - await response.read() # load the body so the poller/deserializer can read it + + response_headers = {} + if response.status_code in [201, 202]: + response_headers["Azure-AsyncOperation"] = self._deserialize( + "str", response.headers.get("Azure-AsyncOperation") + ) + response_headers["Retry-After"] = self._deserialize("int", response.headers.get("Retry-After")) + + deserialized = response.iter_bytes() if _decompress else response.iter_raw() + + if cls: + return cls(pipeline_response, deserialized, response_headers) # type: ignore + + return deserialized # type: ignore + + @overload + async def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: _models.Compute, + *, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncLROPoller[_models.Compute]: ... + @overload + async def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: JSON, + *, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncLROPoller[_models.Compute]: ... + @overload + async def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: IO[bytes], + *, + content_type: str = "application/json", + **kwargs: Any, + ) -> AsyncLROPoller[_models.Compute]: ... + + @distributed_trace_async + @api_version_validation( + method_added_on="2026-03-15-preview", + params_added_on={ + "2026-03-15-preview": [ + "api_version", + "subscription_id", + "resource_group_name", + "account_name", + "compute_name", + "content_type", + "accept", + ] + }, + api_versions_list=["2026-03-15-preview", "2026-05-15-preview"], + ) + async def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: Union[_models.Compute, JSON, IO[bytes]], + **kwargs: Any, + ) -> AsyncLROPoller[_models.Compute]: + """Creates or updates a compute associated with the Cognitive Services account. + + This override accepts the service's ``202 Accepted`` and blocks until the compute reaches a + terminal state, reading ``provisioningState`` from the ``list`` API rather than polling the + operation-status endpoint (which requires ``computeOperations/read``). ``list`` is used because + a ``GET`` on a just-created compute returns ``404 "Cluster not found"`` while it provisions. It + raises on a genuine provisioning failure, so callers get correct results without needing the + operation-status read permission. + + :param resource_group_name: The name of the resource group. The name is case insensitive. + Required. + :type resource_group_name: str + :param account_name: The name of Cognitive Services account. Required. + :type account_name: str + :param compute_name: The name of the compute associated with the Cognitive Services Account. + Required. + :type compute_name: str + :param resource: The compute properties. Is one of the following types: Compute, JSON, + IO[bytes] Required. + :type resource: ~azure.mgmt.cognitiveservices.models.Compute or JSON or IO[bytes] + :return: An instance of AsyncLROPoller that returns Compute. The Compute is compatible with + MutableMapping + :rtype: ~azure.core.polling.AsyncLROPoller[~azure.mgmt.cognitiveservices.models.Compute] + :raises ~azure.core.exceptions.HttpResponseError: + """ + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + _params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None)) + cls: Any = kwargs.pop("cls", None) + polling: Union[bool, AsyncPollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + cont_token: Optional[str] = kwargs.pop("continuation_token", None) + if cont_token is None: + raw_result = await self._create_or_update_initial( + resource_group_name=resource_group_name, + account_name=account_name, + compute_name=compute_name, + resource=resource, + content_type=content_type, + cls=lambda x, y, z: x, + headers=_headers, + params=_params, + **kwargs, + ) + await raw_result.http_response.read() # type: ignore + # A synchronous terminal failure (rare: the accepted body itself is already Failed/Canceled) + # would otherwise slip through to a poller that would just keep listing; surface the + # resource's own error detail immediately instead. + try: + initial_body = raw_result.http_response.json() + except Exception: # pylint: disable=broad-except + initial_body = None + if isinstance(initial_body, MutableMapping) and _state_of(initial_body) in _TERMINAL_FAILED_STATES: + raise _provisioning_error(initial_body) + kwargs.pop("error_map", None) def get_long_running_output(pipeline_response): response = pipeline_response.http_response deserialized = _deserialize(_models.Compute, response.json()) if cls: - return cls(pipeline_response, deserialized, {}) + return cls(pipeline_response, deserialized, {}) # type: ignore return deserialized - path_format_arguments = { - "endpoint": self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), - } - - # Poll the compute resource itself (a GET on the resource URL, watching ``provisioningState``) - # instead of the operation-status endpoint. ``BodyContentPolling`` polls the resource URL, so - # the poller never calls ``.../computeOperations/{id}`` (which requires ``computeOperations/read`` - # that many callers lack). Excluding ``AzureAsyncOperationPolling`` and ``LocationPolling`` - # guarantees the ``Azure-AsyncOperation``/``Location`` headers the service returns are ignored, - # while the poller still blocks until terminal and raises on a failed provisioning. - # ``_AsyncComputeResourcePolling`` additionally tolerates the ~20s read-after-write window where - # the resource GET returns 404 "Cluster not found" right after the create is accepted. if polling is True: polling_method: AsyncPollingMethod = cast( AsyncPollingMethod, - _AsyncComputeResourcePolling( - lro_delay, - lro_algorithms=[BodyContentPolling(), StatusCheckPolling()], - path_format_arguments=path_format_arguments, - **kwargs, - ), + _AsyncComputeListPolling(self, resource_group_name, account_name, compute_name, lro_delay), ) elif polling is False: polling_method = cast(AsyncPollingMethod, AsyncNoPolling()) else: polling_method = polling + if cont_token: + return AsyncLROPoller[_models.Compute].from_continuation_token( + polling_method=polling_method, + continuation_token=cont_token, + client=self._client, + deserialization_callback=get_long_running_output, + ) return AsyncLROPoller[_models.Compute]( - self._client, pipeline_response, get_long_running_output, polling_method # type: ignore + self._client, raw_result, get_long_running_output, polling_method # type: ignore ) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py index acc4a2b830e0..55948b11be22 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py @@ -11,7 +11,7 @@ import time from collections.abc import MutableMapping from io import IOBase -from typing import IO, Any, Optional, Union, cast +from typing import IO, Any, Iterator, Optional, Union, cast, overload from azure.core.exceptions import ( ClientAuthenticationError, @@ -19,152 +19,208 @@ ResourceExistsError, ResourceNotFoundError, ResourceNotModifiedError, + StreamClosedError, + StreamConsumedError, map_error, ) +from azure.core.pipeline import PipelineResponse from azure.core.polling import LROPoller, NoPolling, PollingMethod -from azure.core.polling.base_polling import _raise_if_bad_http_status_and_method from azure.core.tracing.decorator import distributed_trace from azure.core.utils import case_insensitive_dict from azure.mgmt.core.exceptions import ARMErrorFormat -from azure.mgmt.core.polling.arm_polling import ARMPolling, BodyContentPolling, StatusCheckPolling from .. import models as _models from .._utils.model_base import SdkJSONEncoder, _deserialize, _failsafe_deserialize +from .._validation import api_version_validation from ._operations import ComputesOperations as _ComputesOperationsGenerated from ._operations import build_computes_create_or_update_request JSON = MutableMapping[str, Any] _TERMINAL_FAILED_STATES = frozenset({"failed", "canceled", "cancelled"}) +_TERMINAL_SUCCESS_STATES = frozenset({"succeeded"}) -def _compute_provisioning_error(response: Any) -> HttpResponseError: - """Build a descriptive error from a compute resource whose provisioning failed. +def _state_of(compute: Any) -> str: + """Return a compute's ``provisioningState`` as a lowercase string (handles enum or str).""" + props = getattr(compute, "properties", None) + state = getattr(props, "provisioning_state", None) if props is not None else None + if state is None and isinstance(compute, MutableMapping): + state = (compute.get("properties") or {}).get("provisioningState") + if state is None: + return "" + return str(getattr(state, "value", state)).lower() - Because the poller watches the compute resource (not the operation-status endpoint), the failure + +def _provisioning_error(compute: Any) -> HttpResponseError: + """Build a descriptive error from a compute whose provisioning reached a failed state. + + Because the poller reads status from ``list`` (not the operation-status endpoint), the failure detail lives in the compute's ``properties.errors``. Surface that instead of azure-core's generic - ``Operation returned an invalid status 'OK'`` (which happens because the terminal poll is a plain - HTTP 200 with no OData error body). + ``Operation failed or canceled`` / ``invalid status`` messages. """ - state = None - detail = "" - try: - body = response.json() - props = body.get("properties", {}) if isinstance(body, dict) else {} - state = props.get("provisioningState") - parts = [] - for err in props.get("errors") or []: - if isinstance(err, dict): - code, message = err.get("code"), err.get("message") - parts.append(f"{code}: {message}" if code and message else str(message or code)) - detail = "; ".join(p for p in parts if p) - except Exception: # pylint: disable=broad-except - pass - text = f"Compute provisioning {state or 'failed'}." + props = getattr(compute, "properties", None) + errors = getattr(props, "errors", None) if props is not None else None + if errors is None and isinstance(compute, MutableMapping): + errors = (compute.get("properties") or {}).get("errors") + parts = [] + for err in errors or []: + if isinstance(err, MutableMapping): + code, message = err.get("code"), err.get("message") + else: + code, message = getattr(err, "code", None), getattr(err, "message", None) + if code and message: + parts.append(f"{code}: {message}") + elif message or code: + parts.append(str(message or code)) + detail = "; ".join(p for p in parts if p) + state = _state_of(compute) or "failed" + text = f"Compute provisioning {state}." if detail: text = f"{text} {detail}" - return HttpResponseError(response=response, message=text) + return HttpResponseError(message=text) + +class _ComputeListPolling(PollingMethod): + """Blocking LRO poller that reads compute status from the ``list`` API instead of ``get``. -class _ComputeResourcePolling(ARMPolling): - """``ARMPolling`` that tolerates the compute-create read-after-write window. + The compute create is asynchronous (the service returns HTTP 202) and the new compute is **not** + reliably queryable via ``get`` during creation: a ``GET`` on the resource returns + ``404 "Cluster not found"`` for an extended, variable period while the cluster backend + materializes. The compute does, however, appear in the ``list`` response with its + ``provisioningState`` from the moment the create is accepted. This poller therefore polls ``list`` + (matching the resource by name) until the compute reaches a terminal ``provisioningState``. - Right after the service accepts the create (HTTP 202), the new compute is briefly not queryable: a - ``GET`` on the resource returns ``404 "Cluster not found"`` for ~15-20s. Plain resource polling - (:class:`BodyContentPolling`) would treat that 404 as a failure and wrongly fail a create that - actually succeeded. This subclass treats a 404 as "still in progress" and keeps polling until the - resource is queryable and reaches a terminal ``provisioningState``. A generous grace period bounds - the tolerance so a genuinely missing resource still surfaces instead of hanging forever. When the - compute provisions to a failed state it raises the resource's own error detail rather than the - generic ``Operation returned an invalid status 'OK'``. + It blocks like a normal long-running-operation poller, never calls the permission-gated + ``.../computeOperations/{id}`` endpoint (so it does not need ``computeOperations/read``), and + surfaces the compute's own error detail on failure. A bounded grace period bounds how long the + compute may be absent from ``list`` before the poller gives up, so a create that never materializes + does not hang forever. """ _NOT_FOUND_GRACE_SECONDS = 300 - def update_status(self) -> None: - self._pipeline_response = self.request_status(self._operation.get_polling_url()) - response = self._pipeline_response.http_response - if response.status_code == 404: - first_seen = getattr(self, "_not_found_since", None) - if first_seen is None: - self._not_found_since = time.monotonic() - elif time.monotonic() - first_seen > self._NOT_FOUND_GRACE_SECONDS: - _raise_if_bad_http_status_and_method(response) # grace exhausted: surface the 404 - self._status = "InProgress" - return - self._not_found_since = None - _raise_if_bad_http_status_and_method(response) - self._status = self._operation.get_status(self._pipeline_response) - if str(self._status).lower() in _TERMINAL_FAILED_STATES: - raise _compute_provisioning_error(response) + def __init__( + self, + operations: "ComputesOperations", + resource_group_name: str, + account_name: str, + compute_name: str, + interval: float, + ) -> None: + self._operations = operations + self._resource_group_name = resource_group_name + self._account_name = account_name + self._compute_name = compute_name + self._interval = interval + self._status = "InProgress" + self._resource: Optional[_models.Compute] = None + self._first_missing_at: Optional[float] = None + + def initialize(self, client: Any, initial_response: Any, deserialization_callback: Any) -> None: + self._initial_response = initial_response + + def _current_compute(self) -> Optional[_models.Compute]: + for compute in self._operations.list(self._resource_group_name, self._account_name): + if getattr(compute, "name", None) == self._compute_name: + return compute + return None + + def run(self) -> None: + while not self.finished(): + compute = self._current_compute() + if compute is None: + # The create was accepted, so the compute should appear in ``list`` shortly. Tolerate a + # brief absence, but give up after a bounded grace so a create that never materializes + # does not hang forever. + if self._first_missing_at is None: + self._first_missing_at = time.monotonic() + elif time.monotonic() - self._first_missing_at > self._NOT_FOUND_GRACE_SECONDS: + raise HttpResponseError( + message=f"Compute '{self._compute_name}' did not appear in the account's compute " + "list after it was created." + ) + else: + self._first_missing_at = None + self._resource = compute + state = _state_of(compute) + if state in _TERMINAL_SUCCESS_STATES: + self._status = "Succeeded" + elif state in _TERMINAL_FAILED_STATES: + self._status = "Failed" + raise _provisioning_error(compute) + if not self.finished(): + time.sleep(self._interval) + + def status(self) -> str: + return self._status + + def finished(self) -> bool: + return self._status in ("Succeeded", "Failed", "Canceled") + + def resource(self) -> Optional[_models.Compute]: + return self._resource + + def get_continuation_token(self) -> str: + return self._compute_name + + @classmethod + def from_continuation_token(cls, continuation_token: str, **kwargs: Any): + client = kwargs["client"] + deserialization_callback = kwargs["deserialization_callback"] + return client, None, deserialization_callback class ComputesOperations(_ComputesOperationsGenerated): - """Customized ``ComputesOperations`` that polls the compute resource instead of its - operation-status endpoint. + """Customized ``ComputesOperations`` that tracks compute creation via the ``list`` API. The generated :meth:`begin_create_or_update` has two problems against the current service: - * it rejects the ``202 Accepted`` the service returns for the async create with - "Operation returned an invalid status 'Accepted'"; and + * :meth:`_create_or_update_initial` rejects the ``202 Accepted`` the service returns for the async + create with "Operation returned an invalid status 'Accepted'"; and * its poller follows the ``Azure-AsyncOperation`` header to ``.../locations/{location}/computeOperations/{operationId}``, which requires the ``Microsoft.CognitiveServices/locations/computeOperations/read`` permission. Many callers who are allowed to create a compute lack that permission, so the poll fails with ``AuthorizationFailed`` even though the create itself succeeded. - This override accepts the ``202`` and polls the compute resource itself (a ``GET`` on the resource - URL, watching ``provisioningState``) via :class:`BodyContentPolling`, deliberately excluding the - algorithms that would follow the ``Azure-AsyncOperation``/``Location`` headers. As a result the - poller only needs ``computes/read`` (which callers already have), still blocks until the operation - reaches a terminal state, and surfaces a genuine provisioning failure instead of masking it. + This override keeps the generated method's public shape (overloads, api-version validation and + ``continuation_token`` resume path) unchanged. It (1) accepts the ``202`` in + :meth:`_create_or_update_initial` and (2) blocks on a :class:`_ComputeListPolling` poller that + reads ``provisioningState`` from the ``list`` API. ``list`` is used rather than ``get`` because a + ``GET`` on a just-created compute returns ``404 "Cluster not found"`` for an extended period while + it provisions, whereas ``list`` reflects the compute's state from the moment the create is + accepted. The poller only needs ``computes/read``, still blocks until a terminal state, and + surfaces a genuine provisioning failure instead of masking it. """ - @distributed_trace - def begin_create_or_update( + @api_version_validation( + method_added_on="2026-03-15-preview", + params_added_on={ + "2026-03-15-preview": [ + "api_version", + "subscription_id", + "resource_group_name", + "account_name", + "compute_name", + "content_type", + "accept", + ] + }, + api_versions_list=["2026-03-15-preview", "2026-05-15-preview"], + ) + def _create_or_update_initial( self, resource_group_name: str, account_name: str, compute_name: str, resource: Union[_models.Compute, JSON, IO[bytes]], **kwargs: Any, - ) -> LROPoller[_models.Compute]: - """Creates or updates a compute associated with the Cognitive Services account. - - This override accepts the service's ``202 Accepted`` and polls the compute resource itself - (a ``GET`` on the resource URL, watching ``provisioningState``) instead of the operation-status - endpoint that requires ``computeOperations/read``. It still blocks until the operation reaches - a terminal state and raises on a genuine provisioning failure, so callers get correct results - without needing the operation-status read permission. - - :param resource_group_name: The name of the resource group. The name is case insensitive. - Required. - :type resource_group_name: str - :param account_name: The name of Cognitive Services account. Required. - :type account_name: str - :param compute_name: The name of the compute associated with the Cognitive Services Account. - Required. - :type compute_name: str - :param resource: The compute properties. Is one of the following types: Compute, JSON, - IO[bytes] Required. - :type resource: ~azure.mgmt.cognitiveservices.models.Compute or JSON or IO[bytes] - :return: An instance of LROPoller that returns Compute. The Compute is compatible with - MutableMapping - :rtype: ~azure.core.polling.LROPoller[~azure.mgmt.cognitiveservices.models.Compute] - :raises ~azure.core.exceptions.HttpResponseError: - """ - _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) - _params = kwargs.pop("params", {}) or {} - - content_type: Optional[str] = ( - kwargs.pop("content_type", _headers.pop("Content-Type", None)) or "application/json" - ) - cls = kwargs.pop("cls", None) - polling: Union[bool, PollingMethod] = kwargs.pop("polling", True) - lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) - kwargs.pop("continuation_token", None) - - error_map: dict = { + ) -> Iterator[bytes]: + # Mirrors the generated ``_create_or_update_initial`` exactly, except it accepts ``202``. The + # generated code only allows 200/201 and rejects the service's async 202 with + # "Operation returned an invalid status 'Accepted'", failing a create the service accepted. + error_map: MutableMapping = { 401: ClientAuthenticationError, 404: ResourceNotFoundError, 409: ResourceExistsError, @@ -172,6 +228,14 @@ def begin_create_or_update( } error_map.update(kwargs.pop("error_map", {}) or {}) + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + _params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None)) + cls: Any = kwargs.pop("cls", None) + + content_type = content_type or "application/json" + _content = None if isinstance(resource, (IOBase, bytes)): _content = resource else: @@ -188,58 +252,182 @@ def begin_create_or_update( headers=_headers, params=_params, ) - _request.url = self._client.format_url( - _request.url, - endpoint=self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), + path_format_arguments = { + "endpoint": self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), + } + _request.url = self._client.format_url(_request.url, **path_format_arguments) + + _decompress = kwargs.pop("decompress", True) + _stream = True + pipeline_response: PipelineResponse = self._client._pipeline.run( # pylint: disable=protected-access + _request, stream=_stream, **kwargs ) - pipeline_response = self._client._pipeline.run( - _request, stream=False, **kwargs - ) # pylint: disable=protected-access response = pipeline_response.http_response - # Accept 202 (async "Accepted") in addition to 200/201. The generated create rejects 202 with - # "Operation returned an invalid status 'Accepted'", failing a create the service accepted. + if response.status_code not in [200, 201, 202]: + try: + response.read() # Load the body in memory and close the socket + except (StreamConsumedError, StreamClosedError): + pass map_error(status_code=response.status_code, response=response, error_map=error_map) error = _failsafe_deserialize(_models.ErrorResponse, response) raise HttpResponseError(response=response, model=error, error_format=ARMErrorFormat) - response.read() # load the body so the poller/deserializer can read it + + response_headers = {} + if response.status_code in [201, 202]: + response_headers["Azure-AsyncOperation"] = self._deserialize( + "str", response.headers.get("Azure-AsyncOperation") + ) + response_headers["Retry-After"] = self._deserialize("int", response.headers.get("Retry-After")) + + deserialized = response.iter_bytes() if _decompress else response.iter_raw() + + if cls: + return cls(pipeline_response, deserialized, response_headers) # type: ignore + + return deserialized # type: ignore + + @overload + def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: _models.Compute, + *, + content_type: str = "application/json", + **kwargs: Any, + ) -> LROPoller[_models.Compute]: ... + @overload + def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: JSON, + *, + content_type: str = "application/json", + **kwargs: Any, + ) -> LROPoller[_models.Compute]: ... + @overload + def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: IO[bytes], + *, + content_type: str = "application/json", + **kwargs: Any, + ) -> LROPoller[_models.Compute]: ... + + @distributed_trace + @api_version_validation( + method_added_on="2026-03-15-preview", + params_added_on={ + "2026-03-15-preview": [ + "api_version", + "subscription_id", + "resource_group_name", + "account_name", + "compute_name", + "content_type", + "accept", + ] + }, + api_versions_list=["2026-03-15-preview", "2026-05-15-preview"], + ) + def begin_create_or_update( + self, + resource_group_name: str, + account_name: str, + compute_name: str, + resource: Union[_models.Compute, JSON, IO[bytes]], + **kwargs: Any, + ) -> LROPoller[_models.Compute]: + """Creates or updates a compute associated with the Cognitive Services account. + + This override accepts the service's ``202 Accepted`` and blocks until the compute reaches a + terminal state, reading ``provisioningState`` from the ``list`` API rather than polling the + operation-status endpoint (which requires ``computeOperations/read``). ``list`` is used because + a ``GET`` on a just-created compute returns ``404 "Cluster not found"`` while it provisions. It + raises on a genuine provisioning failure, so callers get correct results without needing the + operation-status read permission. + + :param resource_group_name: The name of the resource group. The name is case insensitive. + Required. + :type resource_group_name: str + :param account_name: The name of Cognitive Services account. Required. + :type account_name: str + :param compute_name: The name of the compute associated with the Cognitive Services Account. + Required. + :type compute_name: str + :param resource: The compute properties. Is one of the following types: Compute, JSON, + IO[bytes] Required. + :type resource: ~azure.mgmt.cognitiveservices.models.Compute or JSON or IO[bytes] + :return: An instance of LROPoller that returns Compute. The Compute is compatible with + MutableMapping + :rtype: ~azure.core.polling.LROPoller[~azure.mgmt.cognitiveservices.models.Compute] + :raises ~azure.core.exceptions.HttpResponseError: + """ + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + _params = kwargs.pop("params", {}) or {} + + content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None)) + cls: Any = kwargs.pop("cls", None) + polling: Union[bool, PollingMethod] = kwargs.pop("polling", True) + lro_delay = kwargs.pop("polling_interval", self._config.polling_interval) + cont_token: Optional[str] = kwargs.pop("continuation_token", None) + if cont_token is None: + raw_result = self._create_or_update_initial( + resource_group_name=resource_group_name, + account_name=account_name, + compute_name=compute_name, + resource=resource, + content_type=content_type, + cls=lambda x, y, z: x, + headers=_headers, + params=_params, + **kwargs, + ) + raw_result.http_response.read() # type: ignore + # A synchronous terminal failure (rare: the accepted body itself is already Failed/Canceled) + # would otherwise slip through to a poller that would just keep listing; surface the + # resource's own error detail immediately instead. + try: + initial_body = raw_result.http_response.json() + except Exception: # pylint: disable=broad-except + initial_body = None + if isinstance(initial_body, MutableMapping) and _state_of(initial_body) in _TERMINAL_FAILED_STATES: + raise _provisioning_error(initial_body) + kwargs.pop("error_map", None) def get_long_running_output(pipeline_response): response = pipeline_response.http_response deserialized = _deserialize(_models.Compute, response.json()) if cls: - return cls(pipeline_response, deserialized, {}) + return cls(pipeline_response, deserialized, {}) # type: ignore return deserialized - path_format_arguments = { - "endpoint": self._serialize.url("self._config.base_url", self._config.base_url, "str", skip_quote=True), - } - - # Poll the compute resource itself (a GET on the resource URL, watching ``provisioningState``) - # instead of the operation-status endpoint. ``BodyContentPolling`` polls the resource URL, so - # the poller never calls ``.../computeOperations/{id}`` (which requires ``computeOperations/read`` - # that many callers lack). Excluding ``AzureAsyncOperationPolling`` and ``LocationPolling`` - # guarantees the ``Azure-AsyncOperation``/``Location`` headers the service returns are ignored, - # while the poller still blocks until terminal and raises on a failed provisioning. - # ``_ComputeResourcePolling`` additionally tolerates the ~20s read-after-write window where the - # resource GET returns 404 "Cluster not found" right after the create is accepted. if polling is True: polling_method: PollingMethod = cast( PollingMethod, - _ComputeResourcePolling( - lro_delay, - lro_algorithms=[BodyContentPolling(), StatusCheckPolling()], - path_format_arguments=path_format_arguments, - **kwargs, - ), + _ComputeListPolling(self, resource_group_name, account_name, compute_name, lro_delay), ) elif polling is False: polling_method = cast(PollingMethod, NoPolling()) else: polling_method = polling + if cont_token: + return LROPoller[_models.Compute].from_continuation_token( + polling_method=polling_method, + continuation_token=cont_token, + client=self._client, + deserialization_callback=get_long_running_output, + ) return LROPoller[_models.Compute]( - self._client, pipeline_response, get_long_running_output, polling_method # type: ignore + self._client, raw_result, get_long_running_output, polling_method # type: ignore ) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py index f93beb795089..1219fbb51e8a 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py @@ -7,14 +7,15 @@ These verify the behavior added to work around the async compute-create handling: - the create accepts a 202 (async "Accepted") response (the generated code rejected it), - - the poller polls the compute *resource* (GET on the resource URL, watching - ``provisioningState``) via ``BodyContentPolling`` and never the operation-status endpoint - (``.../computeOperations/{id}``, which requires ``computeOperations/read``), - - the poller still blocks until a terminal state and raises on a genuine provisioning failure, and + - the poller reads status from the ``list`` API (never the operation-status endpoint + ``.../computeOperations/{id}``, which requires ``computeOperations/read``); ``list`` is used + rather than ``get`` because a GET on a just-created compute returns 404 during provisioning, + - the poller blocks until a terminal state and raises the compute's own error on failure, and - genuine non-2xx create errors still propagate. """ import json as _json import time +import typing from types import SimpleNamespace from unittest.mock import MagicMock @@ -23,15 +24,8 @@ from azure.core.credentials import AccessToken from azure.core.exceptions import HttpResponseError from azure.core.polling import NoPolling -from azure.mgmt.core.polling.arm_polling import ( - ARMPolling, - AzureAsyncOperationPolling, - BodyContentPolling, - LocationPolling, - StatusCheckPolling, -) from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient -from azure.mgmt.cognitiveservices.operations._patch import _ComputeResourcePolling +from azure.mgmt.cognitiveservices.operations._patch import _ComputeListPolling RESOURCE_URL = ( "https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000" @@ -40,16 +34,15 @@ ) ACCEPTED_BODY = {"name": "test-compute", "properties": {"provisioningState": "Accepted"}} -SUCCEEDED_BODY = { - "name": "test-compute", - "type": "Microsoft.CognitiveServices/accounts/computes", - "properties": {"provisioningState": "Succeeded"}, -} -NOT_FOUND_BODY = {"error": {"code": "NotFound", "message": "Cluster not found."}} + + +def _compute(name="test-compute", state="Succeeded", errors=None): + """A minimal stand-in for a Compute model as returned by ComputesOperations.list().""" + return SimpleNamespace(name=name, properties=SimpleNamespace(provisioning_state=state, errors=errors)) class _FakeHttpResponse: - """Minimal stand-in for an azure.core.rest HttpResponse used by the polling machinery.""" + """Minimal stand-in for the streamed create (202) response used by _create_or_update_initial.""" def __init__(self, status_code, body, method="PUT", url=RESOURCE_URL, headers=None): self.status_code = status_code @@ -60,7 +53,7 @@ def __init__(self, status_code, body, method="PUT", url=RESOURCE_URL, headers=No self.content_type = "application/json" @property - def content(self): # presence of ``content`` makes azure-core treat this as a "rest" response + def content(self): return _json.dumps(self._body).encode("utf-8") if self._body is not None else b"" def text(self, *args, **kwargs): @@ -72,6 +65,12 @@ def json(self): def read(self, *args, **kwargs): return self.content + def iter_bytes(self, *args, **kwargs): + return iter([self.content]) + + def iter_raw(self, *args, **kwargs): + return iter([self.content]) + def _pipeline_response(status_code, body, method="PUT", url=RESOURCE_URL): return SimpleNamespace(http_response=_FakeHttpResponse(status_code, body, method, url), context={}) @@ -89,59 +88,45 @@ def _make_computes(): return client.computes -def test_create_uses_resource_polling_not_computeoperations(): - """The poller is ARMPolling restricted to resource-based algorithms; the operation-status - algorithms (which would hit ``computeOperations/read``) are excluded.""" +def test_create_uses_list_polling_not_computeoperations(): + """The poller is a _ComputeListPolling that reads status from list(), never computeOperations.""" computes = _make_computes() - # A terminal 200 lets the poller finish during initialize, so no background poll is needed here. - computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(200, SUCCEEDED_BODY)) + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes.list = MagicMock(return_value=[_compute(state="Succeeded")]) poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) - assert isinstance(poller._polling_method, ARMPolling) - assert isinstance(poller._polling_method, _ComputeResourcePolling) - algorithms = [type(a) for a in poller._polling_method._lro_algorithms] - assert BodyContentPolling in algorithms - assert StatusCheckPolling in algorithms - # The whole point of the fix: never follow Azure-AsyncOperation / Location to computeOperations. - assert AzureAsyncOperationPolling not in algorithms - assert LocationPolling not in algorithms - + assert isinstance(poller._polling_method, _ComputeListPolling) result = poller.result() assert result.name == "test-compute" + # Status came from the list API (scoped to the account), not the operation-status endpoint. + computes.list.assert_called_with("rg", "acct") -def test_create_accepts_202_and_polls_resource_until_succeeded(): - """A 202 create is accepted and the poller blocks, polling the resource URL until Succeeded.""" +def test_create_accepts_202_and_polls_list_until_succeeded(): + """A 202 create is accepted and the poller blocks, polling list() until provisioningState=Succeeded.""" computes = _make_computes() computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = MagicMock(return_value=_pipeline_response(200, SUCCEEDED_BODY, method="GET")) + # First poll: still scaling; second poll: succeeded. The poller must block across both. + computes.list = MagicMock(side_effect=[[_compute(state="Scaling")], [_compute(state="Succeeded")]]) poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) result = poller.result() assert result.name == "test-compute" assert result.properties.provisioning_state == "Succeeded" - # The status was polled from the resource itself, not the operation-status endpoint. - computes._client.send_request.assert_called() - polled_url = computes._client.send_request.call_args[0][0].url - assert "computeOperations" not in polled_url - assert "/computes/test-compute" in polled_url + assert computes.list.call_count == 2 # it blocked past the non-terminal state def test_create_surfaces_provisioning_failure(): - """A create that provisions to Failed must raise with the resource's own error detail (not the - generic 'Operation returned an invalid status OK').""" + """A create that provisions to Failed must raise with the resource's own error detail.""" computes = _make_computes() - failed_body = { - "name": "test-compute", - "properties": { - "provisioningState": "Failed", - "errors": [{"code": "QuotaExceeded", "message": "exceeding subscription quota limits."}], - }, - } computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = MagicMock(return_value=_pipeline_response(200, failed_body, method="GET")) + computes.list = MagicMock( + return_value=[ + _compute(state="Failed", errors=[{"code": "QuotaExceeded", "message": "exceeding subscription quota."}]) + ] + ) poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) with pytest.raises(HttpResponseError) as exc_info: @@ -157,36 +142,37 @@ def test_create_propagates_non_2xx_error(): computes._client._pipeline.run = MagicMock( return_value=_pipeline_response(400, {"error": {"code": "Bad", "message": "bad"}}) ) - computes._client.send_request = MagicMock() + computes.list = MagicMock() with pytest.raises(HttpResponseError): computes.begin_create_or_update("rg", "acct", "test-compute", b"{}") - computes._client.send_request.assert_not_called() + computes.list.assert_not_called() def test_create_polling_false_uses_no_polling_escape_hatch(): """Callers can still opt out of blocking with ``polling=False`` (returns the accepted resource).""" computes = _make_computes() computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = MagicMock() + computes.list = MagicMock() poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling=False) assert isinstance(poller._polling_method, NoPolling) result = poller.result() assert result.name == "test-compute" - computes._client.send_request.assert_not_called() + computes.list.assert_not_called() -def test_create_tolerates_read_after_write_404_then_succeeds(): - """Right after the 202, a GET on the compute can return 404 'Cluster not found' for ~20s. The - poller must tolerate that window and keep polling instead of failing the successful create.""" +def test_create_tolerates_compute_absent_from_list_then_appears(): + """Right after the 202 the compute may not be in list() yet; the poller must keep polling until it + appears and reaches a terminal state, instead of failing the successful create.""" computes = _make_computes() computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = MagicMock( + computes.list = MagicMock( side_effect=[ - _pipeline_response(404, NOT_FOUND_BODY, method="GET"), - _pipeline_response(200, SUCCEEDED_BODY, method="GET"), + [], # not visible yet + [_compute(name="other")], # visible, but a different compute + [_compute(state="Succeeded")], # finally our compute, terminal ] ) @@ -194,17 +180,75 @@ def test_create_tolerates_read_after_write_404_then_succeeds(): result = poller.result() assert result.name == "test-compute" - assert computes._client.send_request.call_count == 2 # it kept polling past the 404 + assert computes.list.call_count == 3 -def test_create_surfaces_persistent_not_found(monkeypatch): - """If the resource never becomes queryable, the bounded grace expires and the 404 surfaces, so a - genuinely missing resource does not hang forever.""" - monkeypatch.setattr(_ComputeResourcePolling, "_NOT_FOUND_GRACE_SECONDS", -1) +def test_create_surfaces_persistent_absence(monkeypatch): + """If the compute never appears in list(), the bounded grace expires and the poller gives up rather + than hanging forever.""" + monkeypatch.setattr(_ComputeListPolling, "_NOT_FOUND_GRACE_SECONDS", -1) computes = _make_computes() computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = MagicMock(return_value=_pipeline_response(404, NOT_FOUND_BODY, method="GET")) + computes.list = MagicMock(return_value=[]) poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) with pytest.raises(HttpResponseError): poller.result() + + +def test_public_overloads_are_preserved(): + """The generated method's three public overloads (Compute, JSON, IO[bytes]) must be retained so + APIView and type checkers still see the original public signature.""" + get_overloads = getattr(typing, "get_overloads", None) + if get_overloads is None: # typing.get_overloads is Python 3.11+ + pytest.skip("typing.get_overloads requires Python 3.11+") + from azure.mgmt.cognitiveservices.operations._patch import ComputesOperations as PatchedComputesOperations + + assert len(get_overloads(PatchedComputesOperations.begin_create_or_update)) == 3 + + +def test_old_api_version_is_rejected(): + """The api-version validation guard on the generated method must be preserved: a client on an + API version older than the method's introduction must raise instead of sending an unsupported PUT.""" + client = CognitiveServicesManagementClient( + credential=_FakeCredential(), subscription_id="00000000-0000-0000-0000-000000000000", api_version="2020-01-01" + ) + client.computes._client._pipeline.run = MagicMock() + + with pytest.raises(ValueError): + client.computes.begin_create_or_update("rg", "acct", "test-compute", b"{}") + client.computes._client._pipeline.run.assert_not_called() # never sent the unsupported request + + +def test_continuation_token_does_not_send_new_create_request(): + """Resuming with a continuation_token must rebuild the poller, not re-issue the create PUT (which + would mutate the resource again).""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock() + computes.list = MagicMock(return_value=[_compute(state="Succeeded")]) + + poller = computes.begin_create_or_update( + "rg", "acct", "test-compute", b"{}", continuation_token="resume-token", polling_interval=0 + ) + poller.result() + computes._client._pipeline.run.assert_not_called() # no new create PUT on resume + + +def test_failed_initial_body_surfaces_error(): + """If the accepted initial body is already terminal-failed (a synchronous failure rather than the + usual async 202), the resource's own error detail is surfaced without any polling.""" + computes = _make_computes() + failed_body = { + "name": "test-compute", + "properties": { + "provisioningState": "Failed", + "errors": [{"code": "QuotaExceeded", "message": "exceeding subscription quota limits."}], + }, + } + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(200, failed_body)) + computes.list = MagicMock() + + with pytest.raises(HttpResponseError) as exc_info: + computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + assert "QuotaExceeded" in str(exc_info.value) + computes.list.assert_not_called() # failed synchronously; no polling needed diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py index cdff70cbca92..9f24501b4e9f 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py @@ -5,29 +5,23 @@ # -------------------------------------------------------------------------- """Async unit tests for the customized (patched) ComputesOperations.begin_create_or_update. -Mirror of the sync tests: the create accepts a 202, the poller polls the compute *resource* -(never ``computeOperations``), blocks until terminal, surfaces provisioning failures, and still -propagates genuine non-2xx create errors. +Mirror of the sync tests: the create accepts a 202, the poller reads status from the ``list`` API +(never ``computeOperations``, and not ``get`` which 404s during provisioning), blocks until terminal, +surfaces provisioning failures, and still propagates genuine non-2xx create errors. """ import json as _json import time +import typing from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from azure.core.credentials import AccessToken from azure.core.exceptions import HttpResponseError from azure.core.polling import AsyncNoPolling -from azure.mgmt.core.polling.arm_polling import ( - AzureAsyncOperationPolling, - BodyContentPolling, - LocationPolling, - StatusCheckPolling, -) -from azure.mgmt.core.polling.async_arm_polling import AsyncARMPolling from azure.mgmt.cognitiveservices.aio import CognitiveServicesManagementClient -from azure.mgmt.cognitiveservices.aio.operations._patch import _AsyncComputeResourcePolling +from azure.mgmt.cognitiveservices.aio.operations._patch import _AsyncComputeListPolling RESOURCE_URL = ( "https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000" @@ -36,16 +30,34 @@ ) ACCEPTED_BODY = {"name": "test-compute", "properties": {"provisioningState": "Accepted"}} -SUCCEEDED_BODY = { - "name": "test-compute", - "type": "Microsoft.CognitiveServices/accounts/computes", - "properties": {"provisioningState": "Succeeded"}, -} -NOT_FOUND_BODY = {"error": {"code": "NotFound", "message": "Cluster not found."}} + + +def _compute(name="test-compute", state="Succeeded", errors=None): + """A minimal stand-in for a Compute model as returned by ComputesOperations.list().""" + return SimpleNamespace(name=name, properties=SimpleNamespace(provisioning_state=state, errors=errors)) + + +class _AsyncList: + """An async-iterable stand-in for the AsyncItemPaged returned by the async list().""" + + def __init__(self, items): + self._items = list(items) + + def __aiter__(self): + async def gen(): + for item in self._items: + yield item + + return gen() + + +def _list_mock(*pages): + """A MagicMock whose successive calls return successive async-iterable pages.""" + return MagicMock(side_effect=[_AsyncList(page) for page in pages]) class _FakeAsyncHttpResponse: - """Minimal stand-in for an azure.core.rest AsyncHttpResponse used by the polling machinery.""" + """Minimal stand-in for the streamed create (202) response used by _create_or_update_initial.""" def __init__(self, status_code, body, method="PUT", url=RESOURCE_URL, headers=None): self.status_code = status_code @@ -56,7 +68,7 @@ def __init__(self, status_code, body, method="PUT", url=RESOURCE_URL, headers=No self.content_type = "application/json" @property - def content(self): # presence of ``content`` makes azure-core treat this as a "rest" response + def content(self): return _json.dumps(self._body).encode("utf-8") if self._body is not None else b"" def text(self, *args, **kwargs): @@ -68,6 +80,12 @@ def json(self): async def read(self, *args, **kwargs): return self.content + def iter_bytes(self, *args, **kwargs): + return iter([self.content]) + + def iter_raw(self, *args, **kwargs): + return iter([self.content]) + def _pipeline_response(status_code, body, method="PUT", url=RESOURCE_URL): return SimpleNamespace(http_response=_FakeAsyncHttpResponse(status_code, body, method, url), context={}) @@ -86,65 +104,50 @@ def _make_computes(): @pytest.mark.asyncio -async def test_create_uses_resource_polling_not_computeoperations(): - """The poller is AsyncARMPolling restricted to resource-based algorithms; the operation-status - algorithms (which would hit ``computeOperations/read``) are excluded.""" +async def test_create_uses_list_polling_not_computeoperations(): + """The poller is an _AsyncComputeListPolling that reads status from list(), never computeOperations.""" computes = _make_computes() - computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(200, SUCCEEDED_BODY)) + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes.list = _list_mock([_compute(state="Succeeded")]) poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) - assert isinstance(poller._polling_method, AsyncARMPolling) - assert isinstance(poller._polling_method, _AsyncComputeResourcePolling) - algorithms = [type(a) for a in poller._polling_method._lro_algorithms] - assert BodyContentPolling in algorithms - assert StatusCheckPolling in algorithms - assert AzureAsyncOperationPolling not in algorithms - assert LocationPolling not in algorithms - + assert isinstance(poller._polling_method, _AsyncComputeListPolling) result = await poller.result() assert result.name == "test-compute" + computes.list.assert_called_with("rg", "acct") @pytest.mark.asyncio -async def test_create_accepts_202_and_polls_resource_until_succeeded(): - """A 202 create is accepted and the poller blocks, polling the resource URL until Succeeded.""" +async def test_create_accepts_202_and_polls_list_until_succeeded(): + """A 202 create is accepted and the poller blocks, polling list() until provisioningState=Succeeded.""" computes = _make_computes() computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = AsyncMock(return_value=_pipeline_response(200, SUCCEEDED_BODY, method="GET")) + computes.list = _list_mock([_compute(state="Scaling")], [_compute(state="Succeeded")]) poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) result = await poller.result() assert result.name == "test-compute" assert result.properties.provisioning_state == "Succeeded" - computes._client.send_request.assert_called() - polled_url = computes._client.send_request.call_args[0][0].url - assert "computeOperations" not in polled_url - assert "/computes/test-compute" in polled_url + assert computes.list.call_count == 2 # it blocked past the non-terminal state @pytest.mark.asyncio async def test_create_surfaces_provisioning_failure(): - """A create that provisions to Failed must raise with the resource's own error detail (not the - generic 'Operation returned an invalid status OK').""" + """A create that provisions to Failed must raise with the resource's own error detail.""" computes = _make_computes() - failed_body = { - "name": "test-compute", - "properties": { - "provisioningState": "Failed", - "errors": [{"code": "QuotaExceeded", "message": "exceeding subscription quota limits."}], - }, - } computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = AsyncMock(return_value=_pipeline_response(200, failed_body, method="GET")) + computes.list = _list_mock( + [_compute(state="Failed", errors=[{"code": "QuotaExceeded", "message": "exceeding subscription quota."}])] + ) poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) with pytest.raises(HttpResponseError) as exc_info: await poller.result() message = str(exc_info.value) - assert "QuotaExceeded" in message # the real reason is surfaced - assert "invalid status" not in message # not azure-core's generic fallback message + assert "QuotaExceeded" in message + assert "invalid status" not in message @pytest.mark.asyncio @@ -154,11 +157,11 @@ async def test_create_propagates_non_2xx_error(): computes._client._pipeline.run = AsyncMock( return_value=_pipeline_response(400, {"error": {"code": "Bad", "message": "bad"}}) ) - computes._client.send_request = AsyncMock() + computes.list = MagicMock() with pytest.raises(HttpResponseError): await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}") - computes._client.send_request.assert_not_called() + computes.list.assert_not_called() @pytest.mark.asyncio @@ -166,45 +169,99 @@ async def test_create_polling_false_uses_no_polling_escape_hatch(): """Callers can still opt out of blocking with ``polling=False`` (returns the accepted resource).""" computes = _make_computes() computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = AsyncMock() + computes.list = MagicMock() poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling=False) assert isinstance(poller._polling_method, AsyncNoPolling) result = await poller.result() assert result.name == "test-compute" - computes._client.send_request.assert_not_called() + computes.list.assert_not_called() @pytest.mark.asyncio -async def test_create_tolerates_read_after_write_404_then_succeeds(): - """Right after the 202, a GET on the compute can return 404 'Cluster not found' for ~20s. The - poller must tolerate that window and keep polling instead of failing the successful create.""" +async def test_create_tolerates_compute_absent_from_list_then_appears(): + """Right after the 202 the compute may not be in list() yet; the poller must keep polling until it + appears and reaches a terminal state, instead of failing the successful create.""" computes = _make_computes() computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = AsyncMock( - side_effect=[ - _pipeline_response(404, NOT_FOUND_BODY, method="GET"), - _pipeline_response(200, SUCCEEDED_BODY, method="GET"), - ] - ) + computes.list = _list_mock([], [_compute(name="other")], [_compute(state="Succeeded")]) poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) result = await poller.result() assert result.name == "test-compute" - assert computes._client.send_request.call_count == 2 # it kept polling past the 404 + assert computes.list.call_count == 3 @pytest.mark.asyncio -async def test_create_surfaces_persistent_not_found(monkeypatch): - """If the resource never becomes queryable, the bounded grace expires and the 404 surfaces, so a - genuinely missing resource does not hang forever.""" - monkeypatch.setattr(_AsyncComputeResourcePolling, "_NOT_FOUND_GRACE_SECONDS", -1) +async def test_create_surfaces_persistent_absence(monkeypatch): + """If the compute never appears in list(), the bounded grace expires and the poller gives up rather + than hanging forever.""" + monkeypatch.setattr(_AsyncComputeListPolling, "_NOT_FOUND_GRACE_SECONDS", -1) computes = _make_computes() computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) - computes._client.send_request = AsyncMock(return_value=_pipeline_response(404, NOT_FOUND_BODY, method="GET")) + computes.list = MagicMock(side_effect=lambda *a, **k: _AsyncList([])) poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) with pytest.raises(HttpResponseError): await poller.result() + + +def test_public_overloads_are_preserved(): + """The generated async method's three public overloads must be retained for APIView/type checkers.""" + get_overloads = getattr(typing, "get_overloads", None) + if get_overloads is None: # typing.get_overloads is Python 3.11+ + pytest.skip("typing.get_overloads requires Python 3.11+") + from azure.mgmt.cognitiveservices.aio.operations._patch import ComputesOperations as PatchedComputesOperations + + assert len(get_overloads(PatchedComputesOperations.begin_create_or_update)) == 3 + + +@pytest.mark.asyncio +async def test_old_api_version_is_rejected(): + """The api-version validation guard on the generated async method must be preserved.""" + client = CognitiveServicesManagementClient( + credential=_FakeAsyncCredential(), + subscription_id="00000000-0000-0000-0000-000000000000", + api_version="2020-01-01", + ) + client.computes._client._pipeline.run = AsyncMock() + + with pytest.raises(ValueError): + await client.computes.begin_create_or_update("rg", "acct", "test-compute", b"{}") + client.computes._client._pipeline.run.assert_not_called() + + +@pytest.mark.asyncio +async def test_continuation_token_does_not_send_new_create_request(): + """Resuming with a continuation_token must rebuild the poller, not re-issue the create PUT.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock() + computes.list = _list_mock([_compute(state="Succeeded")]) + + poller = await computes.begin_create_or_update( + "rg", "acct", "test-compute", b"{}", continuation_token="resume-token", polling_interval=0 + ) + await poller.result() + computes._client._pipeline.run.assert_not_called() + + +@pytest.mark.asyncio +async def test_failed_initial_body_surfaces_error(): + """A synchronously terminal-failed initial body surfaces the resource's own error, without polling.""" + computes = _make_computes() + failed_body = { + "name": "test-compute", + "properties": { + "provisioningState": "Failed", + "errors": [{"code": "QuotaExceeded", "message": "exceeding subscription quota limits."}], + }, + } + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(200, failed_body)) + computes.list = MagicMock() + + with pytest.raises(HttpResponseError) as exc_info: + await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + assert "QuotaExceeded" in str(exc_info.value) + computes.list.assert_not_called() From cff382e17dc1efae4e0cabd0818ec3773ce3d901 Mon Sep 17 00:00:00 2001 From: Swarup Chanda Date: Fri, 17 Jul 2026 17:36:55 +0530 Subject: [PATCH 3/5] [cognitiveservices] Regenerate api.md and api.metadata.yml for the patched ComputesOperations --- sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.md | 4 ++-- .../azure-mgmt-cognitiveservices/api.metadata.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.md b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.md index f0a867432302..3d96f0f62ee4 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.md +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.md @@ -1179,7 +1179,7 @@ namespace azure.mgmt.cognitiveservices.aio.operations ) -> ComputeOperationStatus: ... - class azure.mgmt.cognitiveservices.aio.operations.ComputesOperations: + class azure.mgmt.cognitiveservices.aio.operations.ComputesOperations(_ComputesOperationsGenerated): def __init__( self, @@ -9523,7 +9523,7 @@ namespace azure.mgmt.cognitiveservices.operations ) -> ComputeOperationStatus: ... - class azure.mgmt.cognitiveservices.operations.ComputesOperations: + class azure.mgmt.cognitiveservices.operations.ComputesOperations(_ComputesOperationsGenerated): def __init__( self, diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.metadata.yml b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.metadata.yml index 8f9f4c53ce79..755c109a7d04 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.metadata.yml +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/api.metadata.yml @@ -1,3 +1,3 @@ -apiMdSha256: 0012caaf8366e50cf017e38419ca00952dd54561af6c897b4c16b53ba2c122b9 +apiMdSha256: d9ccc3bf57442333ddf38636daad902c89fe92f3b76fc124e3ded208f3ba7af3 parserVersion: 0.3.28 -pythonVersion: 3.13.13 +pythonVersion: 3.12.10 From 1fed83b489b078b8545a5224ace91a0b0e04d839 Mon Sep 17 00:00:00 2001 From: Swarup Chanda Date: Fri, 17 Jul 2026 18:25:05 +0530 Subject: [PATCH 4/5] Address Copilot review: honor cls hook, preserve Canceled status, validate continuation_token - Apply caller-supplied cls to the final resource in the list-based poller. - Preserve terminal Canceled status instead of collapsing it to Failed. - Encode/decode continuation_token so resume uses the token as source of truth and raises on a malformed token. - Add sync+async tests for all three (30 tests total). --- .../aio/operations/_patch.py | 33 ++++++++- .../cognitiveservices/operations/_patch.py | 50 ++++++++++++- .../tests/test_computes_operations_patch.py | 71 +++++++++++++++++- .../test_computes_operations_patch_async.py | 74 ++++++++++++++++++- 4 files changed, 218 insertions(+), 10 deletions(-) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py index 2fc764d672c4..4d7df1db9a4d 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py @@ -33,7 +33,15 @@ from ... import models as _models from ..._utils.model_base import SdkJSONEncoder, _deserialize, _failsafe_deserialize from ..._validation import api_version_validation -from ...operations._patch import _TERMINAL_FAILED_STATES, _TERMINAL_SUCCESS_STATES, _provisioning_error, _state_of +from ...operations._patch import ( + _CANCELED_STATES, + _TERMINAL_FAILED_STATES, + _TERMINAL_SUCCESS_STATES, + _decode_continuation_token, + _encode_continuation_token, + _provisioning_error, + _state_of, +) from ...operations._operations import build_computes_create_or_update_request from ._operations import ComputesOperations as _ComputesOperationsGenerated @@ -65,12 +73,14 @@ def __init__( account_name: str, compute_name: str, interval: float, + cls: Any = None, ) -> None: self._operations = operations self._resource_group_name = resource_group_name self._account_name = account_name self._compute_name = compute_name self._interval = interval + self._cls = cls self._status = "InProgress" self._resource: Optional[_models.Compute] = None self._first_missing_at: Optional[float] = None @@ -105,7 +115,9 @@ async def run(self) -> None: if state in _TERMINAL_SUCCESS_STATES: self._status = "Succeeded" elif state in _TERMINAL_FAILED_STATES: - self._status = "Failed" + # Preserve the service's terminal status (``Canceled`` vs ``Failed``) so callers can + # inspect ``poller.status()`` correctly, while still raising the provisioning error. + self._status = "Canceled" if state in _CANCELED_STATES else "Failed" raise _provisioning_error(compute) if not self.finished(): await asyncio.sleep(self._interval) @@ -117,13 +129,21 @@ def finished(self) -> bool: return self._status in ("Succeeded", "Failed", "Canceled") def resource(self) -> Optional[_models.Compute]: + # Apply the caller's ``cls`` hook (if any) to the final resource, preserving the generated + # method's ``cls=`` contract. There is no final pipeline response for the list-based read-back, + # so ``None`` is passed for that positional argument. + if self._resource is not None and self._cls is not None: + return self._cls(None, self._resource, {}) return self._resource def get_continuation_token(self) -> str: - return self._compute_name + return _encode_continuation_token(self._resource_group_name, self._account_name, self._compute_name) @classmethod def from_continuation_token(cls, continuation_token: str, **kwargs: Any): + # Validate the token is well-formed; the polling scope it encodes is applied when the poller is + # (re)constructed in ``begin_create_or_update``. + _decode_continuation_token(continuation_token) client = kwargs["client"] deserialization_callback = kwargs["deserialization_callback"] return client, None, deserialization_callback @@ -359,6 +379,11 @@ async def begin_create_or_update( initial_body = None if isinstance(initial_body, MutableMapping) and _state_of(initial_body) in _TERMINAL_FAILED_STATES: raise _provisioning_error(initial_body) + else: + # Opaque-token contract: the token is the source of truth for which operation to resume, so + # decode the polling scope from it (raising on a malformed token) rather than trusting the + # method arguments, and do not re-issue the create. + resource_group_name, account_name, compute_name = _decode_continuation_token(cont_token) kwargs.pop("error_map", None) def get_long_running_output(pipeline_response): @@ -371,7 +396,7 @@ def get_long_running_output(pipeline_response): if polling is True: polling_method: AsyncPollingMethod = cast( AsyncPollingMethod, - _AsyncComputeListPolling(self, resource_group_name, account_name, compute_name, lro_delay), + _AsyncComputeListPolling(self, resource_group_name, account_name, compute_name, lro_delay, cls), ) elif polling is False: polling_method = cast(AsyncPollingMethod, AsyncNoPolling()) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py index 55948b11be22..d0c9cf660f3b 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py @@ -7,6 +7,7 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ +import base64 import json import time from collections.abc import MutableMapping @@ -39,6 +40,32 @@ _TERMINAL_FAILED_STATES = frozenset({"failed", "canceled", "cancelled"}) _TERMINAL_SUCCESS_STATES = frozenset({"succeeded"}) +_CANCELED_STATES = frozenset({"canceled", "cancelled"}) + + +def _encode_continuation_token(resource_group_name: str, account_name: str, compute_name: str) -> str: + """Encode the polling scope (resource group, account, compute) into an opaque continuation token.""" + payload = json.dumps( + { + "resource_group_name": resource_group_name, + "account_name": account_name, + "compute_name": compute_name, + } + ) + return base64.urlsafe_b64encode(payload.encode("utf-8")).decode("ascii") + + +def _decode_continuation_token(continuation_token: str): + """Decode a token produced by :func:`_encode_continuation_token` into its polling scope. + + :raises ValueError: if the token is malformed, so a bad or unrelated token fails loudly instead of + silently resuming the wrong operation. + """ + try: + payload = json.loads(base64.urlsafe_b64decode(continuation_token.encode("ascii")).decode("utf-8")) + return payload["resource_group_name"], payload["account_name"], payload["compute_name"] + except Exception as exc: # pylint: disable=broad-except + raise ValueError("Invalid continuation_token for compute begin_create_or_update.") from exc def _state_of(compute: Any) -> str: @@ -107,12 +134,14 @@ def __init__( account_name: str, compute_name: str, interval: float, + cls: Any = None, ) -> None: self._operations = operations self._resource_group_name = resource_group_name self._account_name = account_name self._compute_name = compute_name self._interval = interval + self._cls = cls self._status = "InProgress" self._resource: Optional[_models.Compute] = None self._first_missing_at: Optional[float] = None @@ -147,7 +176,9 @@ def run(self) -> None: if state in _TERMINAL_SUCCESS_STATES: self._status = "Succeeded" elif state in _TERMINAL_FAILED_STATES: - self._status = "Failed" + # Preserve the service's terminal status (``Canceled`` vs ``Failed``) so callers can + # inspect ``poller.status()`` correctly, while still raising the provisioning error. + self._status = "Canceled" if state in _CANCELED_STATES else "Failed" raise _provisioning_error(compute) if not self.finished(): time.sleep(self._interval) @@ -159,13 +190,21 @@ def finished(self) -> bool: return self._status in ("Succeeded", "Failed", "Canceled") def resource(self) -> Optional[_models.Compute]: + # Apply the caller's ``cls`` hook (if any) to the final resource, preserving the generated + # method's ``cls=`` contract. There is no final pipeline response for the list-based read-back, + # so ``None`` is passed for that positional argument. + if self._resource is not None and self._cls is not None: + return self._cls(None, self._resource, {}) return self._resource def get_continuation_token(self) -> str: - return self._compute_name + return _encode_continuation_token(self._resource_group_name, self._account_name, self._compute_name) @classmethod def from_continuation_token(cls, continuation_token: str, **kwargs: Any): + # Validate the token is well-formed; the polling scope it encodes is applied when the poller is + # (re)constructed in ``begin_create_or_update``. + _decode_continuation_token(continuation_token) client = kwargs["client"] deserialization_callback = kwargs["deserialization_callback"] return client, None, deserialization_callback @@ -401,6 +440,11 @@ def begin_create_or_update( initial_body = None if isinstance(initial_body, MutableMapping) and _state_of(initial_body) in _TERMINAL_FAILED_STATES: raise _provisioning_error(initial_body) + else: + # Opaque-token contract: the token is the source of truth for which operation to resume, so + # decode the polling scope from it (raising on a malformed token) rather than trusting the + # method arguments, and do not re-issue the create. + resource_group_name, account_name, compute_name = _decode_continuation_token(cont_token) kwargs.pop("error_map", None) def get_long_running_output(pipeline_response): @@ -413,7 +457,7 @@ def get_long_running_output(pipeline_response): if polling is True: polling_method: PollingMethod = cast( PollingMethod, - _ComputeListPolling(self, resource_group_name, account_name, compute_name, lro_delay), + _ComputeListPolling(self, resource_group_name, account_name, compute_name, lro_delay, cls), ) elif polling is False: polling_method = cast(PollingMethod, NoPolling()) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py index 1219fbb51e8a..40e939fac479 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch.py @@ -25,7 +25,7 @@ from azure.core.exceptions import HttpResponseError from azure.core.polling import NoPolling from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient -from azure.mgmt.cognitiveservices.operations._patch import _ComputeListPolling +from azure.mgmt.cognitiveservices.operations._patch import _ComputeListPolling, _encode_continuation_token RESOURCE_URL = ( "https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000" @@ -227,13 +227,80 @@ def test_continuation_token_does_not_send_new_create_request(): computes._client._pipeline.run = MagicMock() computes.list = MagicMock(return_value=[_compute(state="Succeeded")]) + token = _encode_continuation_token("rg", "acct", "test-compute") poller = computes.begin_create_or_update( - "rg", "acct", "test-compute", b"{}", continuation_token="resume-token", polling_interval=0 + "rg", "acct", "test-compute", b"{}", continuation_token=token, polling_interval=0 ) poller.result() computes._client._pipeline.run.assert_not_called() # no new create PUT on resume +def test_continuation_token_is_source_of_truth_for_scope(): + """The token (not the method arguments) decides which compute is resumed, so a token that encodes a + different scope drives the list() poll to that scope.""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock() + computes.list = MagicMock(return_value=[_compute(name="real-compute", state="Succeeded")]) + + token = _encode_continuation_token("real-rg", "real-acct", "real-compute") + poller = computes.begin_create_or_update( + "ignored-rg", "ignored-acct", "ignored-compute", b"{}", continuation_token=token, polling_interval=0 + ) + result = poller.result() + + assert result.name == "real-compute" + computes.list.assert_called_with("real-rg", "real-acct") # scope came from the token, not the args + + +def test_continuation_token_malformed_raises(): + """A malformed/unrelated continuation_token must fail loudly instead of silently resuming the wrong + (current-argument) operation.""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock() + computes.list = MagicMock() + + with pytest.raises(ValueError): + computes.begin_create_or_update( + "rg", "acct", "test-compute", b"{}", continuation_token="not-a-valid-token", polling_interval=0 + ) + computes._client._pipeline.run.assert_not_called() + computes.list.assert_not_called() + + +def test_create_preserves_canceled_status(): + """A compute that provisions to Canceled must be reported as ``Canceled`` (not collapsed to + ``Failed``) so callers can inspect poller.status(), while still raising.""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes.list = MagicMock(return_value=[_compute(state="Canceled")]) + + poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + with pytest.raises(HttpResponseError): + poller.result() + assert poller.status() == "Canceled" + + +def test_create_applies_cls_hook(): + """A caller-supplied ``cls`` must be applied to the final resource, preserving the generated + method's ``cls=`` contract.""" + computes = _make_computes() + computes._client._pipeline.run = MagicMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes.list = MagicMock(return_value=[_compute(state="Succeeded")]) + + sentinel = object() + calls = [] + + def cls(pipeline_response, deserialized, headers): + calls.append((pipeline_response, deserialized, headers)) + return sentinel + + poller = computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", cls=cls, polling_interval=0) + result = poller.result() + + assert result is sentinel # the cls hook's return value is what the caller receives + assert calls and calls[0][1].name == "test-compute" # cls received the deserialized compute + + def test_failed_initial_body_surfaces_error(): """If the accepted initial body is already terminal-failed (a synchronous failure rather than the usual async 202), the resource's own error detail is surfaced without any polling.""" diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py index 9f24501b4e9f..c2685c8fdf3d 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/tests/test_computes_operations_patch_async.py @@ -22,6 +22,7 @@ from azure.core.polling import AsyncNoPolling from azure.mgmt.cognitiveservices.aio import CognitiveServicesManagementClient from azure.mgmt.cognitiveservices.aio.operations._patch import _AsyncComputeListPolling +from azure.mgmt.cognitiveservices.operations._patch import _encode_continuation_token RESOURCE_URL = ( "https://management.azure.com/subscriptions/00000000-0000-0000-0000-000000000000" @@ -240,13 +241,84 @@ async def test_continuation_token_does_not_send_new_create_request(): computes._client._pipeline.run = AsyncMock() computes.list = _list_mock([_compute(state="Succeeded")]) + token = _encode_continuation_token("rg", "acct", "test-compute") poller = await computes.begin_create_or_update( - "rg", "acct", "test-compute", b"{}", continuation_token="resume-token", polling_interval=0 + "rg", "acct", "test-compute", b"{}", continuation_token=token, polling_interval=0 ) await poller.result() computes._client._pipeline.run.assert_not_called() +@pytest.mark.asyncio +async def test_continuation_token_is_source_of_truth_for_scope(): + """The token (not the method arguments) decides which compute is resumed, so a token that encodes a + different scope drives the list() poll to that scope.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock() + computes.list = _list_mock([_compute(name="real-compute", state="Succeeded")]) + + token = _encode_continuation_token("real-rg", "real-acct", "real-compute") + poller = await computes.begin_create_or_update( + "ignored-rg", "ignored-acct", "ignored-compute", b"{}", continuation_token=token, polling_interval=0 + ) + result = await poller.result() + + assert result.name == "real-compute" + computes.list.assert_called_with("real-rg", "real-acct") # scope came from the token, not the args + + +@pytest.mark.asyncio +async def test_continuation_token_malformed_raises(): + """A malformed/unrelated continuation_token must fail loudly instead of silently resuming the wrong + (current-argument) operation.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock() + computes.list = MagicMock() + + with pytest.raises(ValueError): + await computes.begin_create_or_update( + "rg", "acct", "test-compute", b"{}", continuation_token="not-a-valid-token", polling_interval=0 + ) + computes._client._pipeline.run.assert_not_called() + computes.list.assert_not_called() + + +@pytest.mark.asyncio +async def test_create_preserves_canceled_status(): + """A compute that provisions to Canceled must be reported as ``Canceled`` (not collapsed to + ``Failed``) so callers can inspect poller.status(), while still raising.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes.list = _list_mock([_compute(state="Canceled")]) + + poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", polling_interval=0) + with pytest.raises(HttpResponseError): + await poller.result() + assert poller.status() == "Canceled" + + +@pytest.mark.asyncio +async def test_create_applies_cls_hook(): + """A caller-supplied ``cls`` must be applied to the final resource, preserving the generated + method's ``cls=`` contract.""" + computes = _make_computes() + computes._client._pipeline.run = AsyncMock(return_value=_pipeline_response(202, ACCEPTED_BODY)) + computes.list = _list_mock([_compute(state="Succeeded")]) + + sentinel = object() + calls = [] + + def cls(pipeline_response, deserialized, headers): + calls.append((pipeline_response, deserialized, headers)) + return sentinel + + poller = await computes.begin_create_or_update("rg", "acct", "test-compute", b"{}", cls=cls, polling_interval=0) + result = await poller.result() + + assert result is sentinel # the cls hook's return value is what the caller receives + assert calls and calls[0][1].name == "test-compute" # cls received the deserialized compute + + @pytest.mark.asyncio async def test_failed_initial_body_surfaces_error(): """A synchronously terminal-failed initial body surfaces the resource's own error, without polling.""" From 38e3c4e9b48b22d96e3ecd189166f2c05298a2c4 Mon Sep 17 00:00:00 2001 From: Swarup Chanda Date: Fri, 17 Jul 2026 18:51:05 +0530 Subject: [PATCH 5/5] Make patched ComputesOperations pylint-clean (docstrings, attr init) --- .../aio/operations/_patch.py | 3 +- .../cognitiveservices/operations/_patch.py | 30 +++++++++++++++++-- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py index 4d7df1db9a4d..68ce2b4f5d57 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/aio/operations/_patch.py @@ -7,7 +7,7 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ -import asyncio +import asyncio # pylint: disable=do-not-import-asyncio import json import time from collections.abc import MutableMapping @@ -84,6 +84,7 @@ def __init__( self._status = "InProgress" self._resource: Optional[_models.Compute] = None self._first_missing_at: Optional[float] = None + self._initial_response: Any = None def initialize(self, client: Any, initial_response: Any, deserialization_callback: Any) -> None: self._initial_response = initial_response diff --git a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py index d0c9cf660f3b..488f910660e4 100644 --- a/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py +++ b/sdk/cognitiveservices/azure-mgmt-cognitiveservices/azure/mgmt/cognitiveservices/operations/_patch.py @@ -44,7 +44,17 @@ def _encode_continuation_token(resource_group_name: str, account_name: str, compute_name: str) -> str: - """Encode the polling scope (resource group, account, compute) into an opaque continuation token.""" + """Encode the polling scope (resource group, account, compute) into an opaque continuation token. + + :param resource_group_name: The resource group name of the compute being polled. + :type resource_group_name: str + :param account_name: The Cognitive Services account name of the compute being polled. + :type account_name: str + :param compute_name: The name of the compute being polled. + :type compute_name: str + :return: A url-safe, opaque continuation token encoding the polling scope. + :rtype: str + """ payload = json.dumps( { "resource_group_name": resource_group_name, @@ -58,6 +68,10 @@ def _encode_continuation_token(resource_group_name: str, account_name: str, comp def _decode_continuation_token(continuation_token: str): """Decode a token produced by :func:`_encode_continuation_token` into its polling scope. + :param continuation_token: A token previously produced by :func:`_encode_continuation_token`. + :type continuation_token: str + :return: The decoded scope as ``(resource_group_name, account_name, compute_name)``. + :rtype: tuple[str, str, str] :raises ValueError: if the token is malformed, so a bad or unrelated token fails loudly instead of silently resuming the wrong operation. """ @@ -69,7 +83,13 @@ def _decode_continuation_token(continuation_token: str): def _state_of(compute: Any) -> str: - """Return a compute's ``provisioningState`` as a lowercase string (handles enum or str).""" + """Return a compute's ``provisioningState`` as a lowercase string (handles enum or str). + + :param compute: A compute object (model or mapping) as returned by ``ComputesOperations.list``. + :type compute: any + :return: The compute's ``provisioningState`` lowercased, or an empty string if it is unavailable. + :rtype: str + """ props = getattr(compute, "properties", None) state = getattr(props, "provisioning_state", None) if props is not None else None if state is None and isinstance(compute, MutableMapping): @@ -85,6 +105,11 @@ def _provisioning_error(compute: Any) -> HttpResponseError: Because the poller reads status from ``list`` (not the operation-status endpoint), the failure detail lives in the compute's ``properties.errors``. Surface that instead of azure-core's generic ``Operation failed or canceled`` / ``invalid status`` messages. + + :param compute: A compute whose provisioning reached a failed/canceled terminal state. + :type compute: any + :return: An error describing the failure using the compute's own ``properties.errors`` detail. + :rtype: ~azure.core.exceptions.HttpResponseError """ props = getattr(compute, "properties", None) errors = getattr(props, "errors", None) if props is not None else None @@ -145,6 +170,7 @@ def __init__( self._status = "InProgress" self._resource: Optional[_models.Compute] = None self._first_missing_at: Optional[float] = None + self._initial_response: Any = None def initialize(self, client: Any, initial_response: Any, deserialization_callback: Any) -> None: self._initial_response = initial_response