Adds integration tests for OIDC Endpoints

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

@@ -798,15 +798,18 @@ private OpenIDConnectEndpoints fetchDefaultOidcEndpoints() throws IOException {
     }
 
     if (isAzure() && getAzureClientId() != null) {
-      Request request = new Request("GET", getHost() + "/oidc/oauth2/v2.0/authorize");
-      request.setRedirectionBehavior(false);
-      Response resp = getHttpClient().execute(request);
-      String realAuthUrl = resp.getFirstHeader("location");
-      if (realAuthUrl == null) {
-        return null;
+      ApiClient apiClient =
+          new ApiClient.Builder()
+              .withHttpClient(getHttpClient())
+              .withGetHostFunc(v -> getHost())
+              .build();
+      try {
+        return apiClient.execute(
+            new Request("GET", "/oidc/.well-known/oauth-authorization-server"),
+            OpenIDConnectEndpoints.class);
+      } catch (IOException e) {
+        throw new DatabricksException("IO error: " + e.getMessage(), e);
       }
-      return new OpenIDConnectEndpoints(
-          realAuthUrl.replaceAll("/authorize", "/token"), realAuthUrl);
     }
     if (isAccountClient() && getAccountId() != null) {
       String prefix = getHost() + "/oidc/accounts/" + getAccountId();

databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcAccountEndpointIT.java

@@ -0,0 +1,62 @@
+package com.databricks.sdk.integration;
+
+import com.databricks.sdk.AccountClient;
+import com.databricks.sdk.core.DatabricksConfig;
+import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import com.databricks.sdk.integration.framework.EnvContext;
+import com.databricks.sdk.integration.framework.EnvTest;
+import java.io.IOException;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+@EnvContext("account")
+@ExtendWith(EnvTest.class)
+public class OidcAccountEndpointIT {
+  @Test
+  void checkEndpoints(AccountClient a) throws IOException {
+    OpenIDConnectEndpoints endpoints = a.config().getOidcEndpoints();
+    String host = a.config().getHost();
+    String accountId = a.config().getAccountId();
+    if (a.config().isAzure()) {
+      assert endpoints.getAuthorizationEndpoint().equals(host + "/oidc/v1/authorize");
+      assert endpoints.getTokenEndpoint().equals(host + "/oidc/v1/token");
+    } else {
+      assert endpoints
+          .getAuthorizationEndpoint()
+          .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize");
+      assert endpoints
+          .getTokenEndpoint()
+          .equals(host + "/oidc/accounts/" + accountId + "/v1/token");
+    }
+  }
+
+  @Test
+  void unifiedEndpoints(AccountClient a) throws IOException {
+    DatabricksConfig cfg = a.config();
+    cfg.setExperimentalIsUnifiedHost(true);
+    OpenIDConnectEndpoints endpoints = cfg.getOidcEndpoints();
+    String host = cfg.getHost();
+    String accountId = cfg.getAccountId();
+    assert endpoints
+        .getAuthorizationEndpoint()
+        .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize");
+    assert endpoints.getTokenEndpoint().equals(host + "/oidc/accounts/" + accountId + "/v1/token");
+  }
+
+  @Test
+  @EnabledIfEnvironmentVariable(named = "UNIFIED_HOST", matches = ".+")
+  void unifiedEndpointsForSpog(AccountClient a) throws IOException {
+    String unifiedHost = System.getenv("UNIFIED_HOST");
+    DatabricksConfig cfg = a.config();
+    cfg.setHost(unifiedHost);
+    cfg.setExperimentalIsUnifiedHost(true);
+    OpenIDConnectEndpoints endpoints = cfg.getOidcEndpoints();
+    String host = cfg.getHost();
+    String accountId = cfg.getAccountId();
+    assert endpoints
+        .getAuthorizationEndpoint()
+        .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize");
+    assert endpoints.getTokenEndpoint().equals(host + "/oidc/accounts/" + accountId + "/v1/token");
+  }
+}

databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcUcAccountEndpointIT.java

@@ -0,0 +1,49 @@
+package com.databricks.sdk.integration;
+
+import com.databricks.sdk.AccountClient;
+import com.databricks.sdk.core.DatabricksConfig;
+import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import com.databricks.sdk.integration.framework.EnvContext;
+import com.databricks.sdk.integration.framework.EnvTest;
+import java.io.IOException;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+@EnvContext("ucacct")
+@ExtendWith(EnvTest.class)
+public class OidcUcAccountEndpointIT {
+  @Test
+  void checkEndpoints(AccountClient a) throws IOException {
+    OpenIDConnectEndpoints endpoints = a.config().getOidcEndpoints();
+    String host = a.config().getHost();
+    String accountId = a.config().getAccountId();
+    if (a.config().isAzure()) {
+      assert endpoints.getAuthorizationEndpoint().equals(host + "/oidc/v1/authorize");
+      assert endpoints.getTokenEndpoint().equals(host + "/oidc/v1/token");
+    } else {
+      assert endpoints
+          .getAuthorizationEndpoint()
+          .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize");
+      assert endpoints
+          .getTokenEndpoint()
+          .equals(host + "/oidc/accounts/" + accountId + "/v1/token");
+    }
+  }
+
+  @Test
+  @EnabledIfEnvironmentVariable(named = "UNIFIED_HOST", matches = ".+")
+  void unifiedEndpointsForSpog(AccountClient a) throws IOException {
+    String unifiedHost = System.getenv("UNIFIED_HOST");
+    DatabricksConfig cfg = a.config();
+    cfg.setHost(unifiedHost);
+    cfg.setExperimentalIsUnifiedHost(true);
+    OpenIDConnectEndpoints endpoints = cfg.getOidcEndpoints();
+    String host = cfg.getHost();
+    String accountId = cfg.getAccountId();
+    assert endpoints
+        .getAuthorizationEndpoint()
+        .equals(host + "/oidc/accounts/" + accountId + "/v1/authorize");
+    assert endpoints.getTokenEndpoint().equals(host + "/oidc/accounts/" + accountId + "/v1/token");
+  }
+}

databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcUcWorkspaceEndpointIT.java

@@ -0,0 +1,21 @@
+package com.databricks.sdk.integration;
+
+import com.databricks.sdk.WorkspaceClient;
+import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import com.databricks.sdk.integration.framework.EnvContext;
+import com.databricks.sdk.integration.framework.EnvTest;
+import java.io.IOException;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+@EnvContext("ucws")
+@ExtendWith(EnvTest.class)
+public class OidcUcWorkspaceEndpointIT {
+  @Test
+  void checkEndpoints(WorkspaceClient w) throws IOException {
+    OpenIDConnectEndpoints endpoints = w.config().getOidcEndpoints();
+    String host = w.config().getHost();
+    assert endpoints.getAuthorizationEndpoint().equals(host + "/oidc/v1/authorize");
+    assert endpoints.getTokenEndpoint().equals(host + "/oidc/v1/token");
+  }
+}

databricks-sdk-java/src/test/java/com/databricks/sdk/integration/OidcWorkspaceEndpointIT.java

@@ -0,0 +1,21 @@
+package com.databricks.sdk.integration;
+
+import com.databricks.sdk.WorkspaceClient;
+import com.databricks.sdk.core.oauth.OpenIDConnectEndpoints;
+import com.databricks.sdk.integration.framework.EnvContext;
+import com.databricks.sdk.integration.framework.EnvTest;
+import java.io.IOException;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+@EnvContext("workspace")
+@ExtendWith(EnvTest.class)
+public class OidcWorkspaceEndpointIT {
+  @Test
+  void checkEndpoints(WorkspaceClient w) throws IOException {
+    OpenIDConnectEndpoints endpoints = w.config().getOidcEndpoints();
+    String host = w.config().getHost();
+    assert endpoints.getAuthorizationEndpoint().equals(host + "/oidc/v1/authorize");
+    assert endpoints.getTokenEndpoint().equals(host + "/oidc/v1/token");
+  }
+}