fix(python-cdk): convert expires_in to int when refreshing (#20301)

airbytehq · Mar 9, 2023 · d67afbb · d67afbb
1 parent f472190
commit d67afbb
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 2 deletions.
diff --git a/airbyte-cdk/python/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py b/airbyte-cdk/python/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py
@@ -97,7 +97,7 @@ def refresh_access_token(self) -> Tuple[str, int]:
         :return: a tuple of (access_token, token_lifespan_in_seconds)
         """
         response_json = self._get_refresh_access_token_response()
-        return response_json[self.get_access_token_name()], response_json[self.get_expires_in_name()]
+        return response_json[self.get_access_token_name()], int(response_json[self.get_expires_in_name()])
 
     @abstractmethod
     def get_token_refresh_endpoint(self) -> str:

diff --git a/airbyte-cdk/python/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py b/airbyte-cdk/python/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py
@@ -234,6 +234,6 @@ def refresh_access_token(self) -> Tuple[str, int, str]:
         response_json = self._get_refresh_access_token_response()
         return (
             response_json[self.get_access_token_name()],
-            response_json[self.get_expires_in_name()],
+            int(response_json[self.get_expires_in_name()]),
             response_json[self.get_refresh_token_name()],
         )
diff --git a/.../python/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py b/.../python/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py
@@ -163,6 +163,12 @@ def test_refresh_access_token(self, mocker):
 
         assert ("access_token", 1000) == token
 
+        # Test with expires_in as str
+        mocker.patch.object(resp, "json", return_value={"access_token": "access_token", "expires_in": "2000"})
+        token = oauth.refresh_access_token()
+
+        assert ("access_token", 2000) == token
+
     @pytest.mark.parametrize("error_code", (429, 500, 502, 504))
     def test_refresh_access_token_retry(self, error_code, requests_mock):
         oauth = Oauth2Authenticator(
@@ -260,6 +266,7 @@ def test_refresh_access_token(self, mocker, connector_config):
             connector_config,
             token_refresh_endpoint="foobar",
         )
+
         authenticator._get_refresh_access_token_response = mocker.Mock(
             return_value={
                 authenticator.get_access_token_name(): "new_access_token",
@@ -269,6 +276,16 @@ def test_refresh_access_token(self, mocker, connector_config):
         )
         assert authenticator.refresh_access_token() == ("new_access_token", 42, "new_refresh_token")
 
+        # Test with expires_in as str
+        authenticator._get_refresh_access_token_response = mocker.Mock(
+            return_value={
+                authenticator.get_access_token_name(): "new_access_token",
+                authenticator.get_expires_in_name(): "1000",
+                authenticator.get_refresh_token_name(): "new_refresh_token",
+            }
+        )
+        assert authenticator.refresh_access_token() == ("new_access_token", 1000, "new_refresh_token")
+
 
 def mock_request(method, url, data):
     if url == "refresh_end":