Fix client token credentials issue before Az.Account module is imported. (#27248)

vidai-msft · web-flow · commit 7c039998f484 · 2025-03-04T00:32:10.000+11:00
diff --git a/src/Attestation/Attestation.Test/ScenarioTests/AttestationTestRunner.cs b/src/Attestation/Attestation.Test/ScenarioTests/AttestationTestRunner.cs
@@ -15,10 +15,7 @@
 using Microsoft.Azure.Attestation;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using Microsoft.Azure.Commands.TestFx;
-using Microsoft.Azure.Management.Attestation;
-using Microsoft.Azure.Management.Internal.Resources;
 using Microsoft.Azure.Test.HttpRecorder;
-using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
 using System.Collections.Generic;
 using Xunit.Abstractions;
 
@@ -55,28 +52,14 @@ protected AttestationTestRunner(ITestOutputHelper output)
                         {"Microsoft.Features", null},
                         {"Microsoft.Authorization", null}
                     }
-                ).WithManagementClients(
-                    GetResourceManagementClient,
-                    GetAttestationManagementClient,
-                    GetAttestationClient
+                )
+                .WithManagementClients(context =>
+                    {
+                        var creds = context.GetClientCredentials(AzureEnvironment.ExtendedEndpoint.AzureAttestationServiceEndpointResourceId);
+                        return new AttestationClient(creds, HttpMockServer.CreateInstance());
+                    }
                 )
                 .Build();
         }
-
-        private static ResourceManagementClient GetResourceManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<ResourceManagementClient>();
-        }
-
-        private static AttestationManagementClient GetAttestationManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<AttestationManagementClient>();
-        }
-
-        private static AttestationClient GetAttestationClient(MockContext context)
-        {
-            var creds = TestEnvironmentFactory.GetTestEnvironment().GetAccessToken(AzureEnvironmentConstants.AzureAttestationServiceEndpointResourceId);
-            return new AttestationClient(creds, HttpMockServer.CreateInstance());
-        }
     }
 }
diff --git a/src/DataLakeAnalytics/DataLakeAnalytics.Test/ScenarioTests/DataLakeAnalyticsTestRunner.cs b/src/DataLakeAnalytics/DataLakeAnalytics.Test/ScenarioTests/DataLakeAnalyticsTestRunner.cs
@@ -12,12 +12,9 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.Azure.Commands.TestFx;
 using Microsoft.Azure.Management.DataLake.Analytics;
-using Microsoft.Azure.Management.DataLake.Store;
-using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
 using System.Collections.Generic;
-using NewResourceManagementClient = Microsoft.Azure.Management.Internal.Resources.ResourceManagementClient;
-using Microsoft.Azure.Commands.TestFx;
 using Xunit.Abstractions;
 
 namespace Microsoft.Azure.Commands.DataLakeAnalytics.Test.ScenarioTests
@@ -53,47 +50,18 @@ protected DataLakeAnalyticsTestRunner(ITestOutputHelper output)
                         {"Microsoft.Features", null},
                         {"Microsoft.Authorization", null}
                     }
-                ).WithManagementClients(
-                    GetNewResourceManagementClient,
-                    GetDataLakeStoreAccountManagementClient,
-                    GetDataLakeAnalyticsAccountManagementClient,
-                    GetDataLakeAnalyticsJobManagementClient,
-                    GetDataLakeAnalyticsCatalogManagementClient
+                )
+                .WithManagementClients(
+                    context =>
+                    {
+                        return context.GetServiceClient<DataLakeAnalyticsJobManagementClient>(true);
+                    },
+                    context =>
+                    {
+                        return context.GetServiceClient<DataLakeAnalyticsCatalogManagementClient>(true);
+                    }
                 )
                 .Build();
         }
