Skip to content

Commit 0f128af

Browse files
committed
Add tests for discoveryUrl OIDC endpoint resolution
Adds tests covering: - DATABRICKS_DISCOVERY_URL env var is loaded into Config - getDatabricksOidcEndpoints() short-circuits to discoveryUrl when set The underlying implementation (discoveryUrl config field and getDatabricksOidcEndpoints() short-circuit) was already present in the Java SDK. Mirrors the Python SDK's get_endpoints_from_url() addition.
1 parent d86b772 commit 0f128af

1 file changed

Lines changed: 90 additions & 0 deletions

File tree

databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,4 +421,94 @@ public void testGetClientTypeAccountOnUnified() {
421421
.setExperimentalIsUnifiedHost(true)
422422
.getClientType());
423423
}
424+
425+
// --- HostMetadata tests ---
426+
427+
private static final String DUMMY_ACCOUNT_ID = "00000000-0000-0000-0000-000000000001";
428+
private static final String DUMMY_WORKSPACE_ID = "111111111111111";
429+
430+
private static Environment emptyEnv() {
431+
return new Environment(new HashMap<>(), new ArrayList<>(), System.getProperty("os.name"));
432+
}
433+
434+
@Test
435+
public void testGetHostMetadataWorkspaceStaticOidcEndpoint() throws IOException {
436+
String response =
437+
"{\"oidc_endpoint\":\"https://ws.databricks.com/oidc\","
438+
+ "\"account_id\":\""
439+
+ DUMMY_ACCOUNT_ID
440+
+ "\","
441+
+ "\"workspace_id\":\""
442+
+ DUMMY_WORKSPACE_ID
443+
+ "\"}";
444+
try (FixtureServer server =
445+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
446+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
447+
config.resolve(emptyEnv());
448+
HostMetadata meta = config.getHostMetadata();
449+
assertEquals("https://ws.databricks.com/oidc", meta.getOidcEndpoint());
450+
assertEquals(DUMMY_ACCOUNT_ID, meta.getAccountId());
451+
assertEquals(DUMMY_WORKSPACE_ID, meta.getWorkspaceId());
452+
}
453+
}
454+
455+
@Test
456+
public void testGetHostMetadataAccountRawOidcTemplate() throws IOException {
457+
String response =
458+
"{\"oidc_endpoint\":\"https://acc.databricks.com/oidc/accounts/{account_id}\"}";
459+
try (FixtureServer server =
460+
new FixtureServer().with("GET", "/.well-known/databricks-config", response, 200)) {
461+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
462+
config.resolve(emptyEnv());
463+
HostMetadata meta = config.getHostMetadata();
464+
assertEquals(
465+
"https://acc.databricks.com/oidc/accounts/{account_id}", meta.getOidcEndpoint());
466+
assertNull(meta.getAccountId());
467+
assertNull(meta.getWorkspaceId());
468+
}
469+
}
470+
471+
@Test
472+
public void testGetHostMetadataRaisesOnHttpError() throws IOException {
473+
try (FixtureServer server =
474+
new FixtureServer().with("GET", "/.well-known/databricks-config", "{}", 404)) {
475+
DatabricksConfig config = new DatabricksConfig().setHost(server.getUrl());
476+
config.resolve(emptyEnv());
477+
DatabricksException ex =
478+
assertThrows(DatabricksException.class, () -> config.getHostMetadata());
479+
assertTrue(ex.getMessage().contains("Failed to fetch host metadata"));
480+
}
481+
}
482+
483+
// --- discoveryUrl / OIDC endpoint tests ---
484+
485+
@Test
486+
public void testDiscoveryUrlFromEnv() {
487+
Map<String, String> env = new HashMap<>();
488+
env.put("DATABRICKS_DISCOVERY_URL", "https://custom.idp.example.com/oidc");
489+
DatabricksConfig config = new DatabricksConfig();
490+
config.resolve(new Environment(env, new ArrayList<>(), System.getProperty("os.name")));
491+
assertEquals("https://custom.idp.example.com/oidc", config.getDiscoveryUrl());
492+
}
493+
494+
@Test
495+
public void testDatabricksOidcEndpointsUsesDiscoveryUrl() throws IOException {
496+
String discoveryUrlSuffix = "/oidc";
497+
String discoveryUrlResponse =
498+
"{\"authorization_endpoint\":\"https://ws.databricks.com/oidc/v1/authorize\","
499+
+ "\"token_endpoint\":\"https://ws.databricks.com/oidc/v1/token\"}";
500+
try (FixtureServer server =
501+
new FixtureServer().with("GET", discoveryUrlSuffix, discoveryUrlResponse, 200)) {
502+
String discoveryUrl = server.getUrl() + discoveryUrlSuffix;
503+
OpenIDConnectEndpoints endpoints =
504+
new DatabricksConfig()
505+
.setHost(server.getUrl())
506+
.setDiscoveryUrl(discoveryUrl)
507+
.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build())
508+
.getDatabricksOidcEndpoints();
509+
assertEquals(
510+
"https://ws.databricks.com/oidc/v1/authorize", endpoints.getAuthorizationEndpoint());
511+
assertEquals("https://ws.databricks.com/oidc/v1/token", endpoints.getTokenEndpoint());
512+
}
513+
}
424514
}

0 commit comments

Comments
 (0)