diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index ee2122a8a..da376b995 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -4,6 +4,8 @@ ### New Features and Improvements +* Add native support for Azure DevOps OIDC authentication. + ### Bug Fixes ### Documentation diff --git a/README.md b/README.md index 8a8d2858a..4e6d6af2e 100644 --- a/README.md +++ b/README.md @@ -116,10 +116,11 @@ Depending on the Databricks authentication method, the SDK uses the following in ### Databricks native authentication -By default, the Databricks SDK for Java initially tries [Databricks token authentication](https://docs.databricks.com/dev-tools/api/latest/authentication.html) (`auth_type='pat'` argument). If the SDK is unsuccessful, it then tries Databricks Workload Identity Federation (WIF) authentication using OIDC (`auth_type="github-oidc"` argument). +By default, the Databricks SDK for Java initially tries [Databricks token authentication](https://docs.databricks.com/dev-tools/api/latest/authentication.html) (`auth_type='pat'` argument). If the SDK is unsuccessful, it then tries Workload Identity Federation (WIF). See [Supported WIF](https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation-provider) for the supported JWT token providers. - For Databricks token authentication, you must provide `host` and `token`; or their environment variable or `.databrickscfg` file field equivalents. - For Databricks OIDC authentication, you must provide the `host`, `client_id` and `token_audience` _(optional)_ either directly, through the corresponding environment variables, or in your `.databrickscfg` configuration file. +- For Azure DevOps OIDC authentication, the `token_audience` is irrelevant as the audience is always set to `api://AzureADTokenExchange`. Also, the `System.AccessToken` pipeline variable required for OIDC request must be exposed as the `SYSTEM_ACCESSTOKEN` environment variable, following [Pipeline variables](https://learn.microsoft.com/en-us/azure/devops/pipelines/build/variables?view=azure-devops&tabs=yaml#systemaccesstoken) | Argument | Description | Environment variable | |--------------|-------------|-------------------| diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index 5231a8afc..9cd580a43 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -125,6 +125,16 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) { config.getActionsIdTokenRequestUrl(), config.getActionsIdTokenRequestToken(), config.getHttpClient()))); + + // Try to create Azure DevOps token source - if environment variables are missing, + // skip this provider gracefully. + try { + namedIdTokenSources.add( + new NamedIDTokenSource( + "azure-devops-oidc", new AzureDevOpsIDTokenSource(config.getHttpClient()))); + } catch (DatabricksException e) { + LOG.debug("Azure DevOps OIDC provider not available: {}", e.getMessage()); + } // Add new IDTokenSources and ID providers here. Example: // namedIdTokenSources.add(new NamedIDTokenSource("custom-oidc", new CustomIDTokenSource(...))); diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureDevOpsIDTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureDevOpsIDTokenSource.java new file mode 100644 index 000000000..0e1e747af --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/AzureDevOpsIDTokenSource.java @@ -0,0 +1,177 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Request; +import com.databricks.sdk.core.http.Response; +import com.databricks.sdk.core.utils.Environment; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.common.base.Strings; +import java.io.IOException; + +/** + * AzureDevOpsIDTokenSource retrieves JWT Tokens from Azure DevOps Pipelines. This class implements + * the IDTokenSource interface and provides a method for obtaining ID tokens specifically from Azure + * DevOps Pipeline environment. + * + *