-
-        private static NewResourceManagementClient GetNewResourceManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<NewResourceManagementClient>();
-        }
-
-        private static DataLakeStoreAccountManagementClient GetDataLakeStoreAccountManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<DataLakeStoreAccountManagementClient>();
-        }
-
-        private static DataLakeAnalyticsAccountManagementClient GetDataLakeAnalyticsAccountManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<DataLakeAnalyticsAccountManagementClient>();
-        }
-
-        private static DataLakeAnalyticsJobManagementClient GetDataLakeAnalyticsJobManagementClient(MockContext context)
-        {
-            var currentEnvironment = TestEnvironmentFactory.GetTestEnvironment();
-            var toReturn = context.GetServiceClient<DataLakeAnalyticsJobManagementClient>(true);
-            toReturn.AdlaJobDnsSuffix =
-                currentEnvironment.Endpoints.DataLakeAnalyticsJobAndCatalogServiceUri.OriginalString.Replace("https://", "");
-            return toReturn;
-        }
-
-        private static DataLakeAnalyticsCatalogManagementClient GetDataLakeAnalyticsCatalogManagementClient(MockContext context)
-        {
-            var currentEnvironment = TestEnvironmentFactory.GetTestEnvironment();
-            var toReturn = context.GetServiceClient<DataLakeAnalyticsCatalogManagementClient>(true);
-            toReturn.AdlaCatalogDnsSuffix =
-                currentEnvironment.Endpoints.DataLakeAnalyticsJobAndCatalogServiceUri.OriginalString.Replace("https://", "");
-            return toReturn;
-        }
     }
 }
diff --git a/src/DataLakeStore/DataLakeStore.Test/ScenarioTests/DataLakeStoreTestRunner.cs b/src/DataLakeStore/DataLakeStore.Test/ScenarioTests/DataLakeStoreTestRunner.cs
@@ -56,11 +56,11 @@ protected DataLakeStoreTestRunner(ITestOutputHelper output)
                 .WithRecordMatcher(
                     (ignoreResourcesClient, resourceProviders, userAgentsToIgnore) => new UrlDecodingRecordMatcher(ignoreResourcesClient, resourceProviders, userAgentsToIgnore)
                 )
-                .WithManagementClients(mockContext =>
+                .WithManagementClients(context =>
                     {
                         AdlsClientFactory.IsTest = true;
-                        var creds = TestEnvironmentFactory.GetTestEnvironment().GetAccessToken(AzureEnvironmentConstants.AzureDataLakeStoreFileSystemEndpointSuffix);
-                        AdlsClientFactory.CustomDelegatingHAndler = mockContext.AddHandlers(creds, new AdlMockDelegatingHandler());
+                        var creds = context.GetClientCredentials(AzureEnvironment.Endpoint.AzureDataLakeStoreFileSystemEndpointSuffix);
+                        AdlsClientFactory.CustomDelegatingHAndler = context.AddHandlers(creds, new AdlMockDelegatingHandler());
                         AdlsClientFactory.MockCredentials = creds;
                         return new object();
                     }
diff --git a/src/Network/Network.Test/ScenarioTests/NetworkTestRunner.cs b/src/Network/Network.Test/ScenarioTests/NetworkTestRunner.cs
@@ -16,7 +16,6 @@
 using Microsoft.Azure.Commands.TestFx;
 using Microsoft.Azure.KeyVault;
 using Microsoft.Azure.Test.HttpRecorder;
-using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
 using System.Collections.Generic;
 using Xunit.Abstractions;
 
@@ -67,16 +66,14 @@ protected NetworkTestRunner(ITestOutputHelper output)
                         {"Microsoft.ManagedServiceIdentity", null},
                         {"Microsoft.PrivateDns", null},
                     }
-                ).WithManagementClients(
-                    GetKeyVaultClient
+                )
+                .WithManagementClients(context =>
+                    {
+                        var creds = context.GetClientCredentials(AzureEnvironment.Endpoint.AzureKeyVaultServiceEndpointResourceId);
+                        return new KeyVaultClient(creds, HttpMockServer.CreateInstance());
+                    }
                 )
                 .Build();
         }
