@@ -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