Migrate from ADAL to MSAL for TestFx (#23730)

vidai-msft · web-flow · commit 3df8a49d1289 · 2023-12-14T09:42:09.000+08:00
diff --git a/src/Accounts/Accounts.Test/Mocks/MockSubscriptionClientFactoryVersion2019.cs b/src/Accounts/Accounts.Test/Mocks/MockSubscriptionClientFactoryVersion2019.cs
@@ -16,14 +16,12 @@
 using Microsoft.Azure.Management.ResourceManager.Version2021_01_01.Models;
 using Microsoft.Rest.Azure;
 using Moq;
+using Newtonsoft.Json;
 using System;
 using System.Collections.Generic;
 using System.Linq;
-using System.Management.Automation.Language;
 using System.Threading;
 using System.Threading.Tasks;
-using Newtonsoft.Json;
-using Microsoft.IdentityModel.Tokens;
 
 namespace Microsoft.Azure.Commands.ResourceManager.Common.Test
 {
diff --git a/src/Accounts/Accounts.Test/SubscriptionClientSwitchTest.cs b/src/Accounts/Accounts.Test/SubscriptionClientSwitchTest.cs
@@ -123,7 +123,7 @@ public void SubscriptionClientListTenantFallback()
                     {
                         subscriptionClients.Dequeue();
                     }
-                    throw e;
+                    throw;
                 }
                 return Task.FromResult(result);
             };
@@ -193,7 +193,7 @@ public void SubscriptionClientGetSubFallback()
                     {
                         subscriptionClients.Dequeue();
                     }
-                    throw e;
+                    throw;
                 }
                 return Task.FromResult(result);
             };
@@ -265,7 +265,7 @@ public void SubscriptionClientListSubFallback()
                     {
                         subscriptionClients.Dequeue();
                     }
-                    throw e;
+                    throw;
                 }
                 return Task.FromResult(result);
             };
diff --git a/src/Attestation/Attestation.Test/ScenarioTests/AttestationTestRunner.cs b/src/Attestation/Attestation.Test/ScenarioTests/AttestationTestRunner.cs
@@ -13,14 +13,12 @@
 // ----------------------------------------------------------------------------------
 
 using Microsoft.Azure.Attestation;
+using Microsoft.Azure.Commands.TestFx;
 using Microsoft.Azure.Management.Attestation;
-using Microsoft.Azure.Test.HttpRecorder;
-using System;
-using System.Collections.Generic;
 using Microsoft.Azure.Management.Internal.Resources;
-using Microsoft.IdentityModel.Clients.ActiveDirectory;
+using Microsoft.Azure.Test.HttpRecorder;
 using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
-using Microsoft.Azure.Commands.TestFx;
+using System.Collections.Generic;
 using Xunit.Abstractions;
 
 namespace Microsoft.Azure.Commands.Attestation.Test.ScenarioTests