-
-        private static KeyVaultClient GetKeyVaultClient(MockContext context)
-        {
-            var creds = TestEnvironmentFactory.GetTestEnvironment().GetAccessToken(AzureEnvironmentConstants.AzureKeyVaultServiceEndpointResourceId);
-            return new KeyVaultClient(creds, HttpMockServer.CreateInstance());
-        }
     }
 }
diff --git a/src/Resources/Resources.Test/ScenarioTests/ResourcesTestRunner.cs b/src/Resources/Resources.Test/ScenarioTests/ResourcesTestRunner.cs
@@ -12,12 +12,12 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using Microsoft.Azure.Commands.ResourceManager.Cmdlets.Components;
 using Microsoft.Azure.Commands.ResourceManager.Cmdlets.Extensions;
 using Microsoft.Azure.Commands.TestFx;
 using Microsoft.Azure.Test.HttpRecorder;
 using Microsoft.Rest;
-using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Net.Http;
@@ -56,19 +56,13 @@ protected ResourcesTestRunner(ITestOutputHelper output)
                     },
                     resourceProviders: new Dictionary<string, string>()
                 )
-                .WithManagementClients(mockContext =>
-                {
-                    var credentials = HttpMockServer.Mode != HttpRecorderMode.Playback
-                        ? new Func<ServiceClientCredentials>(() =>
-                        {
-                            var testEnvironment = TestEnvironmentFactory.GetTestEnvironment();
-                            return testEnvironment.TokenInfo[TokenAudience.Management];
-                        })()
-                        : new TokenCredentials("foo");
-
-                    HttpClientHelperFactory.Instance = new TestHttpClientHelperFactory(credentials);
-                    return HttpClientHelperFactory.Instance;
-                })
+                .WithManagementClients(context =>
+                    {
+                        var creds = context.GetClientCredentials(AzureEnvironment.Endpoint.ActiveDirectoryServiceEndpointResourceId);
+                        HttpClientHelperFactory.Instance = new TestHttpClientHelperFactory(creds);
+                        return HttpClientHelperFactory.Instance;
+                    }
+                )
                 .Build();
         }
     }
diff --git a/src/Resources/Resources.Test/ScenarioTests/RoleAssignmentTests.cs b/src/Resources/Resources.Test/ScenarioTests/RoleAssignmentTests.cs
@@ -91,7 +91,7 @@ public void RaByResource()
             TestRunner.RunTestScript("Test-RaByResource");
         }
 
-        [Fact]//(Skip = "Test indicates Graph call is not recorded when it actually is, refer to https://github.com/Azure/azure-powershell/issues/14632 for more details, test passes in record mode")]
+        [Fact]
         [Trait(Category.AcceptanceType, Category.CheckIn)]
         public void RaByServicePrincipal()
         {
diff --git a/src/Sql/Sql.Test/ScenarioTests/SqlTestRunner.cs b/src/Sql/Sql.Test/ScenarioTests/SqlTestRunner.cs
@@ -12,24 +12,10 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using Microsoft.Azure.Management.Network;
-using Microsoft.Azure.Test.HttpRecorder;
-using Microsoft.WindowsAzure.Commands.ScenarioTest;
-using System.Collections.Generic;
-using Microsoft.Azure.Management.Internal.Resources;
-using Microsoft.Azure.Management.Sql;
-using CommonStorage = Microsoft.Azure.Management.Storage.Version2017_10_01;
-using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
-using Xunit.Abstractions;
-using Microsoft.Azure.Management.EventHub;
-using Microsoft.Azure.Management.OperationalInsights;
-using SDKMonitor = Microsoft.Azure.Management.Monitor;
-using CommonMonitor = Microsoft.Azure.Management.Monitor.Version2018_09_01;
-using Microsoft.Azure.Management.KeyVault;
-using Microsoft.Azure.Graph.RBAC;
-using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
 using Microsoft.Azure.Commands.TestFx;
 using Microsoft.Azure.Commands.TestFx.Recorder;
+using System.Collections.Generic;
+using Xunit.Abstractions;
 
 namespace Microsoft.Azure.Commands.ScenarioTest.SqlTests
 {
@@ -43,7 +29,6 @@ public class SqlTestRunner
             "Microsoft.Sql/servers/databases",
             "Microsoft.Sql/servers/elasticPools"
         };
-        private const string TenantIdKey = "TenantId";
 
         protected readonly ITestRunner TestRunner;
 
@@ -93,87 +78,5 @@ protected SqlTestRunner(ITestOutputHelper output)
                 )
                 .Build();
         }
