Add unit tests and comments

Author: emmyzhou-db

databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java

@@ -6,13 +6,25 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+/**
+ * The DefaultCredentialsProvider is the primary authentication handler for the Databricks SDK. It
+ * implements a chain of responsibility pattern to manage multiple authentication methods, including
+ * Personal Access Tokens (PAT), OAuth, Azure, Google, and OpenID Connect (OIDC). The provider
+ * attempts each authentication method in sequence until a valid credential is obtained.
+ */
 public class DefaultCredentialsProvider implements CredentialsProvider {
   private static final Logger LOG = LoggerFactory.getLogger(DefaultCredentialsProvider.class);
 
+  // List of credential providers that will be tried in sequence
   private List<CredentialsProvider> providers = new ArrayList<>();
 
+  // The currently selected authentication type
   private String authType = "default";
 
+  /**
+   * Internal class to associate an ID token source with a name for identification purposes. Used
+   * primarily for OIDC (OpenID Connect) authentication flows.
+   */
   private static class NamedIDTokenSource {
     private final String name;
     private final IDTokenSource idTokenSource;
@@ -33,10 +45,23 @@ public IDTokenSource getIdTokenSource() {
 
   public DefaultCredentialsProvider() {}
 
+  /**
+   * Returns the current authentication type being used
+   *
+   * @return String representing the authentication type
+   */
   public String authType() {
     return authType;
   }
 
+  /**
+   * Configures the credentials provider with the given Databricks configuration. This method tries
+   * each available credential provider in sequence until one succeeds.
+   *
+   * @param config The Databricks configuration containing authentication details
+   * @return HeaderFactory for making authenticated requests
+   * @throws DatabricksException if no valid credentials can be configured
+   */
   @Override
   public synchronized HeaderFactory configure(DatabricksConfig config) {
     addDefaultCredentialsProviders(config);
@@ -69,6 +94,11 @@ public synchronized HeaderFactory configure(DatabricksConfig config) {
             + " to configure credentials for your preferred authentication method");
   }
 
+  /**
+   * Adds OpenID Connect (OIDC) based credential providers to the list of available providers.
+   *
+   * @param config The Databricks configuration containing OIDC settings
+   */
   private void addOIDCCredentialsProviders(DatabricksConfig config) {
     OpenIDConnectEndpoints endpoints = null;
     try {
@@ -88,6 +118,7 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
     // Add new IDTokenSources and ID providers here. Example:
     // namedIdTokenSources.add(new NamedIDTokenSource("custom-oidc", new CustomIDTokenSource(...)));
 
+    // Configure OAuth token sources for each ID token source
     for (NamedIDTokenSource namedIdTokenSource : namedIdTokenSources) {
       DatabricksOAuthTokenSource oauthTokenSource =
           new DatabricksOAuthTokenSource.Builder(
@@ -105,11 +136,18 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
     }
   }
 
+  /**
+   * Initializes all available credential providers in the preferred order. The order of providers
+   * determines the authentication fallback sequence.
+   *
+   * @param config The Databricks configuration to use for provider initialization
+   */
   private void addDefaultCredentialsProviders(DatabricksConfig config) {
     providers.add(new PatCredentialsProvider());
     providers.add(new BasicCredentialsProvider());
     providers.add(new OAuthM2MServicePrincipalCredentialsProvider());
 
+    // Add OIDC-based providers
     addOIDCCredentialsProviders(config);
 
     providers.add(new AzureGithubOidcCredentialsProvider());

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubIDTokenSource.java

@@ -9,11 +9,19 @@
 import com.google.common.base.Strings;
 import java.io.IOException;
 
-/** GithubIDTokenSource retrieves JWT Tokens from GitHub Actions. */
+/**
+ * GithubIDTokenSource retrieves JWT Tokens from GitHub Actions. This class implements the
+ * IDTokenSource interface and provides a method for obtaining ID tokens specifically from GitHub
+ * Actions environment.
+ */
 public class GithubIDTokenSource implements IDTokenSource {
+  // URL endpoint for requesting ID tokens from GitHub Actions
   private final String actionsIDTokenRequestURL;
+  // Authentication token required to request ID tokens from GitHub Actions
   private final String actionsIDTokenRequestToken;
+  // HTTP client for making requests to GitHub Actions
   private final HttpClient httpClient;
+  // JSON mapper for parsing response data
   private final ObjectMapper mapper = new ObjectMapper();
 
   /**
@@ -30,8 +38,18 @@ public GithubIDTokenSource(
     this.httpClient = httpClient;
   }
 
+  /**
+   * Retrieves an ID token from GitHub Actions. This method makes an authenticated request to GitHub
+   * Actions to obtain a JWT token that later can be exchanged for a Databricks access token.
+   *
+   * @param audience Optional audience claim for the token. If provided, it will be included in the
+   *     token request to GitHub Actions.
+   * @return An IDToken object containing the JWT token value
+   * @throws DatabricksException if the token request fails or if required configuration is missing
+   */
   @Override
   public IDToken getIDToken(String audience) {
+    // Validate required configuration
     if (Strings.isNullOrEmpty(actionsIDTokenRequestURL)) {
       throw new DatabricksException("Missing ActionsIDTokenRequestURL");
     }
@@ -59,6 +77,7 @@ public IDToken getIDToken(String audience) {
           "Failed to request ID token from " + requestUrl + ": " + e.getMessage(), e);
     }
 
+    // Validate response status code
     if (resp.getStatusCode() != 200) {
       throw new DatabricksException(
           "Failed to request ID token: status code "
@@ -67,6 +86,7 @@ public IDToken getIDToken(String audience) {
               + resp.getBody().toString());
     }
 
+    // Parse the JSON response
     ObjectNode jsonResp;
     try {
       jsonResp = mapper.readValue(resp.getBody(), ObjectNode.class);
@@ -75,6 +95,7 @@ public IDToken getIDToken(String audience) {
           "Failed to request ID token: corrupted token: " + e.getMessage());
     }
 
+    // Validate response structure and token value
     if (!jsonResp.has("value")) {
       throw new DatabricksException("ID token response missing 'value' field");
     }

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProvider.java

@@ -6,22 +6,38 @@
 import java.util.HashMap;
 import java.util.Map;
 
-/** Base class for token-based credentials providers. */
+/**
+ * A credentials provider that uses a TokenSource to obtain and manage authentication tokens. This
+ * class serves as a base implementation for token-based authentication, handling the conversion of
+ * tokens into HTTP authorization headers.
+ *
+ * <p>The provider validates token availability during configuration and creates. appropriate
+ * authorization headers for API requests.
+ */
 public class TokenSourceCredentialsProvider implements CredentialsProvider {
   private final TokenSource tokenSource;
   private final String authType;
 
   /**
    * Creates a new TokenSourceCredentialsProvider with the specified token source and auth type.
    *
-   * @param tokenSource The token source to use for token exchange
-   * @param authType The authentication type string
+   * @param tokenSource The token source responsible for token acquisition and management.
+   * @param authType The authentication type identifier.
    */
   public TokenSourceCredentialsProvider(TokenSource tokenSource, String authType) {
     this.tokenSource = tokenSource;
     this.authType = authType;
   }
 
+  /**
+   * Configures the credentials provider and creates a HeaderFactory for generating authentication
+   * headers. This method validates token availability by attempting to obtain an access token
+   * before returning the HeaderFactory.
+   *
+   * @param config The Databricks configuration object.
+   * @return A HeaderFactory that generates "Bearer" token authorization headers, or null if token
+   *     acquisition fails.
+   */
   @Override
   public HeaderFactory configure(DatabricksConfig config) {
     try {
@@ -38,6 +54,12 @@ public HeaderFactory configure(DatabricksConfig config) {
     }
   }
 
+  /**
+   * Returns the authentication type identifier for this credentials provider. This is used to
+   * identify the authentication method being used.
+   *
+   * @return The authentication type string
+   */
   @Override
   public String authType() {
     return authType;

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/TokenSourceCredentialsProviderTest.java

@@ -0,0 +1,72 @@
+package com.databricks.sdk.core.oauth;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+import com.databricks.sdk.core.DatabricksConfig;
+import com.databricks.sdk.core.DatabricksException;
+import com.databricks.sdk.core.HeaderFactory;
+import java.time.LocalDateTime;
+import java.util.Map;
+import java.util.stream.Stream;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+/** Tests for TokenSourceCredentialsProvider */
+class TokenSourceCredentialsProviderTest {
+  private TokenSourceCredentialsProvider provider;
+  private static final String TEST_AUTH_TYPE = "test-auth-type";
+  private static final String TEST_ACCESS_TOKEN_VALUE = "test-access-token";
+  private static final Token TEST_TOKEN =
+      new Token(TEST_ACCESS_TOKEN_VALUE, "Bearer", LocalDateTime.now().plusHours(1));
+
+  /** Tests token retrieval scenarios */
+  @ParameterizedTest(name = "{0}")
+  @MethodSource("provideTokenScenarios")
+  void testTokenScenarios(
+      String testName,
+      TokenSource mockTokenSource,
+      String expectedAuthHeader,
+      Exception expectedException) {
+    provider = new TokenSourceCredentialsProvider(mockTokenSource, TEST_AUTH_TYPE);
+    DatabricksConfig config = new DatabricksConfig();
+    HeaderFactory headerFactory = provider.configure(config);
+
+    if (expectedException != null) {
+      assertNull(headerFactory);
+    } else {
+      assertNotNull(headerFactory);
+      Map<String, String> headers = headerFactory.headers();
+      assertEquals(expectedAuthHeader, headers.get("Authorization"));
+    }
+    verify(mockTokenSource).getToken();
+    assertEquals(TEST_AUTH_TYPE, provider.authType());
+  }
+
+  /** Provides test scenarios */
+  private static Stream<Arguments> provideTokenScenarios() {
+    // Mock behaviour of successful token retrieval
+    TokenSource mockSuccessTokenSource = mock(TokenSource.class);
+    when(mockSuccessTokenSource.getToken()).thenReturn(TEST_TOKEN);
+
+    // Mock behaviour of failing token retrieval
+    TokenSource mockFailureTokenSource = mock(TokenSource.class);
+    when(mockFailureTokenSource.getToken())
+        .thenThrow(new DatabricksException("Token retrieval failed"));
+
+    return Stream.of(
+        // Success case
+        Arguments.of(
+            "Successful token retrieval",
+            mockSuccessTokenSource,
+            "Bearer " + TEST_ACCESS_TOKEN_VALUE,
+            null),
+        // Failure case
+        Arguments.of(
+            "Failed token retrieval",
+            mockFailureTokenSource,
+            null,
+            new DatabricksException("Token retrieval failed")));
+  }
+}

examples/docs/src/main/java/com/databricks/example/GithubOIDCAuthExample.java

@@ -8,20 +8,26 @@
 
 /**
  * Example demonstrating how to use GitHub OIDC authentication with Databricks.
+ * 
+ * IMPORTANT: This example only works when running within GitHub Actions.
  */
 public class GithubOIDCAuthExample {
 
     public static void main(String[] args) {
         // Create Databricks configuration with GitHub OIDC authentication
+        // Note: This configuration assumes the code is running in GitHub Actions
         DatabricksConfig config = new DatabricksConfig()
-        .setAuthType("github-oidc")
-        .setHost("https://accounts.cloud.databricks.com/")
-        .setAccountId("968367da-7edd-44f7-9dea-3e0b20b0ec97")
-        .setClientId("dd3b5d58-5a6e-45e3-a3ae-81d8f89fc6fd");
+            .setAuthType("github-oidc")  // Specifies GitHub OIDC as the authentication method
+            .setHost("<INSERT_HOST>")  // Databricks account URL
+            .setAccountId("<INSERT_ACCOUNT_ID>")  // Your Databricks account ID
+            .setClientId("<INSERT_CLIENT_ID>");  // Service Principal ID
 
+        // Initialize the Account client with the OIDC configuration
         AccountClient account = new AccountClient(config);
 
         try {
+            // Example: List all groups in the Databricks account
+            // This demonstrates that the OIDC authentication is working
             System.out.println("\nListing account groups:");
             Iterable<Group> groups = account.groups().list(new ListAccountGroupsRequest());
             for (Group group : groups) {