@@ -724,6 +724,69 @@ func TestEnsureResolved_HostMetadata_MissingAccountIdWithPlaceholder_Warns(t *te
724724 assert .Empty (t , cfg .DiscoveryURL )
725725}
726726
727+ func TestApplyHostMetadata_SetsTokenAudienceForAccountHost (t * testing.T ) {
728+ noopLoader := mockLoader (func (cfg * Config ) error { return nil })
729+ cfg := & Config {
730+ Host : testHMHost ,
731+ Experimental_IsUnifiedHost : true ,
732+ Loaders : []Loader {noopLoader },
733+ HTTPTransport : fixtures.SliceTransport {
734+ {
735+ Method : "GET" ,
736+ Resource : "/.well-known/databricks-config" ,
737+ Status : 200 ,
738+ Response : `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "cloud": "AWS"}` ,
739+ },
740+ },
741+ }
742+ err := cfg .EnsureResolved ()
743+ require .NoError (t , err )
744+ // No workspace_id and has account_id → token audience should be set
745+ assert .Equal (t , testHMAccountID , cfg .TokenAudience )
746+ }
747+
748+ func TestApplyHostMetadata_NoTokenAudienceForWorkspaceHost (t * testing.T ) {
749+ noopLoader := mockLoader (func (cfg * Config ) error { return nil })
750+ cfg := & Config {
751+ Host : testHMHost ,
752+ Experimental_IsUnifiedHost : true ,
753+ Loaders : []Loader {noopLoader },
754+ HTTPTransport : fixtures.SliceTransport {
755+ {
756+ Method : "GET" ,
757+ Resource : "/.well-known/databricks-config" ,
758+ Status : 200 ,
759+ Response : `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "workspace_id": "` + testHMWorkspaceID + `", "cloud": "AWS"}` ,
760+ },
761+ },
762+ }
763+ err := cfg .EnsureResolved ()
764+ require .NoError (t , err )
765+ // Has workspace_id → token audience should NOT be set
766+ assert .Empty (t , cfg .TokenAudience )
767+ }
768+
769+ func TestApplyHostMetadata_DoesNotOverrideExistingTokenAudience (t * testing.T ) {
770+ noopLoader := mockLoader (func (cfg * Config ) error { return nil })
771+ cfg := & Config {
772+ Host : testHMHost ,
773+ TokenAudience : "custom-audience" ,
774+ Experimental_IsUnifiedHost : true ,
775+ Loaders : []Loader {noopLoader },
776+ HTTPTransport : fixtures.SliceTransport {
777+ {
778+ Method : "GET" ,
779+ Resource : "/.well-known/databricks-config" ,
780+ Status : 200 ,
781+ Response : `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "cloud": "AWS"}` ,
782+ },
783+ },
784+ }
785+ err := cfg .EnsureResolved ()
786+ require .NoError (t , err )
787+ assert .Equal (t , "custom-audience" , cfg .TokenAudience )
788+ }
789+
727790func TestConfig_ResolveHostMetadata_Clouds (t * testing.T ) {
728791 tests := []struct {
729792 name string
0 commit comments