-
-        protected SqlManagementClient GetSqlClient(MockContext context)
-        {
-            return context.GetServiceClient<SqlManagementClient>();
-        }
-
-        protected SDKMonitor.IMonitorManagementClient GetMonitorManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<SDKMonitor.MonitorManagementClient>();
-        }
-
-        protected CommonMonitor.IMonitorManagementClient GetCommonMonitorManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<CommonMonitor.MonitorManagementClient>();
-        }
-
-        protected IEventHubManagementClient GetEventHubManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<EventHubManagementClient>();
-        }
-
-        protected IOperationalInsightsManagementClient GetOperationalInsightsManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<OperationalInsightsManagementClient>();
-        }
-
-        protected ResourceManagementClient GetResourcesClient(MockContext context)
-        {
-            return context.GetServiceClient<ResourceManagementClient>();
-        }
-
-        protected GraphRbacManagementClient GetGraphClient(MockContext context)
-        {
-            GraphRbacManagementClient graphClient = context.GetServiceClient<GraphRbacManagementClient>();
-            graphClient.BaseUri = TestEnvironmentFactory.GetTestEnvironment().Endpoints.GraphUri;
-            graphClient.TenantID = TestEnvironmentFactory.GetTestEnvironment().TenantId;
-            return graphClient;
-        }
-
-        protected Microsoft.Azure.Graph.RBAC.Version1_6.GraphRbacManagementClient GetGraphClientVersion1_6(MockContext context)
-        {
-            Microsoft.Azure.Graph.RBAC.Version1_6.GraphRbacManagementClient graphClient = context.GetServiceClient<Microsoft.Azure.Graph.RBAC.Version1_6.GraphRbacManagementClient>();
-            graphClient.BaseUri = TestEnvironmentFactory.GetTestEnvironment().Endpoints.GraphUri;
-            string tenantId = null;
-
-            if (HttpMockServer.Mode == HttpRecorderMode.Record)
-            {
-                tenantId = TestEnvironmentFactory.GetTestEnvironment().TenantId;
-                HttpMockServer.Variables[TenantIdKey] = tenantId;
-            }
-            else if (HttpMockServer.Mode == HttpRecorderMode.Playback)
-            {
-                if (HttpMockServer.Variables.ContainsKey(TenantIdKey))
-                {
-                    tenantId = HttpMockServer.Variables[TenantIdKey];
-                }
-            }
-            graphClient.TenantID = tenantId;
-            if (AzureRmProfileProvider.Instance != null &&
-                AzureRmProfileProvider.Instance.Profile != null &&
-                AzureRmProfileProvider.Instance.Profile.DefaultContext != null &&
-                AzureRmProfileProvider.Instance.Profile.DefaultContext.Tenant != null)
-            {
-                AzureRmProfileProvider.Instance.Profile.DefaultContext.Tenant.Id = tenantId;
-            }
-            return graphClient;
-        }
-
-        protected KeyVaultManagementClient GetKeyVaultClient(MockContext context)
-        {
-            return context.GetServiceClient<KeyVaultManagementClient>();
-        }
-
-        protected NetworkManagementClient GetNetworkClient(MockContext context)
-        {
-            return context.GetServiceClient<NetworkManagementClient>();
-        }
-
-        protected static CommonStorage.StorageManagementClient GetStorageManagementClient(MockContext context)
-        {
-            return context.GetServiceClient<CommonStorage.StorageManagementClient>();
-        }
     }
 }
