diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java index 307fc2cb2..c8ac65ba2 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java @@ -10,27 +10,45 @@ import java.time.LocalDateTime; import java.util.HashMap; import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Implementation of TokenSource that handles OAuth token exchange for Databricks authentication. * This class manages the OAuth token exchange flow using ID tokens to obtain access tokens. */ public class DatabricksOAuthTokenSource implements TokenSource { - /** OAuth client ID used for token exchange */ + private static final Logger LOG = LoggerFactory.getLogger(DatabricksOAuthTokenSource.class); + + /** OAuth client ID used for token exchange. */ private final String clientId; - /** Databricks account ID, used as audience if provided */ + /** Databricks host URL. */ + private final String host; + /** Databricks account ID, used as audience if provided. */ private final String accountId; - /** OpenID Connect endpoints configuration */ + /** OpenID Connect endpoints configuration. */ private final OpenIDConnectEndpoints endpoints; - /** Custom audience value for token exchange */ + /** Custom audience value for token exchange. */ private final String audience; - /** Source of ID tokens used in token exchange */ + /** Source of ID tokens used in token exchange. */ private final IDTokenSource idTokenSource; - /** HTTP client for making token exchange requests */ + /** HTTP client for making token exchange requests. */ private final HttpClient httpClient; + private static final String GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"; + private static final String SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:jwt"; + private static final String SCOPE = "all-apis"; + private static final String GRANT_TYPE_PARAM = "grant_type"; + private static final String SUBJECT_TOKEN_PARAM = "subject_token"; + private static final String SUBJECT_TOKEN_TYPE_PARAM = "subject_token_type"; + private static final String SCOPE_PARAM = "scope"; + private static final String CLIENT_ID_PARAM = "client_id"; + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private DatabricksOAuthTokenSource(Builder builder) { this.clientId = builder.clientId; + this.host = builder.host; this.accountId = builder.accountId; this.endpoints = builder.endpoints; this.audience = builder.audience; @@ -39,8 +57,8 @@ private DatabricksOAuthTokenSource(Builder builder) { } /** - * Builder class for constructing DatabricksOAuthTokenSource instances. Provides a fluent - * interface for setting required and optional parameters. + * Builder class for constructing DatabricksOAuthTokenSource instances. Provides a flexible way to + * set required and optional parameters. */ public static class Builder { private final String clientId; @@ -51,30 +69,14 @@ public static class Builder { private String accountId; private String audience; - /** - * Validates that a value is non-empty and non-null for required fields. - * - * @param value The value to validate - * @param fieldName The name of the field being validated - * @throws IllegalArgumentException if validation fails - */ - private static void validate(Object value, String fieldName) { - if (value == null) { - throw new IllegalArgumentException(fieldName + " must be non-null"); - } - if (value instanceof String && ((String) value).isEmpty()) { - throw new IllegalArgumentException(fieldName + " must be non-empty"); - } - } - /** * Creates a new Builder with required parameters. * - * @param clientId OAuth client ID - * @param host Databricks host URL - * @param endpoints OpenID Connect endpoints configuration - * @param idTokenSource Source of ID tokens - * @param httpClient HTTP client for making requests + * @param clientId OAuth client ID. + * @param host Databricks host URL. + * @param endpoints OpenID Connect endpoints configuration. + * @param idTokenSource Source of ID tokens. + * @param httpClient HTTP client for making requests. */ public Builder( String clientId, @@ -82,12 +84,6 @@ public Builder( OpenIDConnectEndpoints endpoints, IDTokenSource idTokenSource, HttpClient httpClient) { - validate(clientId, "ClientID"); - validate(host, "Host"); - validate(endpoints, "Endpoints"); - validate(idTokenSource, "IDTokenSource"); - validate(httpClient, "HttpClient"); - this.clientId = clientId; this.host = host; this.endpoints = endpoints; @@ -98,11 +94,10 @@ public Builder( /** * Sets the Databricks account ID. * - * @param accountId The account ID - * @return This builder instance + * @param accountId The account ID. + * @return This builder instance. */ public Builder accountId(String accountId) { - validate(accountId, "AccountID"); this.accountId = accountId; return this; } @@ -114,7 +109,6 @@ public Builder accountId(String accountId) { * @return This builder instance */ public Builder audience(String audience) { - validate(audience, "Audience"); this.audience = audience; return this; } @@ -122,36 +116,71 @@ public Builder audience(String audience) { /** * Builds a new DatabricksOAuthTokenSource instance. * - * @return A new DatabricksOAuthTokenSource + * @return A new DatabricksOAuthTokenSource. */ public DatabricksOAuthTokenSource build() { return new DatabricksOAuthTokenSource(this); } } + /** + * Validates that a value is non-null for required fields. If the value is a string, it also + * checks that it is non-empty. + * + * @param value The value to validate. + * @param fieldName The name of the field being validated. + * @throws IllegalArgumentException when the value is null or an empty string. + */ + private static void validate(Object value, String fieldName) { + if (value == null) { + LOG.error("Required parameter '{}' is null", fieldName); + throw new IllegalArgumentException( + String.format("Required parameter '%s' cannot be null", fieldName)); + } + if (value instanceof String && ((String) value).isEmpty()) { + LOG.error("Required parameter '{}' is empty", fieldName); + throw new IllegalArgumentException( + String.format("Required parameter '%s' cannot be empty", fieldName)); + } + } + /** * Retrieves an OAuth token by exchanging an ID token. Implements the OAuth token exchange flow to * obtain an access token. * - * @return A Token containing the access token and related information - * @throws DatabricksException if token exchange fails + * @return A Token containing the access token and related information. + * @throws DatabricksException when the token exchange fails. + * @throws IllegalArgumentException when there is an error code in the response or when required + * parameters are missing. */ @Override public Token getToken() { + // Validate all required parameters + validate(clientId, "ClientID"); + validate(host, "Host"); + validate(endpoints, "Endpoints"); + validate(idTokenSource, "IDTokenSource"); + validate(httpClient, "HttpClient"); + String effectiveAudience = determineAudience(); IDToken idToken = idTokenSource.getIDToken(effectiveAudience); Map params = new HashMap<>(); - params.put("grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"); - params.put("subject_token", idToken.getValue()); - params.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt"); - params.put("scope", "all-apis"); - params.put("client_id", clientId); + params.put(GRANT_TYPE_PARAM, GRANT_TYPE); + params.put(SUBJECT_TOKEN_PARAM, idToken.getValue()); + params.put(SUBJECT_TOKEN_TYPE_PARAM, SUBJECT_TOKEN_TYPE); + params.put(SCOPE_PARAM, SCOPE); + params.put(CLIENT_ID_PARAM, clientId); Response rawResponse; try { rawResponse = httpClient.execute(new FormRequest(endpoints.getTokenEndpoint(), params)); } catch (IOException e) { + LOG.error( + "Failed to exchange ID token for access token at {}: {}", + endpoints.getTokenEndpoint(), + e.getMessage(), + e); throw new DatabricksException( String.format( "Failed to exchange ID token for access token at %s: %s", @@ -161,8 +190,13 @@ public Token getToken() { OAuthResponse response; try { - response = new ObjectMapper().readValue(rawResponse.getBody(), OAuthResponse.class); + response = OBJECT_MAPPER.readValue(rawResponse.getBody(), OAuthResponse.class); } catch (IOException e) { + LOG.error( + "Failed to parse OAuth response from token endpoint {}: {}", + endpoints.getTokenEndpoint(), + e.getMessage(), + e); throw new DatabricksException( String.format( "Failed to parse OAuth response from token endpoint %s: %s", @@ -170,6 +204,10 @@ public Token getToken() { } if (response.getErrorCode() != null) { + LOG.error( + "Token exchange failed with error: {} - {}", + response.getErrorCode(), + response.getErrorSummary()); throw new IllegalArgumentException( String.format( "Token exchange failed with error: %s - %s", diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java index 3ffc9e71a..8d7da8d3a 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java @@ -10,15 +10,14 @@ import com.databricks.sdk.core.http.Response; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.IOException; +import java.net.MalformedURLException; import java.net.URL; import java.util.HashMap; import java.util.Map; import java.util.stream.Stream; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; -import org.mockito.ArgumentCaptor; import org.mockito.Mockito; class DatabricksOAuthTokenSourceTest { @@ -36,6 +35,10 @@ class DatabricksOAuthTokenSourceTest { private static final String TEST_AUDIENCE = "test-audience"; private static final String TEST_ACCOUNT_ID = "test-account-id"; + // Error message constants + private static final String ERROR_NULL = "Required parameter '%s' cannot be null"; + private static final String ERROR_EMPTY = "Required parameter '%s' cannot be empty"; + private IDTokenSource mockIdTokenSource; @BeforeEach @@ -54,25 +57,24 @@ private static class TestCase { final String audience; // Custom audience value if provided final String accountId; // Account ID if provided final String expectedAudience; // Expected audience used in token exchange - final boolean expectError; // Whether this case should result in an error - final int statusCode; // HTTP status code for the response - final String responseBody; // Response body from the token endpoint + final HttpClient mockHttpClient; // Pre-configured mock HTTP client + final Class expectedException; // Expected exception type if any TestCase( String name, String audience, String accountId, String expectedAudience, - boolean expectError, int statusCode, - String responseBody) { + Object responseBody, + HttpClient mockHttpClient, + Class expectedException) { this.name = name; this.audience = audience; this.accountId = accountId; this.expectedAudience = expectedAudience; - this.expectError = expectError; - this.statusCode = statusCode; - this.responseBody = responseBody; + this.mockHttpClient = mockHttpClient; + this.expectedException = expectedException; } @Override @@ -86,75 +88,124 @@ public String toString() { * audience configurations and various error cases. */ private static Stream provideTestCases() { - // Success response with valid token data - Map successResponse = new HashMap<>(); - successResponse.put("access_token", TOKEN); - successResponse.put("token_type", TOKEN_TYPE); - successResponse.put("refresh_token", REFRESH_TOKEN); - successResponse.put("expires_in", EXPIRES_IN); + try { + // Success response with valid token data + Map successResponse = new HashMap<>(); + successResponse.put("access_token", TOKEN); + successResponse.put("token_type", TOKEN_TYPE); + successResponse.put("refresh_token", REFRESH_TOKEN); + successResponse.put("expires_in", EXPIRES_IN); + + // Error response for invalid requests + Map errorResponse = new HashMap<>(); + errorResponse.put("error", "invalid_request"); + errorResponse.put("error_description", "Invalid client ID"); + + ObjectMapper mapper = new ObjectMapper(); + final String errorJson = mapper.writeValueAsString(errorResponse); + final String successJson = mapper.writeValueAsString(successResponse); - // Error response for invalid requests - Map errorResponse = new HashMap<>(); - errorResponse.put("error", "invalid_request"); - errorResponse.put("error_description", "Invalid client ID"); + // Create the expected request that will be used in all test cases + Map formParams = new HashMap<>(); + formParams.put("client_id", TEST_CLIENT_ID); + formParams.put("subject_token", TEST_ID_TOKEN); + formParams.put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt"); + formParams.put("grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"); + formParams.put("scope", "all-apis"); + FormRequest expectedRequest = new FormRequest(TEST_TOKEN_ENDPOINT, formParams); + + return Stream.of( + // Success cases with different audience configurations + new TestCase( + "Default audience from token endpoint", + null, + null, + TEST_TOKEN_ENDPOINT, + 200, + successResponse, + createMockHttpClient(expectedRequest, 200, successJson), + null), + new TestCase( + "Custom audience provided", + TEST_AUDIENCE, + null, + TEST_AUDIENCE, + 200, + successResponse, + createMockHttpClient(expectedRequest, 200, successJson), + null), + new TestCase( + "Custom audience takes precedence over account ID", + TEST_AUDIENCE, + TEST_ACCOUNT_ID, + TEST_AUDIENCE, + 200, + successResponse, + createMockHttpClient(expectedRequest, 200, successJson), + null), + new TestCase( + "Account ID used as audience when no custom audience", + null, + TEST_ACCOUNT_ID, + TEST_ACCOUNT_ID, + 200, + successResponse, + createMockHttpClient(expectedRequest, 200, successJson), + null), + // Error cases + new TestCase( + "Invalid request returns 400", + null, + null, + TEST_TOKEN_ENDPOINT, + 400, + errorJson, + createMockHttpClient(expectedRequest, 400, errorJson), + IllegalArgumentException.class), + new TestCase( + "Network error during token exchange", + null, + null, + TEST_TOKEN_ENDPOINT, + 0, + null, + createMockHttpClientWithError(expectedRequest), + DatabricksException.class), + new TestCase( + "Invalid JSON response from server", + null, + null, + TEST_TOKEN_ENDPOINT, + 200, + "invalid json", + createMockHttpClient(expectedRequest, 200, "invalid json"), + DatabricksException.class)); + } catch (IOException e) { + throw new RuntimeException("Failed to create test cases", e); + } + } - ObjectMapper mapper = new ObjectMapper(); - String successJson; - String errorJson; + private static HttpClient createMockHttpClient( + FormRequest expectedRequest, int statusCode, String responseBody) { try { - successJson = mapper.writeValueAsString(successResponse); - errorJson = mapper.writeValueAsString(errorResponse); + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); + String statusMessage = statusCode == 200 ? "OK" : "Bad Request"; + when(mockHttpClient.execute(expectedRequest)) + .thenReturn(new Response(responseBody, statusCode, statusMessage, new URL(TEST_HOST))); + return mockHttpClient; } catch (IOException e) { - throw new RuntimeException(e); + throw new RuntimeException("Failed to create mock HTTP client", e); } + } - return Stream.of( - // Success cases with different audience configurations - new TestCase( - "Default audience from token endpoint", - null, - null, - TEST_TOKEN_ENDPOINT, - false, - 200, - successJson), - new TestCase( - "Custom audience provided", - TEST_AUDIENCE, - null, - TEST_AUDIENCE, - false, - 200, - successJson), - new TestCase( - "Custom audience takes precedence over account ID", - TEST_AUDIENCE, - TEST_ACCOUNT_ID, - TEST_AUDIENCE, - false, - 200, - successJson), - new TestCase( - "Account ID used as audience when no custom audience", - null, - TEST_ACCOUNT_ID, - TEST_ACCOUNT_ID, - false, - 200, - successJson), - // Error cases - new TestCase( - "Invalid request returns 400", null, null, TEST_TOKEN_ENDPOINT, true, 400, errorJson), - new TestCase( - "Network error during token exchange", null, null, TEST_TOKEN_ENDPOINT, true, 0, null), - new TestCase( - "Invalid JSON response from server", - null, - null, - TEST_TOKEN_ENDPOINT, - true, - 200, - "invalid json")); + private static HttpClient createMockHttpClientWithError(FormRequest expectedRequest) { + try { + HttpClient mockHttpClient = Mockito.mock(HttpClient.class); + when(mockHttpClient.execute(expectedRequest)).thenThrow(new IOException("Network error")); + return mockHttpClient; + } catch (IOException e) { + throw new RuntimeException("Failed to create mock HTTP client with error", e); + } } /** @@ -163,171 +214,177 @@ private static Stream provideTestCases() { */ @ParameterizedTest(name = "testTokenSource: {arguments}") @MethodSource("provideTestCases") - void testTokenSource(TestCase testCase) throws IOException { - // Mock HTTP client with test case specific behavior - HttpClient mockHttpClient = Mockito.mock(HttpClient.class); - if (testCase.expectError) { - if (testCase.statusCode == 0) { - when(mockHttpClient.execute(any())).thenThrow(new IOException("Network error")); - } else { - when(mockHttpClient.execute(any())) - .thenReturn( - new Response( - testCase.responseBody, testCase.statusCode, "Bad Request", new URL(TEST_HOST))); - } - } else { - when(mockHttpClient.execute(any())) - .thenReturn(new Response(testCase.responseBody, new URL(TEST_HOST))); - } - - // Create token source with test configuration - OpenIDConnectEndpoints endpoints = - new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT); + void testTokenSource(TestCase testCase) { + try { + // Create token source with test configuration + OpenIDConnectEndpoints endpoints = + new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT); - DatabricksOAuthTokenSource.Builder builder = - new DatabricksOAuthTokenSource.Builder( - TEST_CLIENT_ID, TEST_HOST, endpoints, mockIdTokenSource, mockHttpClient); + DatabricksOAuthTokenSource.Builder builder = + new DatabricksOAuthTokenSource.Builder( + TEST_CLIENT_ID, TEST_HOST, endpoints, mockIdTokenSource, testCase.mockHttpClient); - if (testCase.audience != null) { - builder.audience(testCase.audience); - } - if (testCase.accountId != null) { - builder.accountId(testCase.accountId); - } + builder.audience(testCase.audience).accountId(testCase.accountId); - DatabricksOAuthTokenSource tokenSource = builder.build(); + DatabricksOAuthTokenSource tokenSource = builder.build(); - if (testCase.expectError) { - if (testCase.statusCode == 400) { - assertThrows(IllegalArgumentException.class, () -> tokenSource.getToken()); + if (testCase.expectedException != null) { + assertThrows(testCase.expectedException, () -> tokenSource.getToken()); } else { - assertThrows(DatabricksException.class, () -> tokenSource.getToken()); - } - } else { - // Verify successful token exchange - Token token = tokenSource.getToken(); - assertEquals(TOKEN, token.getAccessToken()); - assertEquals(TOKEN_TYPE, token.getTokenType()); - assertEquals(REFRESH_TOKEN, token.getRefreshToken()); - assertFalse(token.isExpired()); - - // Verify correct audience was used - verify(mockIdTokenSource).getIDToken(testCase.expectedAudience); + // Verify successful token exchange + Token token = tokenSource.getToken(); + assertEquals(TOKEN, token.getAccessToken()); + assertEquals(TOKEN_TYPE, token.getTokenType()); + assertEquals(REFRESH_TOKEN, token.getRefreshToken()); + assertFalse(token.isExpired()); - // Verify token exchange request - ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(FormRequest.class); - verify(mockHttpClient).execute(requestCaptor.capture()); - - FormRequest capturedRequest = requestCaptor.getValue(); - assertEquals(TEST_TOKEN_ENDPOINT, capturedRequest.getUrl()); - - // Verify request parameters - String body = capturedRequest.getBodyString(); - assertTrue(body.contains("client_id=" + TEST_CLIENT_ID)); - assertTrue(body.contains("subject_token=" + TEST_ID_TOKEN)); - assertTrue( - body.contains("subject_token_type=urn%3Aietf%3Aparams%3Aoauth%3Atoken-type%3Ajwt")); - assertTrue( - body.contains("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Atoken-exchange")); - assertTrue(body.contains("scope=all-apis")); + // Verify correct audience was used + verify(mockIdTokenSource).getIDToken(testCase.expectedAudience); + } + } catch (IOException e) { + throw new RuntimeException("Test failed", e); } } /** - * Tests validation of required fields in the token source builder. Verifies that null or empty - * values for required fields throw IllegalArgumentException. + * Test case data for parameter validation tests. Each case defines a specific validation + * scenario. */ - @Test - void testConstructorValidation() { - // Test null client ID - assertThrows( - IllegalArgumentException.class, - () -> { - new DatabricksOAuthTokenSource.Builder( - null, - TEST_HOST, - new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT), - mockIdTokenSource, - Mockito.mock(HttpClient.class)) - .build(); - }); + private static class ValidationTestCase { + final String name; + final String clientId; + final String host; + final OpenIDConnectEndpoints endpoints; + final IDTokenSource idTokenSource; + final HttpClient httpClient; + final String expectedFieldName; + final boolean isNullTest; - // Test empty client ID - assertThrows( - IllegalArgumentException.class, - () -> { - new DatabricksOAuthTokenSource.Builder( - "", - TEST_HOST, - new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT), - mockIdTokenSource, - Mockito.mock(HttpClient.class)) - .build(); - }); + ValidationTestCase( + String name, + String clientId, + String host, + OpenIDConnectEndpoints endpoints, + IDTokenSource idTokenSource, + HttpClient httpClient, + String expectedFieldName, + boolean isNullTest) { + this.name = name; + this.clientId = clientId; + this.host = host; + this.endpoints = endpoints; + this.idTokenSource = idTokenSource; + this.httpClient = httpClient; + this.expectedFieldName = expectedFieldName; + this.isNullTest = isNullTest; + } - // Test null host - assertThrows( - IllegalArgumentException.class, - () -> { - new DatabricksOAuthTokenSource.Builder( - TEST_CLIENT_ID, - null, - new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT), - mockIdTokenSource, - Mockito.mock(HttpClient.class)) - .build(); - }); + @Override + public String toString() { + return name; + } + } - // Test empty host - assertThrows( - IllegalArgumentException.class, - () -> { - new DatabricksOAuthTokenSource.Builder( - TEST_CLIENT_ID, - "", - new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT), - mockIdTokenSource, - Mockito.mock(HttpClient.class)) - .build(); - }); + private static Stream provideValidationTestCases() + throws MalformedURLException { + OpenIDConnectEndpoints validEndpoints = + new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT); + HttpClient validHttpClient = Mockito.mock(HttpClient.class); + IDTokenSource validIdTokenSource = Mockito.mock(IDTokenSource.class); - // Test null endpoints - assertThrows( - IllegalArgumentException.class, - () -> { - new DatabricksOAuthTokenSource.Builder( - TEST_CLIENT_ID, - TEST_HOST, - null, - mockIdTokenSource, - Mockito.mock(HttpClient.class)) - .build(); - }); + return Stream.of( + // Client ID validation + new ValidationTestCase( + "Null client ID", + null, + TEST_HOST, + validEndpoints, + validIdTokenSource, + validHttpClient, + "ClientID", + true), + new ValidationTestCase( + "Empty client ID", + "", + TEST_HOST, + validEndpoints, + validIdTokenSource, + validHttpClient, + "ClientID", + false), + // Host validation + new ValidationTestCase( + "Null host", + TEST_CLIENT_ID, + null, + validEndpoints, + validIdTokenSource, + validHttpClient, + "Host", + true), + new ValidationTestCase( + "Empty host", + TEST_CLIENT_ID, + "", + validEndpoints, + validIdTokenSource, + validHttpClient, + "Host", + false), + // Endpoints validation + new ValidationTestCase( + "Null endpoints", + TEST_CLIENT_ID, + TEST_HOST, + null, + validIdTokenSource, + validHttpClient, + "Endpoints", + true), + // IDTokenSource validation + new ValidationTestCase( + "Null IDTokenSource", + TEST_CLIENT_ID, + TEST_HOST, + validEndpoints, + null, + validHttpClient, + "IDTokenSource", + true), + // HttpClient validation + new ValidationTestCase( + "Null HttpClient", + TEST_CLIENT_ID, + TEST_HOST, + validEndpoints, + validIdTokenSource, + null, + "HttpClient", + true)); + } - // Test null IDTokenSource - assertThrows( - IllegalArgumentException.class, - () -> { - new DatabricksOAuthTokenSource.Builder( - TEST_CLIENT_ID, - TEST_HOST, - new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT), - null, - Mockito.mock(HttpClient.class)) - .build(); - }); + /** + * Tests validation of required fields in the token source using parameterized test cases. + * Verifies that null or empty values for required fields cause getToken() to throw + * IllegalArgumentException with specific error messages. + */ + @ParameterizedTest(name = "testParameterValidation: {0}") + @MethodSource("provideValidationTestCases") + void testParameterValidation(ValidationTestCase testCase) { + DatabricksOAuthTokenSource tokenSource = + new DatabricksOAuthTokenSource.Builder( + testCase.clientId, + testCase.host, + testCase.endpoints, + testCase.idTokenSource, + testCase.httpClient) + .build(); - // Test null HttpClient - assertThrows( - IllegalArgumentException.class, - () -> { - new DatabricksOAuthTokenSource.Builder( - TEST_CLIENT_ID, - TEST_HOST, - new OpenIDConnectEndpoints(TEST_TOKEN_ENDPOINT, TEST_AUTHORIZATION_ENDPOINT), - mockIdTokenSource, - null) - .build(); - }); + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> tokenSource.getToken()); + + String expectedMessage = + String.format(testCase.isNullTest ? ERROR_NULL : ERROR_EMPTY, testCase.expectedFieldName); + assertEquals(expectedMessage, exception.getMessage()); } }