This implementation relies on the Azure + * DevOps OIDC token API. + */ +public class AzureDevOpsIDTokenSource implements IDTokenSource { + /* Access token for authenticating with Azure DevOps API */ + private final String azureDevOpsAccessToken; + /* Team Foundation Collection URI (e.g., https://dev.azure.com/organization) */ + private final String azureDevOpsTeamFoundationCollectionUri; + /* Plan ID for the current pipeline run */ + private final String azureDevOpsPlanId; + /* Job ID for the current pipeline job */ + private final String azureDevOpsJobId; + /* Team Project ID where the pipeline is running */ + private final String azureDevOpsTeamProjectId; + /* Host type (e.g., "build", "release") */ + private final String azureDevOpsHostType; + /* HTTP client for making requests to Azure DevOps */ + private final HttpClient httpClient; + /* Environment for reading configuration values */ + private final Environment environment; + /* JSON mapper for parsing response data */ + private static final ObjectMapper mapper = new ObjectMapper(); + + /** + * Constructs a new AzureDevOpsIDTokenSource by reading environment variables. This constructor + * implements fail-early validation - if any required environment variables are missing, it will + * throw a DatabricksException immediately. + * + * @param httpClient The HTTP client to use for making requests + * @throws DatabricksException if any required environment variables are missing + */ + public AzureDevOpsIDTokenSource(HttpClient httpClient) { + this(httpClient, createDefaultEnvironment()); + } + + /** + * Constructs a new AzureDevOpsIDTokenSource with a custom environment. This constructor is + * primarily used for testing to inject mock environment variables. + * + * @param httpClient The HTTP client to use for making requests + * @param environment The environment to read configuration from + * @throws DatabricksException if httpClient is null or any required environment variables are + * missing + */ + protected AzureDevOpsIDTokenSource(HttpClient httpClient, Environment environment) { + if (httpClient == null) { + throw new DatabricksException("HttpClient cannot be null"); + } + this.httpClient = httpClient; + this.environment = environment; + + this.azureDevOpsAccessToken = validateEnvironmentVariable("SYSTEM_ACCESSTOKEN"); + this.azureDevOpsTeamFoundationCollectionUri = + validateEnvironmentVariable("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI"); + this.azureDevOpsPlanId = validateEnvironmentVariable("SYSTEM_PLANID"); + this.azureDevOpsJobId = validateEnvironmentVariable("SYSTEM_JOBID"); + this.azureDevOpsTeamProjectId = validateEnvironmentVariable("SYSTEM_TEAMPROJECTID"); + this.azureDevOpsHostType = validateEnvironmentVariable("SYSTEM_HOSTTYPE"); + } + + /** Creates a default Environment using system environment variables. */ + private static Environment createDefaultEnvironment() { + String pathEnv = System.getenv("PATH"); + String[] pathArray = + pathEnv != null ? pathEnv.split(java.io.File.pathSeparator) : new String[0]; + return new Environment(System.getenv(), pathArray, System.getProperty("os.name")); + } + + /** + * Validates that an environment variable is present and not empty. + * + * @param varName The environment variable name + * @return The environment variable value + * @throws DatabricksException if the environment variable is missing or empty + */ + private String validateEnvironmentVariable(String varName) { + String value = environment.get(varName); + if (Strings.isNullOrEmpty(value)) { + if (varName.equals("SYSTEM_ACCESSTOKEN")) { + throw new DatabricksException( + String.format( + "Missing environment variable %s, if calling from Azure DevOps Pipeline, please set this env var following https://learn.microsoft.com/en-us/azure/devops/pipelines/build/variables?view=azure-devops&tabs=yaml#systemaccesstoken", + varName)); + } + throw new DatabricksException( + String.format( + "Missing environment variable %s, likely not calling from Azure DevOps Pipeline", + varName)); + } + return value; + } + + /** + * Retrieves an ID token from Azure DevOps Pipelines. This method makes an authenticated request + * to Azure DevOps to obtain a JWT token that can later be exchanged for a Databricks access + * token. + * + *

Note: The audience parameter is ignored for Azure DevOps OIDC tokens as they have a + * hardcoded audience for Azure AD integration. + * + * @param audience Ignored for Azure DevOps OIDC tokens + * @return An IDToken object containing the JWT token value + * @throws DatabricksException if the token request fails + */ + @Override + public IDToken getIDToken(String audience) { + + // Build Azure DevOps OIDC endpoint URL. + String requestUrl = + String.format( + "%s/%s/_apis/distributedtask/hubs/%s/plans/%s/jobs/%s/oidctoken?api-version=7.2-preview.1", + azureDevOpsTeamFoundationCollectionUri, + azureDevOpsTeamProjectId, + azureDevOpsHostType, + azureDevOpsPlanId, + azureDevOpsJobId); + + Request req = + new Request("POST", requestUrl) + .withHeader("Authorization", "Bearer " + azureDevOpsAccessToken) + .withHeader("Content-Type", "application/json"); + + Response resp; + try { + resp = httpClient.execute(req); + } catch (IOException e) { + throw new DatabricksException( + "Failed to request ID token from Azure DevOps at " + requestUrl + ": " + e.getMessage(), + e); + } + + if (resp.getStatusCode() != 200) { + throw new DatabricksException( + "Failed to request ID token from Azure DevOps: status code " + + resp.getStatusCode() + + ", response body: " + + resp.getBody().toString()); + } + + ObjectNode jsonResp; + try { + jsonResp = mapper.readValue(resp.getBody(), ObjectNode.class); + } catch (IOException e) { + throw new DatabricksException( + "Failed to parse Azure DevOps OIDC token response: " + e.getMessage(), e); + } + + // Azure DevOps returns {"oidcToken":"***"} format, not {"value":"***"} like GitHub Actions. + if (!jsonResp.has("oidcToken")) { + throw new DatabricksException("Azure DevOps OIDC token response missing 'oidcToken' field"); + } + + String tokenValue = jsonResp.get("oidcToken").textValue(); + if (Strings.isNullOrEmpty(tokenValue)) { + throw new DatabricksException("Received empty OIDC token from Azure DevOps"); + } + return new IDToken(tokenValue); + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GitHubOidcTokenSupplier.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GitHubOidcTokenSupplier.java deleted file mode 100644 index 523c0df1f..000000000 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GitHubOidcTokenSupplier.java +++ /dev/null @@ -1,79 +0,0 @@ -package com.databricks.sdk.core.oauth; - -import com.databricks.sdk.core.DatabricksException; -import com.databricks.sdk.core.http.HttpClient; -import com.databricks.sdk.core.http.Request; -import com.databricks.sdk.core.http.Response; -import com.fasterxml.jackson.databind.ObjectMapper; -import com.fasterxml.jackson.databind.node.ObjectNode; -import java.io.IOException; - -public class GitHubOidcTokenSupplier { - - private final ObjectMapper mapper = new ObjectMapper(); - private final HttpClient httpClient; - private final String idTokenRequestUrl; - private final String idTokenRequestToken; - private final String tokenAudience; - - public GitHubOidcTokenSupplier( - HttpClient httpClient, - String idTokenRequestUrl, - String idTokenRequestToken, - String tokenAudience) { - this.httpClient = httpClient; - this.idTokenRequestUrl = idTokenRequestUrl; - this.idTokenRequestToken = idTokenRequestToken; - this.tokenAudience = tokenAudience; - } - - /** Checks if the required parameters are present to request a GitHub's OIDC token. */ - public Boolean enabled() { - return idTokenRequestUrl != null && idTokenRequestToken != null; - } - - /** - * Requests a GitHub's OIDC token. - * - * @return A GitHub OIDC token. - */ - public String getOidcToken() { - if (!enabled()) { - throw new DatabricksException("Failed to request ID token: missing required parameters"); - } - - String requestUrl = idTokenRequestUrl; - if (tokenAudience != null) { - requestUrl += "&audience=" + tokenAudience; - } - - Request req = - new Request("GET", requestUrl).withHeader("Authorization", "Bearer " + idTokenRequestToken); - - Response resp; - try { - resp = httpClient.execute(req); - } catch (IOException e) { - throw new DatabricksException( - "Failed to request ID token from " + requestUrl + ":" + e.getMessage(), e); - } - - if (resp.getStatusCode() != 200) { - throw new DatabricksException( - "Failed to request ID token: status code " - + resp.getStatusCode() - + ", response body: " - + resp.getBody().toString()); - } - - ObjectNode jsonResp; - try { - jsonResp = mapper.readValue(resp.getBody(), ObjectNode.class); - } catch (IOException e) { - throw new DatabricksException( - "Failed to request ID token: corrupted token: " + e.getMessage()); - } - - return jsonResp.get("value").textValue(); - } -} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubOidcCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubOidcCredentialsProvider.java deleted file mode 100644 index c52fcf09d..000000000 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/GithubOidcCredentialsProvider.java +++ /dev/null @@ -1,71 +0,0 @@ -package com.databricks.sdk.core.oauth; - -import com.databricks.sdk.core.CredentialsProvider; -import com.databricks.sdk.core.DatabricksConfig; -import com.databricks.sdk.core.DatabricksException; -import com.databricks.sdk.core.HeaderFactory; -import com.google.common.collect.ImmutableMap; -import java.io.IOException; -import java.util.HashMap; -import java.util.Map; - -/** - * GithubOidcCredentialsProvider uses a Token Supplier to get a GitHub OIDC JWT Token and exchanges - * it for a Databricks Token. - */ -public class GithubOidcCredentialsProvider implements CredentialsProvider { - - @Override - public String authType() { - return "github-oidc"; - } - - @Override - public HeaderFactory configure(DatabricksConfig config) throws DatabricksException { - GitHubOidcTokenSupplier idTokenProvider = - new GitHubOidcTokenSupplier( - config.getHttpClient(), - config.getActionsIdTokenRequestUrl(), - config.getActionsIdTokenRequestToken(), - config.getTokenAudience()); - - if (!idTokenProvider.enabled() || config.getHost() == null || config.getClientId() == null) { - return null; - } - - String endpointUrl; - - try { - endpointUrl = config.getOidcEndpoints().getTokenEndpoint(); - } catch (IOException e) { - throw new DatabricksException("Unable to fetch OIDC endpoint: " + e.getMessage(), e); - } - - ClientCredentials clientCredentials = - new ClientCredentials.Builder() - .withHttpClient(config.getHttpClient()) - .withClientId(config.getClientId()) - .withTokenUrl(endpointUrl) - .withScopes(config.getScopes()) - .withAuthParameterPosition(AuthParameterPosition.HEADER) - .withEndpointParametersSupplier( - () -> - new ImmutableMap.Builder() - .put("subject_token_type", "urn:ietf:params:oauth:token-type:jwt") - .put("subject_token", idTokenProvider.getOidcToken()) - .put("grant_type", "urn:ietf:params:oauth:grant-type:token-exchange") - .build()) - .build(); - - CachedTokenSource cachedTokenSource = - new CachedTokenSource.Builder(clientCredentials) - .setAsyncDisabled(config.getDisableAsyncTokenRefresh()) - .build(); - - return () -> { - Map headers = new HashMap<>(); - headers.put("Authorization", "Bearer " + cachedTokenSource.getToken().getAccessToken()); - return headers; - }; - } -} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/AzureDevOpsIDTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/AzureDevOpsIDTokenSourceTest.java new file mode 100644 index 000000000..6e59cf073 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/AzureDevOpsIDTokenSourceTest.java @@ -0,0 +1,251 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.*; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.http.HttpClient; +import com.databricks.sdk.core.http.Request; +import com.databricks.sdk.core.http.Response; +import com.databricks.sdk.core.utils.Environment; +import java.io.IOException; +import java.net.MalformedURLException; +import java.net.URL; +import java.util.HashMap; +import java.util.Map; +import java.util.function.Predicate; +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 AzureDevOpsIDTokenSource. */ +public class AzureDevOpsIDTokenSourceTest { + + private static final String TEST_ID_TOKEN = "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9..."; + private static final String TEST_ACCESS_TOKEN = "test-access-token"; + private static final String TEST_COLLECTION_URI = "https://dev.azure.com/testorg"; + private static final String TEST_PLAN_ID = "test-plan-id"; + private static final String TEST_JOB_ID = "test-job-id"; + private static final String TEST_PROJECT_ID = "test-project-id"; + private static final String TEST_HOST_TYPE = "build"; + + /** Creates a mock Environment with all required Azure DevOps environment variables. */ + private static Environment createValidEnvironment() { + Map envVars = new HashMap<>(); + envVars.put("SYSTEM_ACCESSTOKEN", TEST_ACCESS_TOKEN); + envVars.put("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", TEST_COLLECTION_URI); + envVars.put("SYSTEM_PLANID", TEST_PLAN_ID); + envVars.put("SYSTEM_JOBID", TEST_JOB_ID); + envVars.put("SYSTEM_TEAMPROJECTID", TEST_PROJECT_ID); + envVars.put("SYSTEM_HOSTTYPE", TEST_HOST_TYPE); + return new Environment(envVars, new String[0], "test"); + } + + /** Creates a mock Environment missing the specified environment variable. */ + private static Environment createEnvironmentMissing(String missingVar) { + Map envVars = new HashMap<>(); + envVars.put("SYSTEM_ACCESSTOKEN", TEST_ACCESS_TOKEN); + envVars.put("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", TEST_COLLECTION_URI); + envVars.put("SYSTEM_PLANID", TEST_PLAN_ID); + envVars.put("SYSTEM_JOBID", TEST_JOB_ID); + envVars.put("SYSTEM_TEAMPROJECTID", TEST_PROJECT_ID); + envVars.put("SYSTEM_HOSTTYPE", TEST_HOST_TYPE); + envVars.remove(missingVar); // Remove the specified variable + return new Environment(envVars, new String[0], "test"); + } + + /** Creates a mock HttpClient that returns the specified response. */ + private static HttpClient createHttpMock( + String responseBody, int statusCode, IOException exception) throws IOException { + HttpClient client = mock(HttpClient.class); + if (exception != null) { + when(client.execute(any(Request.class))).thenThrow(exception); + } else { + when(client.execute(any(Request.class))).thenReturn(makeResponse(responseBody, statusCode)); + } + return client; + } + + /** Creates a Response with the specified body and status code. */ + private static Response makeResponse(String body, int statusCode) throws MalformedURLException { + return new Response(body, statusCode, "OK", new URL("https://databricks.com/")); + } + + /** Creates a mock HttpClient that returns a successful OIDC token response. */ + private static HttpClient createValidHttpMock() throws IOException { + return createHttpMock("{\"oidcToken\":\"" + TEST_ID_TOKEN + "\"}", 200, null); + } + + /** Predicate to validate that the HTTP request is constructed correctly. */ + private static final Predicate REQUEST_VALIDATOR = + request -> + request.getMethod().equals("POST") + && request.getUri().toString().contains("api-version=7.2-preview.1") + && request.getUri().toString().contains(TEST_COLLECTION_URI) + && request.getUri().toString().contains(TEST_PROJECT_ID) + && request.getUri().toString().contains(TEST_HOST_TYPE) + && request.getUri().toString().contains(TEST_PLAN_ID) + && request.getUri().toString().contains(TEST_JOB_ID) + && request.getHeaders().get("Authorization").equals("Bearer " + TEST_ACCESS_TOKEN) + && request.getHeaders().get("Content-Type").equals("application/json"); + + private static Stream provideAllTestScenarios() throws IOException { + return Stream.of( + // Constructor validation tests + Arguments.of( + "Null HttpClient", + null, + createValidEnvironment(), + null, + null, + DatabricksException.class), + Arguments.of( + "Missing SYSTEM_ACCESSTOKEN", + mock(HttpClient.class), + createEnvironmentMissing("SYSTEM_ACCESSTOKEN"), + null, + null, + DatabricksException.class), + Arguments.of( + "Missing SYSTEM_TEAMFOUNDATIONCOLLECTIONURI", + mock(HttpClient.class), + createEnvironmentMissing("SYSTEM_TEAMFOUNDATIONCOLLECTIONURI"), + null, + null, + DatabricksException.class), + Arguments.of( + "Missing SYSTEM_PLANID", + mock(HttpClient.class), + createEnvironmentMissing("SYSTEM_PLANID"), + null, + null, + DatabricksException.class), + Arguments.of( + "Missing SYSTEM_JOBID", + mock(HttpClient.class), + createEnvironmentMissing("SYSTEM_JOBID"), + null, + null, + DatabricksException.class), + Arguments.of( + "Missing SYSTEM_TEAMPROJECTID", + mock(HttpClient.class), + createEnvironmentMissing("SYSTEM_TEAMPROJECTID"), + null, + null, + DatabricksException.class), + Arguments.of( + "Missing SYSTEM_HOSTTYPE", + mock(HttpClient.class), + createEnvironmentMissing("SYSTEM_HOSTTYPE"), + null, + null, + DatabricksException.class), + + // HTTP request/response tests + Arguments.of( + "Successful token retrieval", + createValidHttpMock(), + createValidEnvironment(), + REQUEST_VALIDATOR, + TEST_ID_TOKEN, + null), + Arguments.of( + "HTTP request failure", + createHttpMock(null, 0, new IOException("Network error")), + createValidEnvironment(), + null, + null, + DatabricksException.class), + Arguments.of( + "Non-200 status code", + createHttpMock("Error message", 401, null), + createValidEnvironment(), + null, + null, + DatabricksException.class), + Arguments.of( + "Invalid JSON response", + createHttpMock("invalid json", 200, null), + createValidEnvironment(), + null, + null, + DatabricksException.class), + Arguments.of( + "Missing oidcToken field", + createHttpMock("{\"someOtherField\":\"value\"}", 200, null), + createValidEnvironment(), + null, + null, + DatabricksException.class), + Arguments.of( + "Empty oidcToken field", + createHttpMock("{\"oidcToken\":\"\"}", 200, null), + createValidEnvironment(), + null, + null, + DatabricksException.class), + Arguments.of( + "Null oidcToken field", + createHttpMock("{\"oidcToken\":null}", 200, null), + createValidEnvironment(), + null, + null, + DatabricksException.class), + Arguments.of( + "Non-string oidcToken field (number)", + createHttpMock("{\"oidcToken\":123}", 200, null), + createValidEnvironment(), + null, + null, + DatabricksException.class), + Arguments.of( + "Non-string oidcToken field (object)", + createHttpMock("{\"oidcToken\":{\"nested\":\"value\"}}", 200, null), + createValidEnvironment(), + null, + null, + DatabricksException.class)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideAllTestScenarios") + void testAllScenarios( + String testName, + HttpClient httpClient, + Environment environment, + Predicate requestValidator, + String expectedToken, + Class expectedException) { + + if (expectedException != null) { + // Test constructor or runtime exceptions + assertThrows( + expectedException, + () -> { + AzureDevOpsIDTokenSource tokenSource = + new AzureDevOpsIDTokenSource(httpClient, environment); + // If constructor succeeds, try getIDToken to trigger runtime exceptions + tokenSource.getIDToken("ignored-audience"); + }); + } else { + // Test successful cases + AzureDevOpsIDTokenSource tokenSource = new AzureDevOpsIDTokenSource(httpClient, environment); + IDToken token = tokenSource.getIDToken("ignored-audience"); + assertNotNull(token); + assertEquals(expectedToken, token.getValue()); + + // Verify the HTTP request was made correctly + if (requestValidator != null) { + try { + verify(httpClient).execute(argThat(request -> requestValidator.test(request))); + } catch (IOException e) { + fail("Unexpected IOException during request verification: " + e.getMessage()); + } + } + } + } +}