diff --git a/src/Synapse/Synapse.Test/ScenarioTests/SynapseTestRunner.cs b/src/Synapse/Synapse.Test/ScenarioTests/SynapseTestRunner.cs
@@ -16,8 +16,6 @@
 using Microsoft.Azure.Commands.TestFx;
 using Microsoft.Azure.Synapse;
 using Microsoft.Azure.Test.HttpRecorder;
-using Microsoft.Rest;
-using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
 using System.Collections.Generic;
 using Xunit.Abstractions;
 
@@ -71,16 +69,14 @@ protected SynapseTestRunner(ITestOutputHelper output)
                         {"Microsoft.OperationalInsights", null},
                         {"Microsoft.Storage", null}
                     }
-                ).WithManagementClients(
-                    GetSynapseClient
+                )
+                .WithManagementClients(context =>
+                    {
+                        var creds = context.GetClientCredentials(AzureEnvironment.ExtendedEndpoint.AzureSynapseAnalyticsEndpointResourceId);
+                        return new SynapseClient(creds, HttpMockServer.CreateInstance());
+                    }
                 )
                 .Build();
         }
-
-        protected static SynapseClient GetSynapseClient(MockContext context)
-        {
-            var creds = TestEnvironmentFactory.GetTestEnvironment().GetAccessToken(AzureEnvironmentConstants.AzureSynapseAnalyticsEndpointResourceId);
-            return new SynapseClient(creds, HttpMockServer.CreateInstance());
-        }
     }
 }
diff --git a/tools/TestFx/DelegatingHandlers/HttpMockServer.cs b/tools/TestFx/DelegatingHandlers/HttpMockServer.cs
@@ -17,6 +17,8 @@
 using Microsoft.Azure.Commands.ResourceManager.Common;
 using Microsoft.Azure.Commands.TestFx;
 using Microsoft.Azure.Commands.TestFx.Recorder;
+using Microsoft.Azure.PowerShell.Authenticators;
+using Microsoft.Azure.PowerShell.Authenticators.Factories;
 using System;
 using System.Collections.Generic;
 using System.IO;
