From db238de831630af1feb5798892e4e9a858a6da45 Mon Sep 17 00:00:00 2001
From: weizhichen <weizhichen@microsoft.com>
Date: Wed, 20 Mar 2024 09:07:45 +0000
Subject: [PATCH] add ut

---
 pkg/azuredisk/azuredisk_test.go | 36 +++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/pkg/azuredisk/azuredisk_test.go b/pkg/azuredisk/azuredisk_test.go
index 1eccb9fc85..e34d7f0fd1 100644
--- a/pkg/azuredisk/azuredisk_test.go
+++ b/pkg/azuredisk/azuredisk_test.go
@@ -203,6 +203,42 @@ func TestRun(t *testing.T) {
 				assert.Nil(t, err)
 			},
 		},
+		{
+			name: "Successful run with federated workload identity azure client",
+			testFunc: func(t *testing.T) {
+				if err := os.WriteFile(fakeCredFile, []byte(fakeCredContent), 0666); err != nil {
+					t.Error(err)
+				}
+
+				defer func() {
+					if err := os.Remove(fakeCredFile); err != nil {
+						t.Error(err)
+					}
+				}()
+
+				t.Setenv(consts.DefaultAzureCredentialFileEnv, fakeCredFile)
+				t.Setenv("AZURE_TENANT_ID", "1234")
+				t.Setenv("AZURE_CLIENT_ID", "123456")
+				t.Setenv("AZURE_FEDERATED_TOKEN_FILE", "fake-token-file")
+
+				d := newDriverV1(&DriverOptions{
+					Endpoint: "tcp://127.0.0.1:0",
+				})
+
+				ctx, cancel := context.WithCancel(context.Background())
+				ch := make(chan error)
+				go func() {
+					err := d.Run(ctx)
+					ch <- err
+				}()
+				cancel()
+				assert.Nil(t, <-ch)
+				assert.Equal(t, d.cloud.UseFederatedWorkloadIdentityExtension, true)
+				assert.Equal(t, d.cloud.AADFederatedTokenFile, "fake-token-file")
+				assert.Equal(t, d.cloud.AADClientID, "123456")
+				assert.Equal(t, d.cloud.TenantID, "1234")
+			},
+		},
 	}
 
 	for _, tc := range testCases {