Fixed + Improved all tests in ImdsV2Tests.cs

Author: Robbie-Microsoft

src/client/Microsoft.Identity.Client/ManagedIdentity/V2/ImdsV2ManagedIdentitySource.cs

@@ -66,14 +66,14 @@ public static async Task<CsrMetadata> GetCsrMetadataAsync(RequestContext request
             }
             catch (Exception ex)
             {
-                ThrowProbeFailedException(
+                ThrowCsrMetadataRequestException(
                     "ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed.",
                     ex);
             }
 
             if (response.StatusCode != HttpStatusCode.OK)
             {
-                ThrowProbeFailedException(
+                ThrowCsrMetadataRequestException(
                     $"ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed due to HTTP error. Status code: {response.StatusCode} Body: {response.Body}",
                     null,
                     (int)response.StatusCode);
@@ -87,7 +87,7 @@ public static async Task<CsrMetadata> GetCsrMetadataAsync(RequestContext request
             return TryCreateCsrMetadata(response, requestContext.Logger);
         }
 
-        private static void ThrowProbeFailedException(
+        private static void ThrowCsrMetadataRequestException(
             String errorMessage,
             Exception ex = null,
             int? statusCode = null)
@@ -111,15 +111,15 @@ private static bool ValidateCsrMetadataResponse(
 
             if (serverHeader == null)
             {
-                ThrowProbeFailedException(
+                ThrowCsrMetadataRequestException(
                     $"ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed because response doesn't have server header. Status code: {response.StatusCode} Body: {response.Body}",
                     null,
                     (int)response.StatusCode);
             }
 
             if (!serverHeader.Contains("IMDS", StringComparison.OrdinalIgnoreCase))
             {
-                ThrowProbeFailedException(
+                ThrowCsrMetadataRequestException(
                     $"ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed because the 'server' header format is invalid. Extracted server header: {serverHeader}. Status code: {response.StatusCode} Body: {response.Body}",
                     null,
                     (int)response.StatusCode);
@@ -135,7 +135,7 @@ private static CsrMetadata TryCreateCsrMetadata(
             CsrMetadata csrMetadata = JsonHelper.DeserializeFromJson<CsrMetadata>(response.Body);
             if (!CsrMetadata.ValidateCsrMetadata(csrMetadata))
             {
-                ThrowProbeFailedException(
+                ThrowCsrMetadataRequestException(
                     $"ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed because the CsrMetadata response is invalid. Status code: {response.StatusCode} Body: {response.Body}",
                     null,
                     (int)response.StatusCode);

tests/Microsoft.Identity.Test.Common/Core/Mocks/MockHelpers.cs

@@ -618,7 +618,8 @@ public static MockHttpMessageHandler MockImdsProbe(
             ImdsVersion imdsVersion,
             UserAssignedIdentityId userAssignedIdentityId = UserAssignedIdentityId.None,
             string userAssignedId = null,
-            bool success = true)
+            bool success = true,
+            bool retry = false)
         {
             string apiVersionQueryParam;
             string imdsApiVersion;
@@ -650,7 +651,14 @@ public static MockHttpMessageHandler MockImdsProbe(
             }
             else
             {
-                statusCode = HttpStatusCode.NotFound;
+                if (retry)
+                {
+                    statusCode = HttpStatusCode.InternalServerError;
+                }
+                else
+                {
+                    statusCode = HttpStatusCode.NotFound;
+                }
             }
 
             IDictionary<string, string> expectedQueryParams = new Dictionary<string, string>();
@@ -687,9 +695,10 @@ public static MockHttpMessageHandler MockImdsProbe(
         public static MockHttpMessageHandler MockImdsProbeFailure(
             ImdsVersion imdsVersion,
             UserAssignedIdentityId userAssignedIdentityId = UserAssignedIdentityId.None,
-            string userAssignedId = null)
+            string userAssignedId = null,
+            bool retry = false)
         {
-            return MockImdsProbe(imdsVersion, userAssignedIdentityId, userAssignedId, success: false);
+            return MockImdsProbe(imdsVersion, userAssignedIdentityId, userAssignedId, success: false, retry: retry);
         }
 
         public static MockHttpMessageHandler MockCsrResponse(