@@ -72,6 +74,15 @@ public static void Initialize(string callerIdentity, string testIdentity)
             {
                 AzureSession.Instance.ARMContextSaveMode = ContextSaveMode.CurrentUser;
                 ProtectedProfileProvider.InitializeResourceManagerProfile();
+
+                IAuthenticatorBuilder builder = new DefaultAuthenticatorBuilder();
+                AzureSession.Instance.RegisterComponent(AuthenticatorBuilder.AuthenticatorBuilderKey, () => builder);
+
+                PowerShellTokenCacheProvider provider = new SharedTokenCacheProvider();
+                AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => provider);
+                
+                AzureSession.Instance.RegisterComponent(nameof(AzureCredentialFactory), () => new AzureCredentialFactory());
+                //AzureSession.Instance.RegisterComponent(nameof(MsalAccessTokenAcquirerFactory), () => new MsalAccessTokenAcquirerFactory());
             }
             else if (Mode == HttpRecorderMode.Playback)
             {
diff --git a/tools/TestFx/EnvironmentSetupHelper.cs b/tools/TestFx/EnvironmentSetupHelper.cs
diff --git a/tools/TestFx/Mocks/MockContext.cs b/tools/TestFx/Mocks/MockContext.cs
diff --git a/tools/TestFx/TestEnvironment.cs b/tools/TestFx/TestEnvironment.cs
diff --git a/tools/TestFx/TestFx.csproj b/tools/TestFx/TestFx.csproj

Original file line number	Diff line number	Diff line change
`@@ -56,11 +56,11 @@ protected DataLakeStoreTestRunner(ITestOutputHelper output)`
`56`	`56`	`.WithRecordMatcher(`
`57`	`57`	`(ignoreResourcesClient, resourceProviders, userAgentsToIgnore) => new UrlDecodingRecordMatcher(ignoreResourcesClient, resourceProviders, userAgentsToIgnore)`
`58`	`58`	`)`
`59`		`- .WithManagementClients(mockContext =>`
	`59`	`+ .WithManagementClients(context =>`
`60`	`60`	`{`
`61`	`61`	`AdlsClientFactory.IsTest = true;`
`62`		`- var creds = TestEnvironmentFactory.GetTestEnvironment().GetAccessToken(AzureEnvironmentConstants.AzureDataLakeStoreFileSystemEndpointSuffix);`
`63`		`- AdlsClientFactory.CustomDelegatingHAndler = mockContext.AddHandlers(creds, new AdlMockDelegatingHandler());`
	`62`	`+ var creds = context.GetClientCredentials(AzureEnvironment.Endpoint.AzureDataLakeStoreFileSystemEndpointSuffix);`
	`63`	`+ AdlsClientFactory.CustomDelegatingHAndler = context.AddHandlers(creds, new AdlMockDelegatingHandler());`
`64`	`64`	`AdlsClientFactory.MockCredentials = creds;`
`65`	`65`	`return new object();`
`66`	`66`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,6 @@`
`16`	`16`	`using Microsoft.Azure.Commands.TestFx;`
`17`	`17`	`using Microsoft.Azure.KeyVault;`
`18`	`18`	`using Microsoft.Azure.Test.HttpRecorder;`
`19`		`-using Microsoft.Rest.ClientRuntime.Azure.TestFramework;`
`20`	`19`	`using System.Collections.Generic;`
`21`	`20`	`using Xunit.Abstractions;`
`22`	`21`
`@@ -67,16 +66,14 @@ protected NetworkTestRunner(ITestOutputHelper output)`
`67`	`66`	`{"Microsoft.ManagedServiceIdentity", null},`
`68`	`67`	`{"Microsoft.PrivateDns", null},`
`69`	`68`	`}`
`70`		`- ).WithManagementClients(`
`71`		`- GetKeyVaultClient`
	`69`	`+ )`
	`70`	`+ .WithManagementClients(context =>`
	`71`	`+ {`
	`72`	`+ var creds = context.GetClientCredentials(AzureEnvironment.Endpoint.AzureKeyVaultServiceEndpointResourceId);`
	`73`	`+ return new KeyVaultClient(creds, HttpMockServer.CreateInstance());`
	`74`	`+ }`
`72`	`75`	`)`
`73`	`76`	`.Build();`
`74`	`77`	`}`
`75`		`-`
`76`		`- private static KeyVaultClient GetKeyVaultClient(MockContext context)`
`77`		`- {`
`78`		`- var creds = TestEnvironmentFactory.GetTestEnvironment().GetAccessToken(AzureEnvironmentConstants.AzureKeyVaultServiceEndpointResourceId);`
`79`		`- return new KeyVaultClient(creds, HttpMockServer.CreateInstance());`
`80`		`- }`
`81`	`78`	`}`
`82`	`79`	`}`
Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,7 @@ public void RaByResource()`
`91`	`91`	`TestRunner.RunTestScript("Test-RaByResource");`
`92`	`92`	`}`
`93`	`93`
`94`		`- [Fact]//(Skip = "Test indicates Graph call is not recorded when it actually is, refer to https://github.com/Azure/azure-powershell/issues/14632 for more details, test passes in record mode")]`
	`94`	`+ [Fact]`
`95`	`95`	`[Trait(Category.AcceptanceType, Category.CheckIn)]`
`96`	`96`	`public void RaByServicePrincipal()`
`97`	`97`	`{`