Refactor token expiry handling in DatabricksTokenFederationProvider and enhance unit tests for accurate expiry verification

Author: madhav-db

src/databricks/sql/auth/token_federation.py

@@ -335,45 +335,13 @@ def _exchange_token(self, access_token: str) -> Token:
         token_type = resp_data.get("token_type", "Bearer")
         refresh_token = resp_data.get("refresh_token", "")
 
-        # Determine token expiry - first try from JWT claims
+        # Extract expiry from JWT claims
         expiry = self._get_expiry_from_jwt(new_access_token)
-
-        # If JWT expiry not available, use expires_in from response
         if expiry is None:
-            expiry = self._get_expiry_from_response(resp_data)
-
-        # If we still don't have an expiry, we can't proceed
-        if expiry is None:
-            raise ValueError(
-                "Unable to determine token expiry from response or JWT claims"
-            )
+            raise ValueError("Unable to determine token expiry from JWT claims")
 
         return Token(new_access_token, token_type, refresh_token, expiry)
 
-    def _get_expiry_from_response(
-        self, resp_data: Dict[str, Any]
-    ) -> Optional[datetime]:
-        """
-        Extract expiry datetime from response data.
-
-        Args:
-            resp_data: Response data from token exchange
-
-        Returns:
-            Optional[datetime]: Expiry datetime if found in response, None otherwise
-        """
-        if "expires_in" not in resp_data or not resp_data["expires_in"]:
-            return None
-
-        try:
-            expires_in = int(resp_data["expires_in"])
-            expiry = datetime.now(tz=timezone.utc) + timedelta(seconds=expires_in)
-            logger.debug(f"Using expiry from expires_in: {expiry}")
-            return expiry
-        except (ValueError, TypeError) as e:
-            logger.warning(f"Invalid expires_in value: {str(e)}")
-            return None
-
 
 class SimpleCredentialsProvider(CredentialsProvider):
     """A simple credentials provider that returns a fixed token."""

tests/unit/test_token_federation.py

@@ -344,20 +344,26 @@ def test_token_exchange_success(self, federation_provider):
         """Test successful token exchange."""
         # Mock successful response
         with patch("databricks.sql.auth.token_federation.requests.post") as mock_post:
+            # Create a token with a valid expiry
+            expiry_timestamp = int(
+                (datetime.now(tz=timezone.utc) + timedelta(hours=1)).timestamp()
+            )
+
             # Configure mock response
             mock_response = MagicMock()
             mock_response.status_code = 200
             mock_response.json.return_value = {
                 "access_token": "new_token",
                 "token_type": "Bearer",
                 "refresh_token": "refresh_value",
-                "expires_in": 3600,
             }
             mock_post.return_value = mock_response
 
-            # Patch the _get_expiry_from_jwt method to return None (forcing use of expires_in)
+            # Mock JWT expiry extraction to return a valid expiry
             with patch.object(
-                federation_provider, "_get_expiry_from_jwt", return_value=None
+                federation_provider,
+                "_get_expiry_from_jwt",
+                return_value=datetime.fromtimestamp(expiry_timestamp, tz=timezone.utc),
             ):
                 # Call the exchange method
                 token = federation_provider._exchange_token("original_token")
@@ -367,10 +373,11 @@ def test_token_exchange_success(self, federation_provider):
                 assert token.token_type == "Bearer"
                 assert token.refresh_token == "refresh_value"
 
-                # Verify expiry time (should be ~1 hour in future)
-                now = datetime.now(tz=timezone.utc)
-                assert token.expiry > now
-                assert token.expiry < now + timedelta(seconds=3601)
+                # Verify expiry time is correctly set
+                expiry_datetime = datetime.fromtimestamp(
+                    expiry_timestamp, tz=timezone.utc
+                )
+                assert token.expiry == expiry_datetime
 
     def test_token_exchange_failure(self, federation_provider):
         """Test token exchange failure handling."""