Migrate Identity Live tests deployment to Fed Auth.

eng/pipelines/templates/jobs/live.tests.yml

@@ -4,6 +4,7 @@ parameters:
   EnvVars: {}
   MaxParallel: 0
   PreSteps: []
+  PreTestRunSteps: []
   PostSteps: []
   TimeoutInMinutes: 60
   TestMode: 'LIVE'
@@ -24,6 +25,7 @@ parameters:
   UseHttpFaultInjector: false
   OSName:
   UseFederatedAuth: false
+  PersistOidcToken: false
   VersionOverride: $(VersionOverride)
 
 jobs:
@@ -114,6 +116,7 @@ jobs:
               SubscriptionConfiguration: $(SubscriptionConfiguration)
               ArmTemplateParameters: $(ArmTemplateParameters)
               UseFederatedAuth: ${{ parameters.UseFederatedAuth }}
+              PersistOidcToken: ${{ parameters.PersistOidcToken }}
               ServiceConnection: ${{ parameters.CloudConfig.ServiceConnection }}
               SubscriptionConfigurationFilePaths: ${{ parameters.CloudConfig.SubscriptionConfigurationFilePaths }}
               EnvVars:
@@ -124,6 +127,7 @@ jobs:
 
       - template: /eng/pipelines/templates/steps/build-and-test.yml
         parameters:
+          PreTestRunSteps: ${{ parameters.PreTestRunSteps }}
           ParallelTestPlayback: 'false'
           BuildParallelization: ${{ parameters.BuildParallelization }}
           BuildOptions: ${{ parameters.BuildOptions }}

eng/pipelines/templates/stages/archetype-sdk-tests-isolated.yml

@@ -26,6 +26,9 @@ parameters:
 - name: DisableAzureResourceCreation
   type: boolean
   default: false
+- name: PreTestRunSteps
+  type: object
+  default: []
 - name: PreSteps
   type: object
   default: []
@@ -115,6 +118,9 @@ parameters:
 - name: AdditionalSparseCheckoutPaths
   type: object
   default: []
+- name: PersistOidcToken
+  type: boolean
+  default: false
 
 stages:
 - ${{ each cloud in parameters.CloudConfig }}:
@@ -144,6 +150,8 @@ stages:
               - ${{ parameters.AdditionalSparseCheckoutPaths }}
             JobTemplatePath: /eng/pipelines/templates/jobs/live.tests.yml
             AdditionalParameters:
+              PreTestRunSteps:
+                - ${{ parameters.PreTestRunSteps }}
               PreSteps:
                 - ${{ parameters.PlatformPreSteps }}
                 - ${{ parameters.PreSteps }}
@@ -165,6 +173,7 @@ stages:
               TestOptions: ${{ parameters.TestOptions }}
               UseHttpFaultInjector: ${{ parameters.UseHttpFaultInjector }}
               UseFederatedAuth: ${{ parameters.UseFederatedAuth }}
+              PersistOidcToken: ${{ parameters.PersistOidcToken }}
 
             MatrixConfigs:
               # Enumerate platforms and additional platforms based on supported clouds (sparse platform<-->cloud matrix).

eng/pipelines/templates/stages/archetype-sdk-tests.yml

@@ -106,6 +106,9 @@ parameters:
   - name: UseFederatedAuth
     type: boolean
     default: true
+  - name: PersistOidcToken
+    type: boolean
+    default: false
 
 extends:
   template: /eng/pipelines/templates/stages/1es-redirect.yml

eng/pipelines/templates/steps/build-and-test.yml

@@ -41,6 +41,9 @@ parameters:
   - name: ServiceConnection
     type: string
     default: ''
+  - name: PreTestRunSteps
+    type: object
+    default: [ ]
 
 steps:
   - task: Maven@4
@@ -83,6 +86,8 @@ steps:
       AZURE_VERSION_OVERRIDE_TESTS: ${{ parameters.TestVersionSupport }}
     condition: and(succeeded(), eq(variables['TestVersionSupport'], 'true'))
 
+  - ${{ parameters.PreTestRunSteps }}
+
   - ${{ if eq('true', parameters.UseFederatedAuth) }}:
     # Remove preceeding "1." in the version specified by $(JavaTestVersion)
     # e.g. 1.8 -> 8
