Fixed + improved ImdsTests

Author: Robbie-Microsoft

tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsTests.cs

@@ -39,11 +39,10 @@ public async Task ImdsFails404TwiceThenSucceeds200Async(
                     .WithHttpManager(httpManager)
                     .WithRetryPolicyFactory(_testRetryPolicyFactory);
 
-                // Disable cache to avoid pollution
-                
-
                 IManagedIdentityApplication mi = miBuilder.Build();
 
+                ManagedIdentityTests.MockImdsV1Probe(httpManager, ManagedIdentitySource.Imds, userAssignedIdentityId, userAssignedId);
+
                 // Simulate two 404s (to trigger retries), then a successful response
                 const int Num404Errors = 2;
                 for (int i = 0; i < Num404Errors; i++)
@@ -98,11 +97,10 @@ public async Task ImdsFails410FourTimesThenSucceeds200Async(
                     .WithHttpManager(httpManager)
                     .WithRetryPolicyFactory(_testRetryPolicyFactory);
 
-                // Disable cache to avoid pollution
-                
-
                 IManagedIdentityApplication mi = miBuilder.Build();
 
+                ManagedIdentityTests.MockImdsV1Probe(httpManager, ManagedIdentitySource.Imds, userAssignedIdentityId, userAssignedId);
+
                 // Simulate four 410s (to trigger retries), then a successful response
                 const int Num410Errors = 4;
                 for (int i = 0; i < Num410Errors; i++)
@@ -157,11 +155,10 @@ public async Task ImdsFails410PermanentlyAsync(
                     .WithHttpManager(httpManager)
                     .WithRetryPolicyFactory(_testRetryPolicyFactory);
 
-                // Disable cache to avoid pollution
-                
-
                 IManagedIdentityApplication mi = miBuilder.Build();
 
+                ManagedIdentityTests.MockImdsV1Probe(httpManager, ManagedIdentitySource.Imds, userAssignedIdentityId, userAssignedId);
+
                 // Simulate permanent 410s (to trigger the maximum number of retries)
                 const int Num410Errors = 1 + TestImdsRetryPolicy.LinearStrategyNumRetries; // initial request + maximum number of retries
                 for (int i = 0; i < Num410Errors; i++)
@@ -213,11 +210,10 @@ public async Task ImdsFails504PermanentlyAsync(
                     .WithHttpManager(httpManager)
                     .WithRetryPolicyFactory(_testRetryPolicyFactory);
 
-                // Disable cache to avoid pollution
-                
-
                 IManagedIdentityApplication mi = miBuilder.Build();
 
+                ManagedIdentityTests.MockImdsV1Probe(httpManager, ManagedIdentitySource.Imds, userAssignedIdentityId, userAssignedId);
+
                 // Simulate permanent 504s (to trigger the maximum number of retries)
                 const int Num504Errors = 1 + TestImdsRetryPolicy.ExponentialStrategyNumRetries; // initial request + maximum number of retries
                 for (int i = 0; i < Num504Errors; i++)
@@ -269,11 +265,10 @@ public async Task ImdsFails400WhichIsNonRetriableAndRetryPolicyIsNotTriggeredAsy
                     .WithHttpManager(httpManager)
                     .WithRetryPolicyFactory(_testRetryPolicyFactory);
 
-                // Disable cache to avoid pollution
-                
-
                 IManagedIdentityApplication mi = miBuilder.Build();
 
+                ManagedIdentityTests.MockImdsV1Probe(httpManager, ManagedIdentitySource.Imds, userAssignedIdentityId, userAssignedId);
+
                 httpManager.AddManagedIdentityMockHandler(
                     ManagedIdentityTests.ImdsEndpoint,
                     ManagedIdentityTests.Resource,
@@ -321,11 +316,10 @@ public async Task ImdsFails500AndRetryPolicyIsDisabledAndNotTriggeredAsync(
                     .WithHttpManager(httpManager)
                     .WithRetryPolicyFactory(_testRetryPolicyFactory);
 
-                // Disable cache to avoid pollution
-                
-
                 IManagedIdentityApplication mi = miBuilder.Build();
 
+                ManagedIdentityTests.MockImdsV1Probe(httpManager, ManagedIdentitySource.Imds, userAssignedIdentityId, userAssignedId);
+
                 httpManager.AddManagedIdentityMockHandler(
                     ManagedIdentityTests.ImdsEndpoint,
                     ManagedIdentityTests.Resource,
@@ -367,11 +361,10 @@ public async Task ImdsRetryPolicyLifeTimeIsPerRequestAsync()
                     .WithHttpManager(httpManager)
                     .WithRetryPolicyFactory(_testRetryPolicyFactory);
 
-                // Disable cache to avoid pollution
-                
-
                 IManagedIdentityApplication mi = miBuilder.Build();
 
+                ManagedIdentityTests.MockImdsV1Probe(httpManager, ManagedIdentitySource.Imds);
+
                 // Simulate permanent errors (to trigger the maximum number of retries)
                 const int Num504Errors = 1 + TestImdsRetryPolicy.ExponentialStrategyNumRetries; // initial request + maximum number of retries
                 for (int i = 0; i < Num504Errors; i++)

tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ImdsV2Tests.cs

@@ -90,7 +90,7 @@ private async Task<IManagedIdentityApplication> CreateManagedIdentityAsync(
             bool addProbeMock = true,
             bool addSourceCheck = true,
             ManagedIdentityKeyType managedIdentityKeyType = ManagedIdentityKeyType.InMemory,
-            bool imdsV2 = true) // false indicates imdsV1
+            ImdsVersion imdsVersion = ImdsVersion.V2)
         {
             ManagedIdentityApplicationBuilder miBuilder = null;
 
@@ -108,14 +108,14 @@ private async Task<IManagedIdentityApplication> CreateManagedIdentityAsync(
                 .WithHttpManager(httpManager)
                 .WithRetryPolicyFactory(_testRetryPolicyFactory);
 
-            if (imdsV2)
+            if (imdsVersion == ImdsVersion.V2)
             {
                 miBuilder.WithCsrFactory(_testCsrFactory);
             }
 
             var managedIdentityApp = miBuilder.Build();
 
-            if (!imdsV2)
+            if (imdsVersion == ImdsVersion.V1)
             {
                 httpManager.AddMockHandler(MockHelpers.MockImdsProbeFailure(ImdsVersion.V2, userAssignedIdentityId, userAssignedId));
                 httpManager.AddMockHandler(MockHelpers.MockImdsProbe(ImdsVersion.V1, userAssignedIdentityId, userAssignedId));
@@ -345,7 +345,7 @@ public async Task ImdsV2EndpointsAreNotAvailableButMtlsPopTokenWasRequested(
             {
                 SetEnvironmentVariables(ManagedIdentitySource.Imds, TestConstants.ImdsEndpoint);
 
-                var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId, imdsV2: false).ConfigureAwait(false);
+                var managedIdentityApp = await CreateManagedIdentityAsync(httpManager, userAssignedIdentityId, userAssignedId, imdsVersion: ImdsVersion.V1).ConfigureAwait(false);
 
                 var ex = await Assert.ThrowsExceptionAsync<MsalClientException>(async () =>
                     await managedIdentityApp.AcquireTokenForManagedIdentity(ManagedIdentityTests.Resource)

tests/Microsoft.Identity.Test.Unit/ManagedIdentityTests/ManagedIdentityTests.cs

@@ -43,7 +43,7 @@ public class ManagedIdentityTests : TestBase
         private readonly TestRetryPolicyFactory _testRetryPolicyFactory = new TestRetryPolicyFactory();
 
         // MtlsPop is disabled for all these tests, so no need to mock IMDSv2 probe here
-        private void MockImdsV1Probe(
+        internal static void MockImdsV1Probe(
             MockHttpManager httpManager,
             ManagedIdentitySource managedIdentitySource,
             UserAssignedIdentityId userAssignedIdentityId = UserAssignedIdentityId.None,