[formrecognizer] raise when call quota exceeded (Azure#20062)

kristapratico · web-flow · commit cc7b25ef6477 · 2021-08-05T11:57:45.000-07:00
* add custom quota exceeded policy

* add mock tests

* set content-type in tests

* use async mock transport helper

* debugging

* remove debugging code - mindep problem

* updating min dep for azure-core

* update changelog
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/CHANGELOG.md b/sdk/formrecognizer/azure-ai-formrecognizer/CHANGELOG.md
@@ -2,14 +2,12 @@
 
 ## 3.1.2 (Unreleased)
 
-### Features Added
-
-### Breaking Changes
-
-### Key Bugs Fixed
-
-### Fixed
+### Bugs Fixed
+- A `HttpResponseError` will be immediately raised when the call quota volume is exceeded in a `F0` tier Form Recognizer
+resource.
 
+### Other Changes
+- Bumped `azure-core` minimum dependency version from `1.8.2` to `1.13.0`
 
 ## 3.1.1 (2021-06-08)
 
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/azure/ai/formrecognizer/_form_base_client.py b/sdk/formrecognizer/azure-ai-formrecognizer/azure/ai/formrecognizer/_form_base_client.py
@@ -8,7 +8,7 @@
 from azure.core.pipeline.policies import HttpLoggingPolicy
 from ._generated._form_recognizer_client import FormRecognizerClient as FormRecognizer
 from ._api_versions import FormRecognizerApiVersion, validate_api_version
-from ._helpers import _get_deserialize, get_authentication_policy, POLLING_INTERVAL
+from ._helpers import _get_deserialize, get_authentication_policy, POLLING_INTERVAL, QuotaExceededPolicy
 from ._user_agent import USER_AGENT
 
 if TYPE_CHECKING:
@@ -58,8 +58,9 @@ def __init__(self, endpoint, credential, **kwargs):
             credential=credential,  # type: ignore
             api_version=self._api_version,
             sdk_moniker=USER_AGENT,
-            authentication_policy=authentication_policy,
-            http_logging_policy=http_logging_policy,
+            authentication_policy=kwargs.get("authentication_policy", authentication_policy),
+            http_logging_policy=kwargs.get("http_logging_policy", http_logging_policy),
+            per_retry_policies=kwargs.get("per_retry_policies", QuotaExceededPolicy()),
             polling_interval=polling_interval,
             **kwargs
         )
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/azure/ai/formrecognizer/_helpers.py b/sdk/formrecognizer/azure-ai-formrecognizer/azure/ai/formrecognizer/_helpers.py
@@ -7,8 +7,10 @@
 import re
 import six
 from azure.core.credentials import AzureKeyCredential
-from azure.core.pipeline.policies import AzureKeyCredentialPolicy
+from azure.core.pipeline.policies import AzureKeyCredentialPolicy, SansIOHTTPPolicy
 from azure.core.pipeline.transport import HttpTransport
+from azure.core.exceptions import HttpResponseError
+
 
 POLLING_INTERVAL = 5
 COGNITIVE_KEY_HEADER = "Ocp-Apim-Subscription-Key"
@@ -164,3 +166,23 @@ def __enter__(self):
 
     def __exit__(self, *args):  # pylint: disable=arguments-differ
         pass
