Refresh the NewMode token when needed

parsons/newmode/newmode.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 from Newmode import Client
+from oauthlib.oauth2 import TokenExpiredError
 
 from parsons.etl import Table
 from parsons.utilities import check_env
@@ -429,7 +430,10 @@ def __init__(
         self.client_id: str = check_env.check("NEWMODE_API_CLIENT_ID", client_id)
         self.client_secret: str = check_env.check("NEWMODE_API_CLIENT_SECRET", client_secret)
         self.headers: Dict[str, str] = {"content-type": "application/json"}
-        self.default_client: OAuth2APIConnector = OAuth2APIConnector(
+        self.default_client: OAuth2APIConnector = self.get_default_oauth_client()
+
+    def get_default_oauth_client(self) -> OAuth2APIConnector:
+        return OAuth2APIConnector(
             uri=self.base_url,
             auto_refresh_url=V2_API_AUTH_URL,
             client_id=self.client_id,
@@ -472,10 +476,14 @@ def base_request(
 
         for attempt in range(retries + 1):
             try:
-                response = client.request(
-                    url=url, req_type=method, json=json, data=data, params=params
-                )
-                return self.checked_response(response, client)
+                try:
+                    response = client.request(
+                        url=url, req_type=method, json=json, data=data, params=params
+                    )
+                    return self.checked_response(response, client)
+                except TokenExpiredError as e:
+                    logger.warning(f"Token expired: {e}. Refreshing it...")
+                    self.default_client = self.get_default_oauth_client()
             except Exception as e:
                 if attempt < retries:
                     logger.warning(f"Request failed (attempt {attempt + 1}/{retries}). Retrying...")

test/test_newmode/test_newmode.py

@@ -4,6 +4,7 @@
 from unittest.mock import call, patch
 
 import requests_mock
+from oauthlib.oauth2 import TokenExpiredError
 from requests.exceptions import HTTPError
 
 from parsons import Newmode, Table
@@ -337,3 +338,38 @@ def test_checked_response_http_error(self, m):
         )
         with self.assertRaises(HTTPError):
             self.nm.checked_response(response, self.nm.default_client)
+
+    @requests_mock.Mocker()
+    @patch("parsons.newmode.newmode.NewmodeV2.get_default_oauth_client")
+    def test_token_refresh_on_expired_token(self, m, mock_get_default_oauth_client):
+        m.post(V2_API_AUTH_URL, json={"access_token": "fakeAccessToken"})
+        m.get(f"{V2_API_URL}v2.1/test-endpoint", status_code=401)
+
+        mock_new_client = mock.MagicMock()
+        mock_get_default_oauth_client.return_value = mock_new_client
+        self.nm.default_client.request = mock.MagicMock()
+
+        mock_response = mock.MagicMock()
+        mock_response.raise_for_status = mock.MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"data": "success"}
+
+        # Simulate token expiration and successful response
+        def oauth_side_effect(*args, **kwargs):
+            if not hasattr(self, "call_count"):
+                self.call_count = 0
+            if self.call_count == 0:
+                self.call_count += 1
+                raise TokenExpiredError()
+            return mock_response
+
+        self.nm.default_client.request.side_effect = oauth_side_effect
+
+        response = self.nm.base_request(
+            method="GET",
+            url=f"{V2_API_URL}v2.1/test-endpoint",
+            client=self.nm.default_client,
+        )
+
+        mock_get_default_oauth_client.assert_called_once()
+        self.assertEqual(response, {"data": "success"})