@@ -76,22 +74,12 @@ private static AttestationManagementClient GetAttestationManagementClient(MockCo
 
         private static AttestationClient GetAttestationClient(MockContext context)
         {
-            string environmentConnectionString = Environment.GetEnvironmentVariable("TEST_CSM_ORGID_AUTHENTICATION");
             string accessToken = "fakefakefake";
 
             // When recording, we should have a connection string passed into the code from the environment
-            if (!string.IsNullOrEmpty(environmentConnectionString))
+            if (HttpMockServer.Mode == HttpRecorderMode.Record)
             {
-                // Gather test client credential information from the environment
-                var connectionInfo = new ConnectionString(Environment.GetEnvironmentVariable("TEST_CSM_ORGID_AUTHENTICATION"));
-                string servicePrincipal = connectionInfo.GetValue<string>(ConnectionStringKeys.ServicePrincipalKey);
-                string servicePrincipalSecret = connectionInfo.GetValue<string>(ConnectionStringKeys.ServicePrincipalSecretKey);
-                string aadTenant = connectionInfo.GetValue<string>(ConnectionStringKeys.TenantIdKey);
-
-                // Create credentials
-                var clientCredentials = new ClientCredential(servicePrincipal, servicePrincipalSecret);
-                var authContext = new AuthenticationContext($"https://login.windows.net/{aadTenant}", TokenCache.DefaultShared);
-                accessToken = authContext.AcquireTokenAsync("https://attest.azure.net", clientCredentials).Result.AccessToken;
+                accessToken = TestEnvironmentFactory.GetTestEnvironment().GetServicePrincipalAccessToken(new[] { "https://attest.azure.net/.default" });
             }
 
             return new AttestationClient(new AttestationCredentials(accessToken), HttpMockServer.CreateInstance());
diff --git a/src/Attestation/Attestation.sln b/src/Attestation/Attestation.sln
@@ -84,9 +84,7 @@ Global
 		{3B31A547-8D98-4DE2-B314-885C82474A42} = {C9F34834-F918-4085-A8DE-F9E38FB03AF5}
 		{D3F5BCD1-1B75-421E-8C3A-B16F211680F6} = {C9F34834-F918-4085-A8DE-F9E38FB03AF5}
 		{A52A877E-B6BA-410D-ADEE-0EE969734C50} = {C9F34834-F918-4085-A8DE-F9E38FB03AF5}
-		{0F29534D-1078-4FBC-A65D-913CC6D62D7A} = {E3DA12B6-0025-4624-805C-F3406CA90122}
 		{31B49528-4241-4DDB-8BB9-28F95148C602} = {E3DA12B6-0025-4624-805C-F3406CA90122}
-		{B73D384F-62CA-4FD2-89C6-04DB2DC5AB9D} = {E3DA12B6-0025-4624-805C-F3406CA90122}
 		{D8EAB1F1-D2A9-49F9-BC09-7397F9CCFF10} = {E3DA12B6-0025-4624-805C-F3406CA90122}
 	EndGlobalSection
 EndGlobal
diff --git a/src/Batch/Batch.Test/ScenarioTests/ScenarioTestHelpers.cs b/src/Batch/Batch.Test/ScenarioTests/ScenarioTestHelpers.cs
@@ -238,7 +238,7 @@ public static void CreatePoolIfNotExists(
             {
                 if (e.RequestInformation.BatchError.Code != "PoolExists")
                 {
-                    throw e;
+                    throw;
                 }
             }
         }
diff --git a/src/Network/Network.Test/ScenarioTests/NetworkTestRunner.cs b/src/Network/Network.Test/ScenarioTests/NetworkTestRunner.cs
@@ -15,7 +15,6 @@
 using Microsoft.Azure.Commands.TestFx;
 using Microsoft.Azure.KeyVault;
 using Microsoft.Azure.Test.HttpRecorder;
-using Microsoft.IdentityModel.Clients.ActiveDirectory;
 using Microsoft.Rest;
 using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
 using System;
@@ -77,26 +76,12 @@ protected NetworkTestRunner(ITestOutputHelper output)
 
         private static KeyVaultClient GetKeyVaultClient(MockContext context)
         {
-            string environmentConnectionString = Environment.GetEnvironmentVariable("TEST_CSM_ORGID_AUTHENTICATION");
             string accessToken = "fakefakefake";
 
             // When recording, we should have a connection string passed into the code from the environment
-            if (!string.IsNullOrEmpty(environmentConnectionString))
+            if (HttpMockServer.Mode == HttpRecorderMode.Record)
             {
-                // Gather test client credential information from the environment
-                var connectionInfo = new ConnectionString(Environment.GetEnvironmentVariable("TEST_CSM_ORGID_AUTHENTICATION"));
-                var mode = connectionInfo.GetValue<string>(ConnectionStringKeys.HttpRecorderModeKey);
-                if (mode == HttpRecorderMode.Record.ToString())
-                {
-                    string servicePrincipal = connectionInfo.GetValue<string>(ConnectionStringKeys.ServicePrincipalKey);
-                    string servicePrincipalSecret = connectionInfo.GetValue<string>(ConnectionStringKeys.ServicePrincipalSecretKey);
-                    string aadTenant = connectionInfo.GetValue<string>(ConnectionStringKeys.TenantIdKey);
-
-                    // Create credentials
-                    var clientCredentials = new ClientCredential(servicePrincipal, servicePrincipalSecret);
-                    var authContext = new AuthenticationContext($"https://login.windows.net/{aadTenant}", TokenCache.DefaultShared);
-                    accessToken = authContext.AcquireTokenAsync("https://vault.azure.net", clientCredentials).Result.AccessToken;
-                }
+                accessToken = TestEnvironmentFactory.GetTestEnvironment().GetServicePrincipalAccessToken(new[] { "https://vault.azure.net/.default" });
             }
 
             return new KeyVaultClient(new TokenCredentials(accessToken), HttpMockServer.CreateInstance());
diff --git a/src/Resources/Resources.Test/ScenarioTests/ResourcesTestRunner.cs b/src/Resources/Resources.Test/ScenarioTests/ResourcesTestRunner.cs
@@ -62,7 +62,7 @@ protected ResourcesTestRunner(ITestOutputHelper output)
                         ? new Func<ServiceClientCredentials>(() =>
                         {
                             var testEnvironment = TestEnvironmentFactory.GetTestEnvironment();
-                            return testEnvironment.TokenInfo[TokenAudience.Management];
+                            return testEnvironment.TokenInfo[TokenAudience.Graph];
                         })()
                         : new TokenCredentials("foo");
 
diff --git a/src/Synapse/Synapse.Test/ScenarioTests/SynapseTestRunner.cs b/src/Synapse/Synapse.Test/ScenarioTests/SynapseTestRunner.cs
@@ -12,23 +12,12 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-
+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;
 using System.Collections.Generic;
-using NewResourceManagementClient = Microsoft.Azure.Management.Internal.Resources.ResourceManagementClient;
-using Microsoft.Azure.Management.Synapse;
-using Microsoft.Azure.Synapse;
-using Microsoft.IdentityModel.Clients.ActiveDirectory;
-using Microsoft.Rest;
-using NewStorageManagementClient = Microsoft.Azure.Management.Storage.Version2017_10_01.StorageManagementClient;
-using Microsoft.Azure.Management.Storage;
-using SDKMonitor = Microsoft.Azure.Management.Monitor;
-using CommonMonitor = Microsoft.Azure.Management.Monitor.Version2018_09_01;
-using Microsoft.Azure.Management.OperationalInsights;
-using Microsoft.Azure.Management.EventHub;
-using Microsoft.Azure.Commands.TestFx;
 using Xunit.Abstractions;
 
 namespace Microsoft.Azure.Commands.Synapse.Test.ScenarioTests
@@ -89,22 +78,12 @@ protected SynapseTestRunner(ITestOutputHelper output)
 
         protected static SynapseClient GetSynapseClient(MockContext context)
         {
-            string environmentConnectionString = Environment.GetEnvironmentVariable("TEST_CSM_ORGID_AUTHENTICATION");
             string accessToken = "fakefakefake";
 
             // When recording, we should have a connection string passed into the code from the environment
-            if (!string.IsNullOrEmpty(environmentConnectionString))
+            if (HttpMockServer.Mode == HttpRecorderMode.Record)
             {
-                // Gather test client credential information from the environment
-                var connectionInfo = new ConnectionString(Environment.GetEnvironmentVariable("TEST_CSM_ORGID_AUTHENTICATION"));
-                string servicePrincipal = connectionInfo.GetValue<string>(ConnectionStringKeys.ServicePrincipalKey);
-                string servicePrincipalSecret = connectionInfo.GetValue<string>(ConnectionStringKeys.ServicePrincipalSecretKey);
-                string aadTenant = connectionInfo.GetValue<string>(ConnectionStringKeys.TenantIdKey);
-
-                // Create credentials
-                var clientCredentials = new ClientCredential(servicePrincipal, servicePrincipalSecret);
-                var authContext = new AuthenticationContext($"https://login.windows.net/{aadTenant}", TokenCache.DefaultShared);
-                accessToken = authContext.AcquireTokenAsync("https://dev.azuresynapse.net", clientCredentials).Result.AccessToken;
+                accessToken = TestEnvironmentFactory.GetTestEnvironment().GetServicePrincipalAccessToken(new[] { "https://dev.azuresynapse.net/.default" });
             }
 
             return new SynapseClient(new TokenCredentials(accessToken), HttpMockServer.CreateInstance());
diff --git a/tools/Common.Netcore.Dependencies.Test.targets b/tools/Common.Netcore.Dependencies.Test.targets
@@ -10,7 +10,6 @@
 
   <ItemGroup>
     <PackageReference Include="Microsoft.Azure.Management.ResourceManager" Version="3.17.3-preview" />
-    <PackageReference Include="Microsoft.Rest.ClientRuntime.Azure.Authentication" Version="2.3.7" />
     <PackageReference Include="xunit" Version="2.4.2" />
   </ItemGroup>
 
@@ -40,7 +39,7 @@
   </ItemGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.IdentityModel.Clients.ActiveDirectory" Version="3.19.2" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.5.0" />
+    <PackageReference Include="Microsoft.Identity.Client" Version="4.49.1" />
   </ItemGroup>
 </Project>
diff --git a/tools/TestFx/EnvironmentSetupHelper.cs b/tools/TestFx/EnvironmentSetupHelper.cs
@@ -406,6 +406,7 @@ public void SetupAzureEnvironmentFromEnvironmentVariables(AzureModule mode)
             environment.StorageEndpointSuffix = AzureEnvironmentConstants.AzureStorageEndpointSuffix;
             environment.AzureKeyVaultDnsSuffix = AzureEnvironmentConstants.AzureKeyVaultDnsSuffix;
             environment.AzureKeyVaultServiceEndpointResourceId = AzureEnvironmentConstants.AzureKeyVaultServiceEndpointResourceId;
+            environment.ExtendedProperties.SetProperty(AzureEnvironment.ExtendedEndpoint.MicrosoftGraphUrl, currentEnvironment.Endpoints.GraphUri.AbsoluteUri);
             environment.ExtendedProperties.SetProperty(AzureEnvironment.ExtendedEndpoint.OperationalInsightsEndpoint, "https://api.loganalytics.io/v1");
             environment.ExtendedProperties.SetProperty(AzureEnvironment.ExtendedEndpoint.OperationalInsightsEndpointResourceId, "https://api.loganalytics.io");
             if (!AzureRmProfileProvider.Instance.GetProfile<AzureRmProfile>().EnvironmentTable.ContainsKey(TestEnvironmentName))
diff --git a/tools/TestFx/Mocks/MockClientFactory.cs b/tools/TestFx/Mocks/MockClientFactory.cs
@@ -31,6 +31,8 @@
 using Microsoft.Azure.Commands.TestFx.DelegatingHandlers;
 using Microsoft.Rest.ClientRuntime.Azure.TestFramework;
 using Microsoft.Azure.Test.HttpRecorder;
+using Microsoft.Azure.Commands.Common.MSGraph.Version1_0;
+
 #if NETSTANDARD
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Core;
 #endif
@@ -81,9 +83,16 @@ public TClient CreateArmClient<TClient>(IAzureContext context, string endpoint)
 
         public TClient CreateCustomArmClient<TClient>(params object[] parameters) where TClient : Microsoft.Rest.ServiceClient<TClient>
         {
-            if (!(_serviceClients.FirstOrDefault(o => o is TClient) is TClient client))
+            if (_serviceClients.FirstOrDefault(o => o is TClient) is not TClient client)
             {
-                client = _mockContext?.GetServiceClient<TClient>();
+                if (typeof(TClient) == typeof(MicrosoftGraphClient))
+                {
+                    client = _mockContext?.GetGraphServiceClient<TClient>();
+                }
+                else
+                {
+                    client = _mockContext?.GetServiceClient<TClient>();
+                }
             }
 
             if (client == null)
diff --git a/tools/TestFx/Mocks/MockContext.cs b/tools/TestFx/Mocks/MockContext.cs
diff --git a/tools/TestFx/TestEndpoints.cs b/tools/TestFx/TestEndpoints.cs
diff --git a/tools/TestFx/TestEnvironment.cs b/tools/TestFx/TestEnvironment.cs

Original file line number	Diff line number	Diff line change
`@@ -123,7 +123,7 @@ public void SubscriptionClientListTenantFallback()`
`123`	`123`	`{`
`124`	`124`	`subscriptionClients.Dequeue();`
`125`	`125`	`}`
`126`		`- throw e;`
	`126`	`+ throw;`
`127`	`127`	`}`
`128`	`128`	`return Task.FromResult(result);`
`129`	`129`	`};`
`@@ -193,7 +193,7 @@ public void SubscriptionClientGetSubFallback()`
`193`	`193`	`{`
`194`	`194`	`subscriptionClients.Dequeue();`
`195`	`195`	`}`
`196`		`- throw e;`
	`196`	`+ throw;`
`197`	`197`	`}`
`198`	`198`	`return Task.FromResult(result);`
`199`	`199`	`};`
`@@ -265,7 +265,7 @@ public void SubscriptionClientListSubFallback()`
`265`	`265`	`{`
`266`	`266`	`subscriptionClients.Dequeue();`
`267`	`267`	`}`
`268`		`- throw e;`
	`268`	`+ throw;`
`269`	`269`	`}`
`270`	`270`	`return Task.FromResult(result);`
`271`	`271`	`};`
Original file line number	Diff line number	Diff line change
`@@ -238,7 +238,7 @@ public static void CreatePoolIfNotExists(`
`238`	`238`	`{`
`239`	`239`	`if (e.RequestInformation.BatchError.Code != "PoolExists")`
`240`	`240`	`{`
`241`		`- throw e;`
	`241`	`+ throw;`
`242`	`242`	`}`
`243`	`243`	`}`
`244`	`244`	`}`
Original file line number	Diff line number	Diff line change
`@@ -62,7 +62,7 @@ protected ResourcesTestRunner(ITestOutputHelper output)`
`62`	`62`	`? new Func<ServiceClientCredentials>(() =>`
`63`	`63`	`{`
`64`	`64`	`var testEnvironment = TestEnvironmentFactory.GetTestEnvironment();`
`65`		`- return testEnvironment.TokenInfo[TokenAudience.Management];`
	`65`	`+ return testEnvironment.TokenInfo[TokenAudience.Graph];`
`66`	`66`	`})()`
`67`	`67`	`: new TokenCredentials("foo");`
`68`	`68`