add ut

pkg/provider/azure_test.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	"math"
 	"net/http"
+	"os"
 	"strings"
 	"testing"
 	"time"
@@ -2276,20 +2277,29 @@ publicIPCacheTTLInSeconds: 100
 plsCacheTTLInSeconds: 100
 vmType: vmss
 disableAvailabilitySetNodes: true
+aadFederatedTokenFile: --aad-federated-token-file--
 `
+	os.Setenv("AZURE_TENANT_ID", "123")
+	os.Setenv("AZURE_CLIENT_ID", "456")
+	os.Setenv("AZURE_FEDERATED_TOKEN_FILE", "789")
+	defer func() {
+		os.Unsetenv("AZURE_TENANT_ID")
+		os.Unsetenv("AZURE_CLIENT_ID")
+		os.Unsetenv("AZURE_FEDERATED_TOKEN_FILE")
+	}()
 	validateConfig(t, config)
 }
 
 func validateConfig(t *testing.T, config string) { //nolint
 	azureCloud := getCloudFromConfig(t, config)
 
-	if azureCloud.TenantID != "--tenant-id--" {
+	if azureCloud.TenantID != "123" {
 		t.Errorf("got incorrect value for TenantID")
 	}
 	if azureCloud.SubscriptionID != "--subscription-id--" {
 		t.Errorf("got incorrect value for SubscriptionID")
 	}
-	if azureCloud.AADClientID != "--aad-client-id--" {
+	if azureCloud.AADClientID != "456" {
 		t.Errorf("got incorrect value for AADClientID")
 	}
 	if azureCloud.AADClientSecret != "--aad-client-secret--" {
@@ -2391,6 +2401,12 @@ func validateConfig(t *testing.T, config string) { //nolint
 	if !azureCloud.DisableAvailabilitySetNodes {
 		t.Errorf("got incorrect value for disableAvailabilitySetNodes")
 	}
+	if azureCloud.AADFederatedTokenFile != "789" {
+		t.Errorf("got incorrect value for AADFederatedTokenFile")
+	}
+	if !azureCloud.UseFederatedWorkloadIdentityExtension {
+		t.Errorf("got incorrect value for UseFederatedWorkloadIdentityExtension")
+	}
 }
 
 func getCloudFromConfig(t *testing.T, config string) *Cloud {

pkg/provider/config/azure_auth_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package config
 
 import (
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -205,6 +206,38 @@ func TestGetServicePrincipalTokenFromMSI(t *testing.T) {
 
 }
 
+func TestGetServicePrincipalTokenFromWorkloadIdentity(t *testing.T) {
+	config := &AzureAuthConfig{
+		TenantID:                              "TenantID",
+		AADClientID:                           "AADClientID",
+		AADFederatedTokenFile:                 "/tmp/federated-token",
+		UseFederatedWorkloadIdentityExtension: true,
+	}
+	env := &azure.PublicCloud
+
+	token, err := GetServicePrincipalToken(config, env, "")
+	assert.NoError(t, err)
+	marshalToken, _ := token.MarshalJSON()
+
+	oauthConfig, err := adal.NewOAuthConfigWithAPIVersion(env.ActiveDirectoryEndpoint, config.TenantID, nil)
+	assert.NoError(t, err)
+
+	jwtCallback := func() (string, error) {
+		jwt, err := os.ReadFile(config.AADFederatedTokenFile)
+		if err != nil {
+			return "", fmt.Errorf("failed to read a file with a federated token: %w", err)
+		}
+		return string(jwt), nil
+	}
+
+	spt, err := adal.NewServicePrincipalTokenFromFederatedTokenCallback(*oauthConfig, config.AADClientID, jwtCallback, env.ResourceManagerEndpoint)
+	assert.NoError(t, err)
+
+	marshalSpt, _ := spt.MarshalJSON()
+
+	assert.Equal(t, marshalToken, marshalSpt)
+}
+
 func TestGetServicePrincipalToken(t *testing.T) {
 	config := &AzureAuthConfig{
 		TenantID:        "TenantID",