@@ -176,6 +181,8 @@ steps:
       publishJUnitResults: false
     condition: and(succeeded(), eq(variables['TestFromSource'], 'true'))
 
+  - ${{ parameters.PreTestRunSteps }}
+
   - ${{ if eq('true', parameters.UseFederatedAuth) }}:
     - ${{ if eq(parameters.IsLiveTest, 'true') }}:
       - pwsh: |

sdk/identity/azure-identity/src/test/java/LiveManagedIdentityTests.java

@@ -7,12 +7,12 @@
 import com.azure.core.http.HttpMethod;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpResponse;
-import com.azure.core.test.TestProxyTestBase;
 import com.azure.core.util.Configuration;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.logging.LogLevel;
 import com.azure.identity.ClientSecretCredential;
 import com.azure.identity.ClientSecretCredentialBuilder;
+import com.azure.identity.implementation.Retry;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
@@ -28,11 +28,13 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
-public class LiveManagedIdentityTests extends TestProxyTestBase {
+public class LiveManagedIdentityTests {
 
     @Test
     @EnabledIfEnvironmentVariable(named = "AZURE_TEST_MODE", matches = "LIVE")
+    @Retry(maxRetries = 3)
     public void testManagedIdentityFuncDeployment() {
+
         HttpClient client = HttpClient.createDefault();
         String functionUrl = "https://" + System.getenv("IDENTITY_FUNCTION_NAME") + ".azurewebsites.net/api/mitest";
         HttpRequest request = new HttpRequest(HttpMethod.GET, functionUrl);
@@ -45,6 +47,7 @@ public void testManagedIdentityFuncDeployment() {
 
     @Test
     @EnabledIfEnvironmentVariable(named = "AZURE_TEST_MODE", matches = "LIVE")
+    @Retry(maxRetries = 3)
     public void testManagedIdentityWebAppDeployment() {
         HttpClient client = HttpClient.createDefault();
         String functionUrl = "https://" + System.getenv("IDENTITY_WEBAPP_NAME") + ".azurewebsites.net/mitest";
@@ -62,30 +65,19 @@ public void testManagedIdentityWebAppDeployment() {
     @EnabledIfEnvironmentVariable(named = "AZURE_TEST_MODE", matches = "LIVE")
     @EnabledIfSystemProperty(named = "os.name", matches = "Linux")
     @Timeout(value = 15, unit = TimeUnit.MINUTES)
+    @Retry(maxRetries = 3)
     public void testManagedIdentityAksDeployment() {
-
+        System.out.println("Environment: " + System.getenv("IDENTITY_ENVIRONMENT"));
         String os = System.getProperty("os.name");
         System.out.println("OS: " + os);
         //Setup Env
         Configuration configuration = Configuration.getGlobalConfiguration().clone();
 
-        String spClientId = configuration.get("IDENTITY_CLIENT_ID");
-        String secret = configuration.get("IDENTITY_CLIENT_SECRET");
-        String tenantId = configuration.get("IDENTITY_TENANT_ID");
-        String resourceGroup = configuration.get("IDENTITY_RESOURCE_GROUP");
-        String aksCluster = configuration.get("IDENTITY_AKS_CLUSTER_NAME");
-        String subscriptionId = configuration.get("IDENTITY_SUBSCRIPTION_ID");
         String podName = configuration.get("IDENTITY_AKS_POD_NAME");
         String pathCommand = os.contains("Windows") ? "where" : "which";
 
-        String azPath = runCommand(pathCommand, "az").trim();
         String kubectlPath = runCommand(pathCommand, "kubectl").trim();
 
-        runCommand(azPath, "login", "--service-principal", "-u", spClientId, "-p", secret, "--tenant", tenantId);
-        runCommand(azPath, "account", "set", "--subscription", subscriptionId);
-        runCommand(azPath, "aks", "get-credentials", "--resource-group", resourceGroup, "--name", aksCluster,
-            "--overwrite-existing");
-
         String podOutput = runCommand(kubectlPath, "get", "pods", "-o", "jsonpath='{.items[0].metadata.name}'");
         assertTrue(podOutput.contains(podName), "Pod name not found in the output");
 
@@ -96,18 +88,20 @@ public void testManagedIdentityAksDeployment() {
 
     @Test
     @EnabledIfEnvironmentVariable(named = "AZURE_TEST_MODE", matches = "LIVE")
+    @EnabledIfEnvironmentVariable(named = "IDENTITY_ENVIRONMENT", matches = "AzureCloud")
     @EnabledIfSystemProperty(named = "os.name", matches = "Linux")
     @Timeout(value = 15, unit = TimeUnit.MINUTES)
+    @Retry(maxRetries = 3)
     public void testManagedIdentityVmDeployment() {
-
+        System.out.println("Environment: " + System.getenv("IDENTITY_ENVIRONMENT"));
         String os = System.getProperty("os.name");
         System.out.println("OS: " + os);
         //Setup Env
         Configuration configuration = Configuration.getGlobalConfiguration().clone();
 
-        String spClientId = configuration.get("IDENTITY_CLIENT_ID");
-        String secret = configuration.get("IDENTITY_CLIENT_SECRET");
-        String tenantId = configuration.get("IDENTITY_TENANT_ID");
+        String spClientId = configuration.get("AZURESUBSCRIPTION_CLIENT_ID");
+        String oidc = configuration.get("AZ_OIDC_TOKEN");
+        String tenantId = configuration.get("AZURESUBSCRIPTION_TENANT_ID");
         String resourceGroup = configuration.get("IDENTITY_RESOURCE_GROUP");
         String subscriptionId = configuration.get("IDENTITY_SUBSCRIPTION_ID");
         String vmName = configuration.get("IDENTITY_VM_NAME");
@@ -118,7 +112,8 @@ public void testManagedIdentityVmDeployment() {
         String azPath = runCommand(isWindows ? "where" : "which", "az").trim();
         azPath = isWindows ? extractAzCmdPath(azPath) : azPath;
 
-        runCommand(azPath, "login", "--service-principal", "-u", spClientId, "-p", secret, "--tenant", tenantId);
+        runCommand(azPath, "login", "--federated-token", oidc, "--service-principal", "-u", spClientId, "--tenant",
+            tenantId);
         runCommand(azPath, "account", "set", "--subscription", subscriptionId);
 
         String storageKey = runCommand(azPath, "storage", "account", "keys", "list", "--account-name",
@@ -143,11 +138,12 @@ public void testManagedIdentityVmDeployment() {
 
     @Test
     @EnabledIfEnvironmentVariable(named = "AZURE_TEST_MODE", matches = "LIVE")
+    @Retry(maxRetries = 3)
     public void callGraphWithClientSecret() {
         Configuration configuration = Configuration.getGlobalConfiguration().clone();
 
-        String multiTenantId = configuration.get("AZURE_IDENTITY_MULTI_TENANT_TENANT_ID");
-        String multiClientId = configuration.get("AZURE_IDENTITY_MULTI_TENANT_CLIENT_ID");
+        String multiTenantId = "54826b22-38d6-4fb2-bad9-b7b93a3e9c5a";
+        String multiClientId = "4fc2b07b-9d91-4a4a-86e0-96d5b9145075";
         String multiClientSecret = configuration.get("AZURE_IDENTITY_MULTI_TENANT_CLIENT_SECRET");
 
         ClientSecretCredential credential = new ClientSecretCredentialBuilder().tenantId(multiTenantId)

sdk/identity/azure-identity/src/test/java/com/azure/identity/implementation/Retry.java

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.identity.implementation;
+
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * TODO: @g2vinay, move this to azure-core-test.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.METHOD)
+@ExtendWith(RetryExtension.class)
+public @interface Retry {
+    int maxRetries() default 3; // Default to 3 retries
+}

sdk/identity/azure-identity/src/test/java/com/azure/identity/implementation/RetryExtension.java

@@ -0,0 +1,84 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.identity.implementation;
+
+import org.junit.jupiter.api.extension.BeforeTestExecutionCallback;
+import org.junit.jupiter.api.extension.ExtensionContext;
+import org.junit.jupiter.api.extension.TestExecutionExceptionHandler;
+
+import java.lang.reflect.InvocationTargetException;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Supplier;
+
+/**
+ * RetryExtension: The extension to retry tests upon failure.
+ * TODO: @g2vinay, move this to azure-core-test.
+ */
+public class RetryExtension implements TestExecutionExceptionHandler, BeforeTestExecutionCallback {
+    private static final int DEFAULT_MAX_RETRIES = 3;
+
+    @Override
+    public void handleTestExecutionException(ExtensionContext context, Throwable throwable) {
+        int maxRetries = getMaxRetries(context);
+
+        executeWithRetries(() -> {
+            try {
+                return context.getTestMethod().get().invoke(context.getRequiredTestInstance());
+            } catch (IllegalAccessException | InvocationTargetException e) {
+                throw new RuntimeException(e);
+            }
+        }, maxRetries, 1000, 2);
+
+    }
+
+    @Override
+    public void beforeTestExecution(ExtensionContext context) {
+        // Initialize the retry counter before a test execution starts
+        getStore(context).put(getTestKey(context), new AtomicInteger(0));
+    }
+
+    private int getMaxRetries(ExtensionContext context) {
+        return context.getTestMethod()
+            .flatMap(method -> Optional.ofNullable(method.getAnnotation(Retry.class)).map(Retry::maxRetries))
+            .orElse(DEFAULT_MAX_RETRIES);
+    }
+
+    private ExtensionContext.Store getStore(ExtensionContext context) {
+        return context.getStore(ExtensionContext.Namespace.create(getClass(), context.getUniqueId()));
+    }
+
+    private String getTestKey(ExtensionContext context) {
+        return "retries-" + context.getUniqueId();
+    }
+
+    private static <T> T executeWithRetries(Supplier<T> logic, int maxRetries, long initialDelayMillis,
+        double backoffFactor) {
+        int attempt = 0;
+        long delay = initialDelayMillis;
+        Exception lastException = null;
+
+        while (attempt < maxRetries) {
+            try {
+                return logic.get();
+            } catch (Exception e) {
+                lastException = e;
+                attempt++;
+                if (attempt >= maxRetries) {
+                    break;
+                }
+                System.out.printf("Attempt %d failed. Retrying in %d ms...%n", attempt, delay);
+                try {
+                    TimeUnit.MILLISECONDS.sleep(delay);
+                } catch (InterruptedException ex) {
+                    throw new RuntimeException(ex);
+                }
+                delay *= backoffFactor;
+            }
+        }
+
+        throw new RuntimeException(lastException);
+    }
+}

sdk/identity/ci.yml

@@ -73,16 +73,30 @@ extends:
     LiveTestStages:
       - template: /eng/pipelines/templates/stages/archetype-sdk-tests-isolated.yml
         parameters:
+          # TODO: Remove/Refactor this to work in test code, as here it will only work in public cloud.
+          PreTestRunSteps:
+            - task: AzureCLI@2
+              displayName: Refresh OIDC token
+              env:
+                ARM_OIDC_TOKEN: $(ARM_OIDC_TOKEN)
+              inputs:
+                azureSubscription: azure-sdk-tests-public
+                addSpnToEnvironment: true
+                scriptLocation: inlineScript
+                scriptType: pscore
+                inlineScript: |
+                  Write-Host "##vso[task.setvariable variable=ARM_OIDC_TOKEN;issecret=true]$($env:idToken)"
+          PersistOidcToken: true
+          Location: 'westus2'
           CalledFromClient: true
           ServiceDirectory: identity
           SupportedClouds: 'Public,UsGov,China'
-          UseFederatedAuth: false
           EnvVars:
-            SYSTEM_ACCESSTOKEN: $(System.AccessToken)
+            AZ_OIDC_TOKEN: $(ARM_OIDC_TOKEN)
+            AZURE_IDENTITY_MULTI_TENANT_CLIENT_SECRET: $(AZURE-IDENTITY-MULTI-TENANT-CLIENT-SECRET)
           CloudConfig:
             Public:
               SubscriptionConfigurations:
-                - $(sub-config-azure-cloud-test-resources)
                 - $(sub-config-identity-test-resources)
           Artifacts:
             - name: azure-identity