77import com .databricks .sdk .core .DatabricksException ;
88import com .databricks .sdk .core .FixtureServer ;
99import com .databricks .sdk .core .HeaderFactory ;
10+ import com .databricks .sdk .core .OpenIDConnectEndpoints ;
1011import com .databricks .sdk .core .commons .CommonsHttpClient ;
1112import com .databricks .sdk .core .http .HttpClient ;
1213import com .databricks .sdk .core .http .Request ;
@@ -303,7 +304,8 @@ void cacheWithValidRefreshableTokenTest() throws IOException {
303304 any (DatabricksConfig .class ),
304305 any (String .class ),
305306 any (String .class ),
306- any (TokenCache .class ));
307+ any (TokenCache .class ),
308+ any (OpenIDConnectEndpoints .class ));
307309
308310 // Verify token was NOT saved back to cache (we're using the cached one as-is).
309311 Mockito .verify (mockTokenCache , Mockito .never ()).save (any (Token .class ));
@@ -363,7 +365,7 @@ void cacheWithValidNonRefreshableTokenTest() throws IOException {
363365
364366 // Verify performBrowserAuth was NOT called.
365367 Mockito .verify (provider , Mockito .never ())
366- .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ));
368+ .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ), any ( OpenIDConnectEndpoints . class ) );
367369
368370 // Verify no token was saved (we're using the cached one as-is).
369371 Mockito .verify (mockTokenCache , Mockito .never ()).save (any (Token .class ));
@@ -430,7 +432,9 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException {
430432 any (DatabricksConfig .class ),
431433 any (String .class ),
432434 any (String .class ),
433- any (TokenCache .class ));
435+ any (TokenCache .class ),
436+ any (OpenIDConnectEndpoints .class )
437+ );
434438
435439 // Verify token was saved back to cache
436440 Mockito .verify (mockTokenCache , Mockito .times (1 )).save (any (Token .class ));
@@ -508,7 +512,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException {
508512 Mockito .spy (new ExternalBrowserCredentialsProvider (mockTokenCache ));
509513 Mockito .doReturn (cachedTokenSource )
510514 .when (provider )
511- .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ));
515+ .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ), any ( OpenIDConnectEndpoints . class ) );
512516
513517 // Spy on the config to inject the endpoints
514518 DatabricksConfig spyConfig = Mockito .spy (config );
@@ -527,7 +531,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException {
527531
528532 // Verify performBrowserAuth was called since refresh failed
529533 Mockito .verify (provider , Mockito .times (1 ))
530- .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ));
534+ .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ), any ( OpenIDConnectEndpoints . class ) );
531535
532536 // Verify token was saved after browser auth (for the new token)
533537 Mockito .verify (mockTokenCache , Mockito .times (1 )).save (any (Token .class ));
@@ -572,17 +576,26 @@ void cacheWithInvalidTokensTest() throws IOException {
572576 new DatabricksConfig ()
573577 .setAuthType ("external-browser" )
574578 .setHost ("https://test.databricks.com" )
575- .setClientId ("test-client-id" );
579+ .setClientId ("test-client-id" )
580+ .setHttpClient (mockHttpClient );
576581
577582 // Create our provider and mock the browser auth method
578583 ExternalBrowserCredentialsProvider provider =
579584 Mockito .spy (new ExternalBrowserCredentialsProvider (mockTokenCache ));
580585 Mockito .doReturn (cachedTokenSource )
581586 .when (provider )
582- .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ));
587+ .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ), any (OpenIDConnectEndpoints .class ));
588+
589+ // Spy on the config to inject the endpoints
590+ OpenIDConnectEndpoints endpoints =
591+ new OpenIDConnectEndpoints (
592+ "https://test.databricks.com/oidc/v1/token" ,
593+ "https://test.databricks.com/oidc/v1/authorize" );
594+ DatabricksConfig spyConfig = Mockito .spy (config );
595+ Mockito .doReturn (endpoints ).when (spyConfig ).getOidcEndpoints ();
583596
584597 // Configure provider
585- HeaderFactory headerFactory = provider .configure (config );
598+ HeaderFactory headerFactory = provider .configure (spyConfig );
586599 assertNotNull (headerFactory );
587600 // Verify headers contain the browser auth token (fallback)
588601 Map <String , String > headers = headerFactory .headers ();
@@ -593,7 +606,7 @@ void cacheWithInvalidTokensTest() throws IOException {
593606
594607 // Verify performBrowserAuth was called since we had an invalid token
595608 Mockito .verify (provider , Mockito .times (1 ))
596- .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ));
609+ .performBrowserAuth (any (DatabricksConfig .class ), any (), any (), any (TokenCache .class ), any ( OpenIDConnectEndpoints . class ) );
597610
598611 // Verify token was saved after browser auth (for the new token)
599612 Mockito .verify (mockTokenCache , Mockito .times (1 )).save (any (Token .class ));
@@ -609,7 +622,7 @@ void doNotAddOfflineAccessScopeWhenDisableOauthRefreshTokenIsTrue() {
609622 .setScopes (Arrays .asList ("my-test-scope" ));
610623
611624 ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider ();
612- List <String > scopes = provider .getScopes (config );
625+ List <String > scopes = provider .getScopes (config , null );
613626
614627 assertEquals (1 , scopes .size ());
615628 assertTrue (scopes .contains ("my-test-scope" ));
@@ -625,7 +638,7 @@ void doNotRemoveUserProvidedScopesWhenDisableOauthRefreshTokenIsTrue() {
625638 .setScopes (Arrays .asList ("my-test-scope" , "offline_access" ));
626639
627640 ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider ();
628- List <String > scopes = provider .getScopes (config );
641+ List <String > scopes = provider .getScopes (config , null );
629642
630643 assertEquals (2 , scopes .size ());
631644 assertTrue (scopes .contains ("offline_access" ));
@@ -641,7 +654,7 @@ void addOfflineAccessScopeWhenDisableOauthRefreshTokenIsFalse() {
641654 .setScopes (Arrays .asList ("my-test-scope" ));
642655
643656 ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider ();
644- List <String > scopes = provider .getScopes (config );
657+ List <String > scopes = provider .getScopes (config , null );
645658
646659 assertEquals (2 , scopes .size ());
647660 assertTrue (scopes .contains ("offline_access" ));
0 commit comments