Merge pull request #12 from NYPL/noref-improve-oauth2-api-client

Fix Oauth2ApiClient token refresh, responses

Author: nonword

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## v0.0.9 - 3/16/23
+- Improve Oauth2ApiClient token refresh and method responses
+
 ## v0.0.8 - 3/3/23
 - Pass in all kwargs from PostgreSQLClient to ConnectionPool so that all ConnectionPool settings
 can be set from the wrapper

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

@@ -21,8 +21,7 @@ def __init__(self, client_id=None, client_secret=None, base_url=None,
         self.base_url = base_url \
             or os.environ.get('NYPL_API_BASE_URL', None)
 
-        self.client = None
-        self.token = None
+        self.oauth_client = None
 
         self.logger = create_log('oauth2_api_client')
 
@@ -56,39 +55,52 @@ def _do_http_method(self, method, request_path, **kwargs):
         """
         Issue an HTTP method call on on the given request_path
         """
-        if not self.client:
+        if not self.oauth_client:
             self._create_oauth_client()
 
         url = f'{self.base_url}/{request_path}'
         self.logger.debug(f'{method} {url}')
 
         try:
-            return self.oauth_client.request(method, url, **kwargs).json()
-        except TokenExpiredError as e:
-            self.logger.debug(f'TokenExpiredError encountered: {e}')
-            self._generate_access_token()
+            # Build kwargs cleaned of local variables:
+            kwargs_cleaned = {k: kwargs[k] for k in kwargs
+                              if not k.startswith('_do_http_method_')}
+            resp = self.oauth_client.request(method, url, **kwargs_cleaned)
+            resp.raise_for_status()
+            return resp
+        except TokenExpiredError:
+            self.logger.debug('TokenExpiredError encountered')
+
+            # 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)
-        except TimeoutError as e:
-            self.logger.error(f'TimeoutError encountered: {e}')
-            return {}
 
     def _create_oauth_client(self):
         """
         Creates an authenticated a OAuth2Session instance for later requests
         """
+        client = BackendApplicationClient(client_id=self.client_id)
+        self.oauth_client = OAuth2Session(client=client)
         self._generate_access_token()
-        self.oauth_client = OAuth2Session(self.client_id, token=self.token)
 
     def _generate_access_token(self):
         """
         Fetch and store a fresh token
         """
-        client = BackendApplicationClient(client_id=self.client_id)
-        oauth = OAuth2Session(client=client)
         self.logger.debug(f'Refreshing token via @{self.token_url}')
-        self.token = oauth.fetch_token(
+        self.oauth_client.fetch_token(
             token_url=self.token_url,
             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

@@ -1,32 +1,36 @@
 import os
+import time
 import json
 import pytest
 from requests_oauthlib import OAuth2Session
+from requests import HTTPError
 
-from nypl_py_utils import Oauth2ApiClient
-# from requests.exceptions import ConnectTimeout
+from nypl_py_utils import (Oauth2ApiClient, Oauth2ApiClientError)
 
 _TOKEN_RESPONSE = {
     'access_token': 'super-secret-token',
-    'expires_in': 3600,
+    'expires_in': 1,
     'token_type': 'Bearer',
     'scope': ['offline_access', 'openid', 'login:staff', 'admin'],
-    'id_token': 'super-secret-token',
-    'expires_at': 1677599823.3180869
+    'id_token': 'super-secret-token'
 }
 
 BASE_URL = 'https://example.com/api/v0.1'
+TOKEN_URL = 'https://oauth.example.com/oauth/token'
 
 
 class TestOauth2ApiClient:
 
     @pytest.fixture
-    def test_instance(self, requests_mock):
-        token_url = 'https://oauth.example.com/oauth/token'
-        requests_mock.post(token_url, text=json.dumps(_TOKEN_RESPONSE))
+    def token_server_post(self, requests_mock):
+        token_url = TOKEN_URL
+        token_response = dict(_TOKEN_RESPONSE)
+        return requests_mock.post(token_url, text=json.dumps(token_response))
 
+    @pytest.fixture
+    def test_instance(self, requests_mock):
         return Oauth2ApiClient(base_url=BASE_URL,
-                               token_url=token_url,
+                               token_url=TOKEN_URL,
                                client_id='clientid',
                                client_secret='clientsecret'
                                )
@@ -50,42 +54,84 @@ def test_uses_env_vars(self):
         for key, value in env.items():
             os.environ[key] = ''
 
-    def test_generate_access_token(self, test_instance):
+    def test_generate_access_token(self, test_instance, token_server_post):
+        test_instance._create_oauth_client()
         test_instance._generate_access_token()
-        assert test_instance.token['access_token']\
+        assert test_instance.oauth_client.token['access_token']\
             == _TOKEN_RESPONSE['access_token']
 
-    def test_create_oauth_client(self, test_instance):
+    def test_create_oauth_client(self, token_server_post, test_instance):
         test_instance._create_oauth_client()
         assert type(test_instance.oauth_client) is OAuth2Session
 
-    def test_do_http_method(self, requests_mock, test_instance):
+    def test_do_http_method(self, requests_mock, token_server_post,
+                            test_instance):
+        requests_mock.get(f'{BASE_URL}/foo', json={'foo': 'bar'})
+
         requests_mock.get(f'{BASE_URL}/foo', json={'foo': 'bar'})
         resp = test_instance._do_http_method('GET', 'foo')
-        assert resp == {'foo': 'bar'}
+        assert resp.status_code == 200
+        assert resp.json() == {'foo': 'bar'}
 
-    def test_token_expiration(self, requests_mock, test_instance):
+    def test_token_expiration(self, requests_mock, test_instance,
+                              token_server_post, mocker):
         api_get_mock = requests_mock.get(f'{BASE_URL}/foo',
                                          json={'foo': 'bar'})
 
-        # Perform first request to auto-authenticate:
-        resp = test_instance._do_http_method('GET', 'foo')
+        # Perform first request:
+        test_instance._do_http_method('GET', 'foo')
+        # Expect this first call triggered a single token server call:
+        assert len(token_server_post.request_history) == 1
+        # And the GET request used the supplied Bearer token:
         assert api_get_mock.request_history[0]._request\
             .headers['Authorization'] == 'Bearer super-secret-token'
 
-        # Emulate token expiration:
-        test_instance.token['expires_at'] = 0
-
-        token_post_mock = requests_mock.post(
-            test_instance.token_url,
-            text=json.dumps(_TOKEN_RESPONSE)
-        )
+        # The token obtained above expires in 1s, so wait out expiration:
+        time.sleep(1.1)
+
+        # Register new token response:
+        second_token_response = dict(_TOKEN_RESPONSE)
+        second_token_response['id_token'] = 'super-secret-second-token'
+        second_token_response['access_token'] = 'super-secret-second-token'
+        second_token_server_post = requests_mock\
+            .post(TOKEN_URL, text=json.dumps(second_token_response))
+
+        # Perform second request:
+        test_instance._do_http_method('GET', 'foo')
+        # Expect a call on the second token server:
+        assert len(second_token_server_post.request_history) == 1
+        # Expect the second GET request to carry the new Bearer token:
+        assert api_get_mock.request_history[1]._request\
+            .headers['Authorization'] == 'Bearer super-secret-second-token'
+
+    def test_error_status_raises_error(self, requests_mock, test_instance,
+                                       token_server_post):
+        requests_mock.get(f'{BASE_URL}/foo', status_code=400)
+
+        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'})
 
-        # Perform second request, which should detect token expiration and
-        # re-authenticate:
-        resp = test_instance._do_http_method('GET', 'foo')
+        token_response = dict(_TOKEN_RESPONSE)
+        token_response['expires_in'] = 0
+        token_server_post = requests_mock\
+            .post(TOKEN_URL, text=json.dumps(token_response))
 
-        assert token_post_mock.called is True
-        assert api_get_mock.request_history[1]._request\
-            .headers['Authorization'] == 'Bearer super-secret-token'
-        assert resp == {'foo': 'bar'}
+        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