diff --git a/intuitlib/utils.py b/intuitlib/utils.py index c1fd832..cd606a7 100644 --- a/intuitlib/utils.py +++ b/intuitlib/utils.py @@ -24,16 +24,19 @@ from base64 import b64encode, b64decode from datetime import datetime from requests.sessions import Session +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry from intuitlib.config import DISCOVERY_URL, ACCEPT_HEADER from intuitlib.enums import Scopes from intuitlib.exceptions import AuthClientError -def get_discovery_doc(environment, session=None): +def get_discovery_doc(environment, session=None, max_retries=3): """Gets discovery doc based on environment specified. :param environment: App environment, accepted values: 'sandbox','production','prod','e2e' :param session: `requests.Session` object if a session is already being used, defaults to None + :param max_retries: Maximum number of retry attempts, defaults to 3 :return: Discovery doc response :raises HTTPError: if response status != 200 """ @@ -43,11 +46,23 @@ def get_discovery_doc(environment, session=None): discovery_url = DISCOVERY_URL['sandbox'] else: discovery_url = environment - + + retry_strategy = Retry( + total=max_retries, + backoff_factor=1, + status_forcelist=[429, 500, 502, 503, 504], + ) + if session is not None and isinstance(session, Session): + adapter = HTTPAdapter(max_retries=retry_strategy) + session.mount("https://", adapter) response = session.get(url=discovery_url) else: - response = requests.get(url=discovery_url) + adapter = HTTPAdapter(max_retries=retry_strategy) + with requests.Session() as temp_session: + temp_session.mount("https://", adapter) + response = temp_session.get(url=discovery_url) + if response.status_code != 200: raise AuthClientError(response) return response.json() diff --git a/tests/test_utils.py b/tests/test_utils.py index 9cf2059..4f434a2 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -58,10 +58,12 @@ def test_get_discovery_doc_custom_url_input(self): assert discovery_doc['issuer'] =='https://oauth.platform.intuit.com/op/v1' assert discovery_doc['userinfo_endpoint'] == 'https://sandbox-accounts.platform.intuit.com/v1/openid_connect/userinfo' - @mock.patch('intuitlib.utils.requests.get') - def test_get_discovery_doc_bad_response(self, mock_get): + @mock.patch('intuitlib.utils.requests.Session') + def test_get_discovery_doc_bad_response(self, mock_session_class): + mock_session = mock.Mock() + mock_session_class.return_value.__enter__.return_value = mock_session mock_resp = self.mock_request(status=400) - mock_get.return_value = mock_resp + mock_session.get.return_value = mock_resp with pytest.raises(AuthClientError): get_discovery_doc('sandbox')