diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java index 77045df97..57c97a2ff 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/Consent.java @@ -155,6 +155,131 @@ public String getClientSecret() { return clientSecret; } + /** + * Launch a browser to collect an authorization code and exchange the code for an OAuth token. + * + * @return A {@code SessionCredentials} instance representing the retrieved OAuth token. + * @throws IOException if the webserver cannot be started, or if the browser cannot be opened. + */ + public SessionCredentials launchExternalBrowser() throws IOException { + Map params = getOAuthCallbackParameters(); + return exchangeCallbackParameters(params); + } + + /** + * Exchange callback parameters for OAuth credentials. + * + * @param query The callback parameters from the OAuth flow + * @return A {@code SessionCredentials} instance representing the retrieved OAuth token. + */ + public SessionCredentials exchangeCallbackParameters(Map query) { + validateCallbackParameters(query); + Token token = exchange(query.get("code"), query.get("state")); + return new SessionCredentials.Builder() + .withHttpClient(this.hc) + .withClientId(this.clientId) + .withClientSecret(this.clientSecret) + .withTokenUrl(this.tokenUrl) + .withToken(token) + .build(); + } + + /** + * Launches an external browser to collect OAuth callback parameters and exchanges them for an + * OAuth token. + * + * @return A {@code Token} instance containing the OAuth access token and related credentials + * @throws IOException if the local HTTP server cannot be started, the browser cannot be opened, + * or there are network issues during the token exchange + * @throws DatabricksException if the OAuth callback contains an error, missing required + * parameters, or if there's a state mismatch during the token exchange. + */ + Token getTokenFromExternalBrowser() throws IOException { + Map params = getOAuthCallbackParameters(); + validateCallbackParameters(params); + return exchange(params.get("code"), params.get("state")); + } + + protected void desktopBrowser() throws IOException { + Desktop.getDesktop().browse(URI.create(this.authUrl)); + } + + /** + * Handles the OAuth callback by setting up a local HTTP server, launching the browser, and + * collecting the callback parameters. + * + * @return A map containing the callback parameters from the OAuth flow. + * @throws IOException if the webserver cannot be started, or if the browser cannot be opened. + */ + private Map getOAuthCallbackParameters() throws IOException { + URL redirect = new URL(getRedirectUrl()); + if (!Arrays.asList("localhost", "127.0.0.1").contains(redirect.getHost())) { + throw new IllegalArgumentException( + "cannot listen on " + + redirect.getHost() + + ", redirectUrl host must be one of: localhost, 127.0.0.1"); + } + CallbackResponseHandler handler = new CallbackResponseHandler(); + HttpServer httpServer = + HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0); + httpServer.createContext("/", handler); + httpServer.start(); + desktopBrowser(); + Map params = handler.getParams(); + httpServer.stop(0); + return params; + } + + /** + * Validates the OAuth callback parameters to ensure they contain the required fields and no error + * conditions. + * + * @param query The callback parameters to validate + * @throws DatabricksException if validation fails due to error conditions or missing required + * parameters + */ + private void validateCallbackParameters(Map query) { + if (query.containsKey("error")) { + throw new DatabricksException(query.get("error") + ": " + query.get("error_description")); + } + if (!query.containsKey("code") || !query.containsKey("state")) { + throw new DatabricksException("No code returned in callback"); + } + } + + /** + * Exchange authorization code for OAuth token. + * + * @param code The authorization code from the OAuth callback + * @param state The state parameter from the OAuth callback + * @return A {@code Token} instance representing the OAuth token + */ + private Token exchange(String code, String state) { + if (!this.state.equals(state)) { + throw new DatabricksException( + "state mismatch: original state: " + this.state + "; retrieved state: " + state); + } + Map params = new HashMap<>(); + params.put("grant_type", "authorization_code"); + params.put("code", code); + params.put("code_verifier", this.verifier); + params.put("redirect_uri", this.redirectUrl); + Map headers = new HashMap<>(); + if (this.tokenUrl.contains("microsoft")) { + headers.put("Origin", this.redirectUrl); + } + Token token = + RefreshableTokenSource.retrieveToken( + this.hc, + this.clientId, + this.clientSecret, + this.tokenUrl, + params, + headers, + AuthParameterPosition.BODY); + return token; + } + static class CallbackResponseHandler implements HttpHandler { private final Logger LOG = LoggerFactory.getLogger(getClass().getName()); // Protects params @@ -258,75 +383,4 @@ public Map getParams() { } } } - - /** - * Launch a browser to collect an authorization code and exchange the code for an OAuth token. - * - * @return A {@code SessionCredentials} instance representing the retrieved OAuth token. - * @throws IOException if the webserver cannot be started, or if the browser cannot be opened - */ - public SessionCredentials launchExternalBrowser() throws IOException { - URL redirect = new URL(getRedirectUrl()); - if (!Arrays.asList("localhost", "127.0.0.1").contains(redirect.getHost())) { - throw new IllegalArgumentException( - "cannot listen on " - + redirect.getHost() - + ", redirectUrl host must be one of: localhost, 127.0.0.1"); - } - CallbackResponseHandler handler = new CallbackResponseHandler(); - HttpServer httpServer = - HttpServer.create(new InetSocketAddress(redirect.getHost(), redirect.getPort()), 0); - httpServer.createContext("/", handler); - httpServer.start(); - desktopBrowser(); - Map params = handler.getParams(); - httpServer.stop(0); - return exchangeCallbackParameters(params); - } - - protected void desktopBrowser() throws IOException { - Desktop.getDesktop().browse(URI.create(this.authUrl)); - } - - public SessionCredentials exchangeCallbackParameters(Map query) { - if (query.containsKey("error")) { - throw new DatabricksException(query.get("error") + ": " + query.get("error_description")); - } - if (!query.containsKey("code") || !query.containsKey("state")) { - throw new DatabricksException("No code returned in callback"); - } - return exchange(query.get("code"), query.get("state")); - } - - public SessionCredentials exchange(String code, String state) { - if (!this.state.equals(state)) { - throw new DatabricksException( - "state mismatch: original state: " + this.state + "; retrieved state: " + state); - } - Map params = new HashMap<>(); - params.put("grant_type", "authorization_code"); - params.put("code", code); - params.put("code_verifier", this.verifier); - params.put("redirect_uri", this.redirectUrl); - Map headers = new HashMap<>(); - if (this.tokenUrl.contains("microsoft")) { - headers.put("Origin", this.redirectUrl); - } - Token token = - RefreshableTokenSource.retrieveToken( - this.hc, - this.clientId, - this.clientSecret, - this.tokenUrl, - params, - headers, - AuthParameterPosition.BODY); - return new SessionCredentials.Builder() - .withHttpClient(this.hc) - .withClientId(this.clientId) - .withClientSecret(this.clientSecret) - .withTokenUrl(this.tokenUrl) - .withToken(token) - .build(); - } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java index 7bae60022..af67daeba 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProvider.java @@ -6,6 +6,7 @@ import java.io.IOException; import java.nio.file.Path; import java.util.Objects; +import java.util.Optional; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -66,21 +67,20 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { LOGGER.debug("Found cached token for {}:{}", config.getHost(), clientId); try { - // Create SessionCredentials with the cached token and try to refresh if needed - SessionCredentials cachedCreds = - new SessionCredentials.Builder() - .withToken(cachedToken) - .withHttpClient(config.getHttpClient()) - .withClientId(clientId) - .withClientSecret(clientSecret) - .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .withTokenCache(tokenCache) - .build(); + // Create SessionCredentialsTokenSource with the cached token and try to refresh if needed + SessionCredentialsTokenSource tokenSource = + new SessionCredentialsTokenSource( + cachedToken, + config.getHttpClient(), + config.getOidcEndpoints().getTokenEndpoint(), + clientId, + clientSecret, + Optional.of(config.getEffectiveOAuthRedirectUrl()), + Optional.of(tokenCache)); LOGGER.debug("Using cached token, will immediately refresh"); - cachedCreds.token = cachedCreds.refresh(); - return cachedCreds.configure(config); + tokenSource.token = tokenSource.refresh(); + return OAuthHeaderFactory.fromTokenSource(tokenSource); } catch (Exception e) { // If token refresh fails, log and continue to browser auth LOGGER.info("Token refresh failed: {}, falling back to browser auth", e.getMessage()); @@ -88,17 +88,17 @@ public OAuthHeaderFactory configure(DatabricksConfig config) { } // If no cached token or refresh failed, perform browser auth - SessionCredentials credentials = + SessionCredentialsTokenSource tokenSource = performBrowserAuth(config, clientId, clientSecret, tokenCache); - tokenCache.save(credentials.getToken()); - return credentials.configure(config); + tokenCache.save(tokenSource.getToken()); + return OAuthHeaderFactory.fromTokenSource(tokenSource); } catch (IOException | DatabricksException e) { LOGGER.error("Failed to authenticate: {}", e.getMessage()); return null; } } - SessionCredentials performBrowserAuth( + SessionCredentialsTokenSource performBrowserAuth( DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache) throws IOException { LOGGER.debug("Performing browser authentication"); @@ -113,18 +113,17 @@ SessionCredentials performBrowserAuth( .build(); Consent consent = client.initiateConsent(); - // Use the existing browser flow to get credentials - SessionCredentials credentials = consent.launchExternalBrowser(); - - // Create a new SessionCredentials with the same token but with our token cache - return new SessionCredentials.Builder() - .withToken(credentials.getToken()) - .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withClientSecret(config.getClientSecret()) - .withTokenUrl(config.getOidcEndpoints().getTokenEndpoint()) - .withRedirectUrl(config.getEffectiveOAuthRedirectUrl()) - .withTokenCache(tokenCache) - .build(); + // Use the existing browser flow to get credentials. + Token token = consent.getTokenFromExternalBrowser(); + + // Create a SessionCredentialsTokenSource with the token from browser auth. + return new SessionCredentialsTokenSource( + token, + config.getHttpClient(), + config.getOidcEndpoints().getTokenEndpoint(), + config.getClientId(), + config.getClientSecret(), + Optional.ofNullable(config.getEffectiveOAuthRedirectUrl()), + Optional.ofNullable(tokenCache)); } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java index 4d2d512e3..5a05b8751 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentials.java @@ -2,11 +2,9 @@ import com.databricks.sdk.core.CredentialsProvider; import com.databricks.sdk.core.DatabricksConfig; -import com.databricks.sdk.core.DatabricksException; import com.databricks.sdk.core.http.HttpClient; import java.io.Serializable; -import java.util.HashMap; -import java.util.Map; +import java.util.Optional; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -17,11 +15,24 @@ * requests to an API, and a long-lived refresh token, which can be used to fetch new access tokens. * Calling refresh() uses the refresh token to retrieve a new access token to authenticate to APIs. */ -public class SessionCredentials extends RefreshableTokenSource - implements CredentialsProvider, Serializable { +public class SessionCredentials implements CredentialsProvider, Serializable { private static final long serialVersionUID = 3083941540130596650L; private static final Logger LOGGER = LoggerFactory.getLogger(SessionCredentials.class); + private final SessionCredentialsTokenSource tokenSource; + + private SessionCredentials(Builder b) { + this.tokenSource = + new SessionCredentialsTokenSource( + b.token, + b.hc, + b.tokenUrl, + b.clientId, + b.clientSecret, + Optional.ofNullable(b.redirectUrl), + Optional.ofNullable(b.tokenCache)); + } + @Override public String authType() { return "oauth-u2m"; @@ -29,7 +40,7 @@ public String authType() { @Override public OAuthHeaderFactory configure(DatabricksConfig config) { - return OAuthHeaderFactory.fromTokenSource(this); + return OAuthHeaderFactory.fromTokenSource(tokenSource); } static class Builder { @@ -80,52 +91,4 @@ public SessionCredentials build() { return new SessionCredentials(this); } } - - private final HttpClient hc; - private final String tokenUrl; - private final String redirectUrl; - private final String clientId; - private final String clientSecret; - private final TokenCache tokenCache; - - private SessionCredentials(Builder b) { - super(b.token); - this.hc = b.hc; - this.tokenUrl = b.tokenUrl; - this.redirectUrl = b.redirectUrl; - this.clientId = b.clientId; - this.clientSecret = b.clientSecret; - this.tokenCache = b.tokenCache; - } - - @Override - protected Token refresh() { - if (this.token == null) { - throw new DatabricksException("oauth2: token is not set"); - } - String refreshToken = this.token.getRefreshToken(); - if (refreshToken == null) { - throw new DatabricksException("oauth2: token expired and refresh token is not set"); - } - - Map params = new HashMap<>(); - params.put("grant_type", "refresh_token"); - params.put("refresh_token", refreshToken); - Map headers = new HashMap<>(); - if (tokenUrl.contains("microsoft")) { - // Tokens issued for the 'Single-Page Application' client-type may only be redeemed via - // cross-origin requests - headers.put("Origin", redirectUrl); - } - Token newToken = - retrieveToken( - hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); - - // Save the refreshed token directly to cache - if (tokenCache != null) { - tokenCache.save(newToken); - LOGGER.debug("Saved refreshed token to cache"); - } - return newToken; - } } diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java new file mode 100644 index 000000000..8c71f2eb3 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/SessionCredentialsTokenSource.java @@ -0,0 +1,102 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * TokenSource that handles OAuth token refresh for SessionCredentials. + * + *

Implements the refresh_token OAuth grant type with optional token caching. + */ +public class SessionCredentialsTokenSource extends RefreshableTokenSource { + private static final Logger LOGGER = LoggerFactory.getLogger(SessionCredentialsTokenSource.class); + + // HTTP client for making token refresh requests + private final HttpClient hc; + // OAuth token endpoint URL for refresh requests + private final String tokenUrl; + // OAuth redirect URL, used for Microsoft OAuth endpoints + private final Optional redirectUrl; + // OAuth client ID for authentication + private final String clientId; + // OAuth client secret for authentication + private final String clientSecret; + // Optional token cache for persisting refreshed tokens + private final Optional tokenCache; + + /** + * Constructs a new SessionCredentialsTokenSource. + * + * @param token The initial token to use + * @param hc The HTTP client for making token refresh requests + * @param tokenUrl The OAuth token endpoint URL + * @param clientId The OAuth client ID + * @param clientSecret The OAuth client secret + * @param redirectUrl The OAuth redirect URL (optional) + * @param tokenCache The token cache for persisting refreshed tokens (optional) + */ + public SessionCredentialsTokenSource( + Token token, + HttpClient hc, + String tokenUrl, + String clientId, + String clientSecret, + Optional redirectUrl, + Optional tokenCache) { + super(token); + this.hc = hc; + this.tokenUrl = tokenUrl; + this.clientId = clientId; + this.clientSecret = clientSecret; + this.redirectUrl = redirectUrl; + this.tokenCache = tokenCache; + } + + /** + * Refreshes the OAuth token using the refresh_token grant type. + * + *

This method attempts to refresh the current token using the refresh token. If successful, + * the new token is automatically saved to the token cache if one is configured. For Microsoft + * OAuth endpoints, it includes the Origin header. + * + * @return A new {@link Token} with updated access and refresh tokens. + * @throws DatabricksException if the token is not set, refresh token is missing, or the refresh + * request fails. + */ + @Override + protected Token refresh() { + if (this.token == null) { + throw new DatabricksException("oauth2: token is not set"); + } + String refreshToken = this.token.getRefreshToken(); + if (refreshToken == null) { + throw new DatabricksException("oauth2: token expired and refresh token is not set"); + } + + Map params = new HashMap<>(); + params.put("grant_type", "refresh_token"); + params.put("refresh_token", refreshToken); + Map headers = new HashMap<>(); + if (tokenUrl.contains("microsoft")) { + // Tokens issued for the 'Single-Page Application' client-type may only be redeemed via + // cross-origin requests + redirectUrl.ifPresent(url -> headers.put("Origin", url)); + } + Token newToken = + retrieveToken( + hc, clientId, clientSecret, tokenUrl, params, headers, AuthParameterPosition.BODY); + + // Save the refreshed token directly to cache + tokenCache.ifPresent( + cache -> { + cache.save(newToken); + LOGGER.debug("Saved refreshed token to cache"); + }); + return newToken; + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java index 932690bd7..b7e237ddd 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java @@ -17,6 +17,7 @@ import java.util.Arrays; import java.util.HashMap; import java.util.Map; +import java.util.Optional; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import org.mockito.Mockito; @@ -133,7 +134,9 @@ void exchangeTest() throws IOException { URL url = new URL("https://databricks.com/"); // Mock because it's a POST Request to http client - Mockito.doReturn(new Response(response, url)).when(hc).execute(any(Request.class)); + Mockito.doAnswer(invocation -> new Response(response, url)) + .when(hc) + .execute(any(Request.class)); Consent testConsent = new Consent.Builder() @@ -166,9 +169,21 @@ void exchangeTest() throws IOException { Map queryCreds = new HashMap<>(); queryCreds.put("code", "testCode"); queryCreds.put("state", "testState"); + + // Verify that SessionCredentials is created successfully with the exchange SessionCredentials creds = testConsent.exchangeCallbackParameters(queryCreds); - assertEquals("accessTokenFromServer", creds.token.getAccessToken()); - assertEquals("refreshTokenFromServer", creds.token.getRefreshToken()); + assertNotNull(creds); + + // Create a minimal config for testing the configure method + DatabricksConfig testConfig = new DatabricksConfig(); + + // Configure the SessionCredentials to get the OAuthHeaderFactory + OAuthHeaderFactory headerFactory = creds.configure(testConfig); + assertNotNull(headerFactory); + + // Verify the headers are correctly formatted + Map headers = headerFactory.headers(); + assertEquals("tokenTypeFromServer accessTokenFromServer", headers.get("Authorization")); } @Test @@ -177,7 +192,9 @@ void clientCredentials() throws IOException { String response = "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"10\", \"refresh_token\": \"refreshTokenFromServer\"}"; URL url = new URL("https://databricks.com/"); - Mockito.doReturn(new Response(response, url)).when(hc).execute(any(Request.class)); + Mockito.doAnswer(invocation -> new Response(response, url)) + .when(hc) + .execute(any(Request.class)); ClientCredentials clientCredentials = new ClientCredentials.Builder() @@ -197,22 +214,21 @@ void sessionCredentials() throws IOException { String response = "{\"access_token\": \"accessTokenFromServer\", \"token_type\": \"tokenTypeFromServer\", \"expires_in\": \"10\", \"refresh_token\": \"refreshTokenFromServer\"}"; URL url = new URL("https://databricks.com/"); - Mockito.doReturn(new Response(response, url)).when(hc).execute(any(Request.class)); + Mockito.doAnswer(invocation -> new Response(response, url)) + .when(hc) + .execute(any(Request.class)); - SessionCredentials sessionCredentials = - new SessionCredentials.Builder() - .withHttpClient(hc) - .withClientId("testClientId") - .withClientSecret("abc") - .withTokenUrl("https://tokenUrl") - .withToken( - new Token( - "originalAccessToken", - "originalTokenType", - "originalRefreshToken", - Instant.MAX)) - .build(); - Token token = sessionCredentials.refresh(); + SessionCredentialsTokenSource sessionCredentialsTokenSource = + new SessionCredentialsTokenSource( + new Token( + "originalAccessToken", "originalTokenType", "originalRefreshToken", Instant.MAX), + hc, + "https://tokenUrl", + "testClientId", + "abc", + Optional.empty(), + Optional.empty()); + Token token = sessionCredentialsTokenSource.refresh(); // We check that we are actually getting the token from server response (that is defined // above) rather than what was given while creating session credentials @@ -229,7 +245,7 @@ void cacheWithValidTokenTest() throws IOException { String refreshResponse = "{\"access_token\": \"refreshed_access_token\", \"token_type\": \"Bearer\", \"expires_in\": \"3600\", \"refresh_token\": \"new_refresh_token\"}"; URL url = new URL("https://test.databricks.com/"); - Mockito.doReturn(new Response(refreshResponse, url)) + Mockito.doAnswer(invocation -> new Response(refreshResponse, url)) .when(mockHttpClient) .execute(any(Request.class)); @@ -309,7 +325,7 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException { String refreshResponse = "{\"access_token\": \"refreshed_access_token\", \"token_type\": \"Bearer\", \"expires_in\": \"3600\", \"refresh_token\": \"new_refresh_token\"}"; URL url = new URL("https://test.databricks.com/"); - Mockito.doReturn(new Response(refreshResponse, url)) + Mockito.doAnswer(invocation -> new Response(refreshResponse, url)) .when(mockHttpClient) .execute(any(Request.class)); @@ -408,12 +424,15 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { "browser_refresh_token", Instant.now().plusSeconds(3600)); - SessionCredentials browserAuthCreds = - new SessionCredentials.Builder() - .withToken(browserAuthToken) - .withClientId("test-client-id") - .withTokenUrl("https://test-token-url") - .build(); + SessionCredentialsTokenSource browserAuthTokenSource = + new SessionCredentialsTokenSource( + browserAuthToken, + mockHttpClient, + "https://test-token-url", + "test-client-id", + "test-client-secret", + Optional.empty(), + Optional.empty()); // Create config with failing HTTP client and mock token cache DatabricksConfig config = @@ -431,7 +450,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - Mockito.doReturn(browserAuthCreds) + Mockito.doReturn(browserAuthTokenSource) .when(provider) .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class)); @@ -459,6 +478,9 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException { @Test void cacheWithInvalidTokensTest() throws IOException { + // Create mock HTTP client + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); + // Create completely invalid token (no refresh token) Instant pastTime = Instant.now().minusSeconds(3600); Token invalidToken = new Token("expired_access_token", "Bearer", null, pastTime); @@ -475,12 +497,15 @@ void cacheWithInvalidTokensTest() throws IOException { "browser_refresh_token", Instant.now().plusSeconds(3600)); - SessionCredentials browserAuthCreds = - new SessionCredentials.Builder() - .withToken(browserAuthToken) - .withClientId("test-client-id") - .withTokenUrl("https://test-token-url") - .build(); + SessionCredentialsTokenSource browserAuthTokenSource = + new SessionCredentialsTokenSource( + browserAuthToken, + mockHttpClient, + "https://test-token-url", + "test-client-id", + "test-client-secret", + Optional.empty(), + Optional.empty()); // Create simple config DatabricksConfig config = @@ -492,7 +517,7 @@ void cacheWithInvalidTokensTest() throws IOException { // Create our provider and mock the browser auth method ExternalBrowserCredentialsProvider provider = Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache)); - Mockito.doReturn(browserAuthCreds) + Mockito.doReturn(browserAuthTokenSource) .when(provider) .performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));