+
+
+class QuotaExceededPolicy(SansIOHTTPPolicy):
+    """Raises an exception immediately when the call quota volume has been exceeded in a F0
+    tier form recognizer 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 in [403, 429] and \
+                "Out of call volume quota for FormRecognizer F0 pricing tier" in http_response.text():
+            raise HttpResponseError(http_response.text(), response=http_response)
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/azure/ai/formrecognizer/aio/_form_base_client_async.py b/sdk/formrecognizer/azure-ai-formrecognizer/azure/ai/formrecognizer/aio/_form_base_client_async.py
@@ -14,7 +14,7 @@
     FormRecognizerClient as FormRecognizer,
 )
 from .._api_versions import FormRecognizerApiVersion, validate_api_version
-from .._helpers import _get_deserialize, get_authentication_policy, POLLING_INTERVAL
+from .._helpers import _get_deserialize, get_authentication_policy, POLLING_INTERVAL, QuotaExceededPolicy
 from .._user_agent import USER_AGENT
 
 if TYPE_CHECKING:
@@ -68,8 +68,9 @@ def __init__(
             credential=credential,  # type: ignore
             api_version=self._api_version,
             sdk_moniker=USER_AGENT,
-            authentication_policy=authentication_policy,
-            http_logging_policy=http_logging_policy,
+            authentication_policy=kwargs.get("authentication_policy", authentication_policy),
+            http_logging_policy=kwargs.get("http_logging_policy", http_logging_policy),
+            per_retry_policies=kwargs.get("per_retry_policies", QuotaExceededPolicy()),
             polling_interval=polling_interval,
             **kwargs
         )
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/setup.py b/sdk/formrecognizer/azure-ai-formrecognizer/setup.py
@@ -80,7 +80,7 @@
         'azure.ai',
     ]),
     install_requires=[
-        "azure-core<2.0.0,>=1.8.2",
+        "azure-core<2.0.0,>=1.13.0",
         "msrest>=0.6.21",
         'six>=1.11.0',
         'azure-common~=1.1',
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/tests/test_logging.py b/sdk/formrecognizer/azure-ai-formrecognizer/tests/test_logging.py
@@ -7,8 +7,15 @@
 
 import logging
 import pytest
+import json
+try:
+    from unittest import mock
+except ImportError:  # python < 3.3
+    import mock  # type: ignore
+
 from azure.ai.formrecognizer import FormRecognizerClient, FormTrainingClient
 from azure.core.credentials import AzureKeyCredential
+from azure.core.exceptions import HttpResponseError
 from testcase import FormRecognizerTest
 from preparers import FormRecognizerPreparer
 
@@ -63,3 +70,47 @@ def test_logging_info_ft_client(self, formrecognizer_test_endpoint, formrecogniz
                     assert message.message.find("REDACTED") != -1
                 else:
                     assert message.message.find("REDACTED") == -1
+
+    @FormRecognizerPreparer()
+    def test_mock_quota_exceeded_403(self, formrecognizer_test_endpoint, formrecognizer_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 FormRecognizer F0 pricing tier. "
+            "Please retry after 1 day. To increase your call volume switch to a paid tier."}}
+        )
+        response.content_type = "application/json"
+        transport = mock.Mock(send=lambda request, **kwargs: response)
+
+        client = FormRecognizerClient(formrecognizer_test_endpoint, AzureKeyCredential(formrecognizer_test_api_key), transport=transport)
+
+        with pytest.raises(HttpResponseError) as e:
+            poller = client.begin_recognize_receipts_from_url(self.receipt_url_jpg)
+        assert e.value.status_code == 403
+        assert e.value.error.message == 'Out of call volume quota for FormRecognizer F0 pricing tier. Please retry after 1 day. To increase your call volume switch to a paid tier.'
+
+    @FormRecognizerPreparer()
+    def test_mock_quota_exceeded_429(self, formrecognizer_test_endpoint, formrecognizer_test_api_key):
+
+        response = mock.Mock(
+            status_code=429,
+            headers={"Retry-After": 186688, "Content-Type": "application/json"},
+            reason="Bad Request"
+        )
+        response.text = lambda encoding=None: json.dumps(
+            {"error": {"code": "429", "message": "Out of call volume quota for FormRecognizer F0 pricing tier. "
+            "Please retry after 1 day. To increase your call volume switch to a paid tier."}}
+        )
+        response.content_type = "application/json"
+        transport = mock.Mock(send=lambda request, **kwargs: response)
+
+        client = FormRecognizerClient(formrecognizer_test_endpoint, AzureKeyCredential(formrecognizer_test_api_key), transport=transport)
+
+        with pytest.raises(HttpResponseError) as e:
+            poller = client.begin_recognize_receipts_from_url(self.receipt_url_jpg)
+        assert e.value.status_code == 429
+        assert e.value.error.message == 'Out of call volume quota for FormRecognizer F0 pricing tier. Please retry after 1 day. To increase your call volume switch to a paid tier.'
diff --git a/sdk/formrecognizer/azure-ai-formrecognizer/tests/test_logging_async.py b/sdk/formrecognizer/azure-ai-formrecognizer/tests/test_logging_async.py
@@ -7,8 +7,17 @@
 
 import logging
 import pytest
+import json
+import sys
+import asyncio
+import functools
+try:
+    from unittest import mock
+except ImportError:  # python < 3.3
+    import mock  # type: ignore
 from azure.ai.formrecognizer.aio import FormRecognizerClient, FormTrainingClient
 from azure.core.credentials import AzureKeyCredential
+from azure.core.exceptions import HttpResponseError
 from preparers import FormRecognizerPreparer
 from asynctestcase import AsyncFormRecognizerTest
 
@@ -22,6 +31,38 @@ def emit(self, record):
         self.messages.append(record)
 
 
+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 TestLogging(AsyncFormRecognizerTest):
 
     @FormRecognizerPreparer()
@@ -64,3 +105,45 @@ async def test_logging_info_ft_client(self, formrecognizer_test_endpoint, formre
                     assert message.message.find("REDACTED") != -1
                 else:
                     assert message.message.find("REDACTED") == -1
+
+    @FormRecognizerPreparer()
+    async def test_mock_quota_exceeded_403(self, formrecognizer_test_endpoint, formrecognizer_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 FormRecognizer F0 pricing tier. "
+            "Please retry after 1 day. 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 = FormRecognizerClient(formrecognizer_test_endpoint, AzureKeyCredential(formrecognizer_test_api_key), transport=transport)
+
+        with pytest.raises(HttpResponseError) as e:
+            poller = await client.begin_recognize_receipts_from_url(self.receipt_url_jpg)
+        assert e.value.status_code == 403
+        assert e.value.error.message == 'Out of call volume quota for FormRecognizer F0 pricing tier. Please retry after 1 day. To increase your call volume switch to a paid tier.'
+
+    @FormRecognizerPreparer()
+    async def test_mock_quota_exceeded_429(self, formrecognizer_test_endpoint, formrecognizer_test_api_key):
+        response = mock.Mock(
+            status_code=429,
+            headers={"Retry-After": 186688, "Content-Type": "application/json"},
+            reason="Bad Request"
+        )
+        response.text = lambda encoding=None: json.dumps(
+            {"error": {"code": "429", "message": "Out of call volume quota for FormRecognizer F0 pricing tier. "
+            "Please retry after 1 day. 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 = FormRecognizerClient(formrecognizer_test_endpoint, AzureKeyCredential(formrecognizer_test_api_key), transport=transport)
+        with pytest.raises(HttpResponseError) as e:
+            poller = await client.begin_recognize_receipts_from_url(self.receipt_url_jpg)
+        assert e.value.status_code == 429
+        assert e.value.error.message == 'Out of call volume quota for FormRecognizer F0 pricing tier. Please retry after 1 day. To increase your call volume switch to a paid tier.'
diff --git a/shared_requirements.txt b/shared_requirements.txt
@@ -153,7 +153,7 @@ chardet<5,>=3.0.2
 #override azure-ai-language-questionanswering msrest>=0.6.21
 #override azure-search-documents azure-core<2.0.0,>=1.14.0
 #override azure-ai-formrecognizer msrest>=0.6.21
-#override azure-ai-formrecognizer azure-core<2.0.0,>=1.8.2
+#override azure-ai-formrecognizer azure-core<2.0.0,>=1.13.0
 #override azure-storage-blob azure-core<2.0.0,>=1.10.0
 #override azure-storage-blob msrest>=0.6.21
 #override azure-storage-blob-changefeed azure-storage-blob>=12.5.0,<13.0.0
