Oauth2ApiClient raise custom error

Author: nonword

src/nypl_py_utils/__init__.py

@@ -2,7 +2,7 @@
 from .classes.kinesis_client import KinesisClient, KinesisClientError  # noqa
 from .classes.kms_client import KmsClient, KmsClientError # noqa
 from .classes.mysql_client import MySQLClient, MySQLClientError  # noqa
-from .classes.oauth2_api_client import Oauth2ApiClient  # noqa
+from .classes.oauth2_api_client import Oauth2ApiClient, Oauth2ApiClientError  # noqa
 from .classes.postgresql_client import PostgreSQLClient, PostgreSQLClientError  # noqa
 from .classes.redshift_client import RedshiftClient, RedshiftClientError  # noqa
 from .classes.s3_client import S3Client, S3ClientError  # noqa

src/nypl_py_utils/classes/oauth2_api_client.py

@@ -71,12 +71,12 @@ def _do_http_method(self, method, request_path, **kwargs):
         except TokenExpiredError:
             self.logger.debug('TokenExpiredError encountered')
 
-            if '_do_http_method_token_refreshes' not in kwargs.keys():
-                kwargs['_do_http_method_token_refreshes'] = 1
-            else:
-                kwargs['_do_http_method_token_refreshes'] += 1
-                if kwargs['_do_http_method_token_refreshes'] > 3:
-                    raise Exception('Exhausted token refreshes')
+            # Raise error after 3 successive token refreshes
+            kwargs['_do_http_method_token_refreshes'] = \
+                kwargs.get('_do_http_method_token_refreshes', 0) + 1
+            if kwargs['_do_http_method_token_refreshes'] > 3:
+                raise Oauth2ApiClientError('Exhausted token refreshes') \
+                    from None
 
             self._generate_access_token()
             return self._do_http_method(method, request_path, **kwargs)
@@ -102,3 +102,8 @@ def _generate_access_token(self):
             client_id=self.client_id,
             client_secret=self.client_secret
         )
+
+
+class Oauth2ApiClientError(Exception):
+    def __init__(self, message=None):
+        self.message = message

tests/test_oauth2_api_client.py

@@ -5,7 +5,7 @@
 from requests_oauthlib import OAuth2Session
 from requests import HTTPError
 
-from nypl_py_utils import Oauth2ApiClient
+from nypl_py_utils import (Oauth2ApiClient, Oauth2ApiClientError)
 
 _TOKEN_RESPONSE = {
     'access_token': 'super-secret-token',
@@ -110,3 +110,28 @@ def test_error_status_raises_error(self, requests_mock, test_instance,
 
         with pytest.raises(HTTPError):
             test_instance._do_http_method('GET', 'foo')
+
+    def test_token_refresh_failure_raises_error(self, requests_mock,
+            test_instance, token_server_post):
+        """
+        Failure to fetch a token can raise a number of errors including:
+         - requests.exceptions.HTTPError for invalid access_token
+         - oauthlib.oauth2.rfc6749.errors.InvalidClientError for invalid
+           credentials
+         - oauthlib.oauth2.rfc6749.errors.MissingTokenError for failure to
+           fetch a token
+        One error that can arise from this client itself is failure to fetch
+        a new valid token in response to token expiration. This test asserts
+        that the client will not allow more than successive 3 retries.
+        """
+        requests_mock.get(f'{BASE_URL}/foo', json={'foo': 'bar'})
+
+        token_response = dict(_TOKEN_RESPONSE)
+        token_response['expires_in'] = 0
+        token_server_post = requests_mock\
+            .post(TOKEN_URL, text=json.dumps(token_response))
+
+        with pytest.raises(Oauth2ApiClientError):
+            test_instance._do_http_method('GET', 'foo')
+        # Expect 1 initial token fetch, plus 3 retries:
+        assert len(token_server_post.request_history) == 4