inject custom policy to raise exception on quota limit exceeded (Azure#25631)

Author: kristapratico

sdk/textanalytics/azure-ai-textanalytics/CHANGELOG.md

@@ -25,6 +25,10 @@ The version of this client library defaults to the API version `2022-05-01`.
 - Renamed `SingleCategoryClassifyAction` to `SingleLabelClassifyAction`
 - Renamed `MultiCategoryClassifyAction` to `MultiLabelClassifyAction`.
 
+### Bugs Fixed
+
+- A `HttpResponseError` will be immediately raised when the call quota volume is exceeded in a `F0` tier Language resource.
+
 ### Other Changes
 
 - Python 3.6 is no longer supported. Please use Python version 3.7 or later. For more details, see [Azure SDK for Python version support policy](https://github.com/Azure/azure-sdk-for-python/wiki/Azure-SDKs-Python-version-support-policy).

sdk/textanalytics/azure-ai-textanalytics/azure/ai/textanalytics/_base_client.py

@@ -9,7 +9,7 @@
 from azure.core.pipeline.policies import AzureKeyCredentialPolicy, HttpLoggingPolicy
 from azure.core.credentials import AzureKeyCredential, TokenCredential
 from ._generated import TextAnalyticsClient as _TextAnalyticsClient
-from ._policies import TextAnalyticsResponseHookPolicy
+from ._policies import TextAnalyticsResponseHookPolicy, QuotaExceededPolicy
 from ._user_agent import USER_AGENT
 from ._version import DEFAULT_API_VERSION
 
@@ -84,6 +84,7 @@ def __init__(
             authentication_policy=kwargs.pop("authentication_policy", _authentication_policy(credential)),
             custom_hook_policy=kwargs.pop("custom_hook_policy", TextAnalyticsResponseHookPolicy(**kwargs)),
             http_logging_policy=kwargs.pop("http_logging_policy", http_logging_policy),
+            per_retry_policies=kwargs.get("per_retry_policies", QuotaExceededPolicy()),
             **kwargs
         )
 

sdk/textanalytics/azure-ai-textanalytics/azure/ai/textanalytics/_policies.py

@@ -5,6 +5,7 @@
 
 from azure.core.pipeline.policies import ContentDecodePolicy
 from azure.core.pipeline.policies import SansIOHTTPPolicy
+from azure.core.exceptions import HttpResponseError
 from ._models import TextDocumentBatchStatistics
 from ._lro import _FINISHED
 
@@ -43,3 +44,23 @@ def on_response(self, request, response):
                 response.model_version = model_version
             response.raw_response = data
             self._response_callback(response)
+
+
+class QuotaExceededPolicy(SansIOHTTPPolicy):
+    """Raises an exception immediately when the call quota volume has been exceeded in a F0
+    tier language resource. This is to avoid waiting the Retry-After time returned in
+    the response.
+    """
+
+    def on_response(self, request, response):
+        """Is executed after the request comes back from the policy.
+
+        :param request: Request to be modified after returning from the policy.
+        :type request: ~azure.core.pipeline.PipelineRequest
+        :param response: Pipeline response object
+        :type response: ~azure.core.pipeline.PipelineResponse
+        """
+        http_response = response.http_response
+        if http_response.status_code == 403 and \
+                "Out of call volume quota for TextAnalytics F0 pricing tier" in http_response.text():
+            raise HttpResponseError(http_response.text(), response=http_response)

sdk/textanalytics/azure-ai-textanalytics/azure/ai/textanalytics/aio/_base_client_async.py

@@ -7,7 +7,7 @@
 from azure.core.credentials_async import AsyncTokenCredential
 from azure.core.pipeline.policies import AzureKeyCredentialPolicy, HttpLoggingPolicy
 from .._generated.aio import TextAnalyticsClient as _TextAnalyticsClient
-from .._policies import TextAnalyticsResponseHookPolicy
+from .._policies import TextAnalyticsResponseHookPolicy, QuotaExceededPolicy
 from .._user_agent import USER_AGENT
 from .._version import DEFAULT_API_VERSION
 
@@ -71,6 +71,7 @@ def __init__(
             authentication_policy=kwargs.pop("authentication_policy", _authentication_policy(credential)),
             custom_hook_policy=kwargs.pop("custom_hook_policy", TextAnalyticsResponseHookPolicy(**kwargs)),
             http_logging_policy=kwargs.pop("http_logging_policy", http_logging_policy),
+            per_retry_policies=kwargs.get("per_retry_policies", QuotaExceededPolicy()),
             **kwargs
         )
 

sdk/textanalytics/azure-ai-textanalytics/tests/test_analyze_sentiment.py

@@ -7,6 +7,7 @@
 import platform
 import functools
 import json
+from unittest import mock
 from azure.core.exceptions import HttpResponseError, ClientAuthenticationError
 from azure.core.credentials import AzureKeyCredential
 from testcase import TextAnalyticsTest, TextAnalyticsPreparer
@@ -854,3 +855,25 @@ def test_sentiment_multiapi_validate_args_v3_0(self, **kwargs):
         with pytest.raises(ValueError) as e:
             res = client.analyze_sentiment(["I'm tired"], show_opinion_mining=True, disable_service_logs=True, string_index_type="UnicodeCodePoint")
         assert str(e.value) == "'show_opinion_mining' is only available for API version v3.1 and up.\n'disable_service_logs' is only available for API version v3.1 and up.\n'string_index_type' is only available for API version v3.1 and up.\n"
+
+    @TextAnalyticsPreparer()
+    def test_mock_quota_exceeded(self, **kwargs):
+        textanalytics_test_endpoint = kwargs.pop("textanalytics_test_endpoint")
+        textanalytics_test_api_key = kwargs.pop("textanalytics_test_api_key")
+        response = mock.Mock(
+            status_code=403,
+            headers={"Retry-After": 186688, "Content-Type": "application/json"},
+            reason="Bad Request"
+        )
+        response.text = lambda encoding=None: json.dumps(
+            {"error": {"code": "403", "message": "Out of call volume quota for TextAnalytics F0 pricing tier. Please retry after 15 days. To increase your call volume switch to a paid tier."}}
+        )
+        response.content_type = "application/json"
+        transport = mock.Mock(send=lambda request, **kwargs: response)
+
+        client = TextAnalyticsClient(textanalytics_test_endpoint, AzureKeyCredential(textanalytics_test_api_key), transport=transport)
+
+        with pytest.raises(HttpResponseError) as e:
+            result = client.analyze_sentiment(["I'm tired"])
+        assert e.value.status_code == 403
+        assert e.value.error.message == 'Out of call volume quota for TextAnalytics F0 pricing tier. Please retry after 15 days. To increase your call volume switch to a paid tier.'

sdk/textanalytics/azure-ai-textanalytics/tests/test_analyze_sentiment_async.py

@@ -7,6 +7,10 @@
 import platform
 import functools
 import json
+import sys
+import asyncio
+import functools
+from unittest import mock
 from azure.core.exceptions import HttpResponseError, ClientAuthenticationError
 from azure.core.credentials import AzureKeyCredential
 from azure.ai.textanalytics.aio import TextAnalyticsClient
@@ -25,6 +29,36 @@
 # pre-apply the client_cls positional argument so it needn't be explicitly passed below
 TextAnalyticsClientPreparer = functools.partial(_TextAnalyticsClientPreparer, TextAnalyticsClient)
 
+def get_completed_future(result=None):
+    future = asyncio.Future()
+    future.set_result(result)
+    return future
+
+
+def wrap_in_future(fn):
+    """Return a completed Future whose result is the return of fn.
+    Added to simplify using unittest.Mock in async code. Python 3.8's AsyncMock would be preferable.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        result = fn(*args, **kwargs)
+        return get_completed_future(result)
+    return wrapper
+
+
+class AsyncMockTransport(mock.MagicMock):
+    """Mock with do-nothing aenter/exit for mocking async transport.
+    This is unnecessary on 3.8+, where MagicMocks implement aenter/exit.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if sys.version_info < (3, 8):
+            self.__aenter__ = mock.Mock(return_value=get_completed_future())
+            self.__aexit__ = mock.Mock(return_value=get_completed_future())
+
 
 class TestAnalyzeSentiment(TextAnalyticsTest):
 
@@ -865,3 +899,26 @@ async def test_sentiment_multiapi_validate_args_v3_0(self, **kwargs):
         with pytest.raises(ValueError) as e:
             res = await client.analyze_sentiment(["I'm tired"], show_opinion_mining=True, disable_service_logs=True, string_index_type="UnicodeCodePoint")
         assert str(e.value) == "'show_opinion_mining' is only available for API version v3.1 and up.\n'disable_service_logs' is only available for API version v3.1 and up.\n'string_index_type' is only available for API version v3.1 and up.\n"
+
+    @TextAnalyticsPreparer()
+    async def test_mock_quota_exceeded(self, **kwargs):
+        textanalytics_test_endpoint = kwargs.pop("textanalytics_test_endpoint")
+        textanalytics_test_api_key = kwargs.pop("textanalytics_test_api_key")
+
+        response = mock.Mock(
+            status_code=403,
+            headers={"Retry-After": 186688, "Content-Type": "application/json"},
+            reason="Bad Request"
+        )
+        response.text = lambda encoding=None: json.dumps(
+            {"error": {"code": "403", "message": "Out of call volume quota for TextAnalytics F0 pricing tier. Please retry after 15 days. To increase your call volume switch to a paid tier."}}
+        )
+        response.content_type = "application/json"
+        transport = AsyncMockTransport(send=wrap_in_future(lambda request, **kwargs: response))
+
+        client = TextAnalyticsClient(textanalytics_test_endpoint, AzureKeyCredential(textanalytics_test_api_key), transport=transport)
+
+        with pytest.raises(HttpResponseError) as e:
+            result = await client.analyze_sentiment(["I'm tired"])
+        assert e.value.status_code == 403
+        assert e.value.error.message == 'Out of call volume quota for TextAnalytics F0 pricing tier. Please retry after 15 days. To increase your call volume switch to a paid tier.'