Use common OIDC token env vars in post scripts (#23549)

* Use common OIDC token env vars in post scripts * Pass cloud into azcontainerregistry test config
Azure · Oct 10, 2024 · 5df372c · 5df372c
1 parent 1c2a108
commit 5df372c
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 25 deletions.
diff --git a/eng/pipelines/templates/jobs/archetype-sdk-client.yml b/eng/pipelines/templates/jobs/archetype-sdk-client.yml
@@ -99,6 +99,9 @@ parameters:
   - name: UseFederatedAuth
     type: boolean
     default: true
+  - name: PersistOidcToken
+    type: boolean
+    default: false
 
 extends:
   template: /eng/pipelines/templates/stages/1es-redirect.yml
@@ -194,6 +197,7 @@ extends:
                                 PostSteps:
                                   - ${{ parameters.PostSteps }}
                                 UseFederatedAuth: ${{ parameters.UseFederatedAuth }}
+                                PersistOidcToken: ${{ parameters.PersistOidcToken }}
                               MatrixConfigs:
                                 # Enumerate platforms and additional platforms based on supported clouds (sparse platform<-->cloud matrix).
                                 - ${{ each config in parameters.MatrixConfigs }}:

diff --git a/eng/pipelines/templates/jobs/live.tests.yml b/eng/pipelines/templates/jobs/live.tests.yml
@@ -36,6 +36,9 @@ parameters:
   - name: UseFederatedAuth
     type: boolean
     default: true
+  - name: PersistOidcToken
+    type: boolean
+    default: false
 
 jobs:
   - job:
@@ -92,6 +95,7 @@ jobs:
           UseFederatedAuth: ${{ parameters.UseFederatedAuth }}
           ServiceConnection: ${{ parameters.CloudConfig.ServiceConnection }}
           SubscriptionConfigurationFilePaths: ${{ parameters.CloudConfig.SubscriptionConfigurationFilePaths}}
+          PersistOidcToken: ${{ parameters.PersistOidcToken }}
           EnvVars:
             Pool: $(Pool)
             ${{ insert }}: ${{ parameters.EnvVars }}

diff --git a/sdk/azidentity/ci.yml b/sdk/azidentity/ci.yml
@@ -38,15 +38,7 @@ extends:
       UsePipelineProxy: false
 
       ${{ if endsWith(variables['Build.DefinitionName'], 'weekly') }}:
-        PreSteps:
-        - task: AzureCLI@2
-          displayName: Set OIDC token
-          inputs:
-            addSpnToEnvironment: true
-            azureSubscription: azure-sdk-tests
-            inlineScript: Write-Host "##vso[task.setvariable variable=OIDC_TOKEN;]$($env:idToken)"
-            scriptLocation: inlineScript
-            scriptType: pscore
+        PersistOidcToken: true
         MatrixConfigs:
           - Name: managed_identity_matrix
             GenerateVMJobs: true

diff --git a/sdk/azidentity/test-resources-post.ps1 b/sdk/azidentity/test-resources-post.ps1
@@ -7,6 +7,10 @@ param (
   [hashtable] $AdditionalParameters = @{},
   [hashtable] $DeploymentOutputs,
 
+  [Parameter(Mandatory = $true)]
+  [ValidateNotNullOrEmpty()]
+  [string] $SubscriptionId,
+
   [Parameter(ParameterSetName = 'Provisioner', Mandatory = $true)]
   [ValidateNotNullOrEmpty()]
   [string] $TenantId,
@@ -15,6 +19,10 @@ param (
   [ValidatePattern('^[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}$')]
   [string] $TestApplicationId,
 
+  [Parameter(Mandatory = $true)]
+  [ValidateNotNullOrEmpty()]
+  [string] $Environment,
+
   # Captures any arguments from eng/New-TestResources.ps1 not declared here (no parameter errors).
   [Parameter(ValueFromRemainingArguments = $true)]
   $RemainingArguments
@@ -28,8 +36,9 @@ if ($CI) {
     Write-Host "Skipping post-provisioning script because resources weren't deployed"
     return
   }
-  az login --federated-token $env:OIDC_TOKEN --service-principal -t $TenantId -u $TestApplicationId
-  az account set --subscription $DeploymentOutputs['AZIDENTITY_SUBSCRIPTION_ID']
+  az cloud set -n $Environment
+  az login --federated-token $env:ARM_OIDC_TOKEN --service-principal -t $TenantId -u $TestApplicationId
+  az account set --subscription $SubscriptionId
 }
 
 Write-Host "Building container"

diff --git a/sdk/containers/azcontainerregistry/ci.yml b/sdk/containers/azcontainerregistry/ci.yml
@@ -24,18 +24,10 @@ pr:
 extends:
   template: /eng/pipelines/templates/jobs/archetype-sdk-client.yml
   parameters:
-    PreSteps:
-    - task: AzureCLI@2
-      displayName: Set OIDC token
-      inputs:
-        addSpnToEnvironment: true
-        azureSubscription: azure-sdk-tests
-        inlineScript: Write-Host "##vso[task.setvariable variable=OIDC_TOKEN;]$($env:idToken)"
-        scriptType: pscore
-        scriptLocation: inlineScript
     ServiceDirectory: 'containers/azcontainerregistry'
     RunLiveTests: true
     UseFederatedAuth: true
     UsePipelineProxy: false
     TestRunTime: '30m'
     SupportedClouds: 'Public,UsGov'
+    PersistOidcToken: true
diff --git a/sdk/containers/azcontainerregistry/test-resources-post.ps1 b/sdk/containers/azcontainerregistry/test-resources-post.ps1
@@ -4,22 +4,32 @@
 # IMPORTANT: Do not invoke this file directly. Please instead run eng/common/TestResources/New-TestResources.ps1 from the repository root.
 
 param (
-    [hashtable] $DeploymentOutputs,
+    [Parameter(Mandatory = $true)]
+    [ValidateNotNullOrEmpty()]
+    [string] $SubscriptionId,
 
-    [Parameter(ParameterSetName = 'Provisioner', Mandatory = $true)]
+    [Parameter(Mandatory = $true)]
     [ValidateNotNullOrEmpty()]
     [string] $TenantId,
 
-    [Parameter()]
+    [Parameter(Mandatory = $true)]
     [ValidatePattern('^[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}$')]
     [string] $TestApplicationId,
 
+    [Parameter(Mandatory = $true)]
+    [ValidateNotNullOrEmpty()]
+    [string] $Environment,
+
     # Captures any arguments from eng/New-TestResources.ps1 not declared here (no parameter errors).
     [Parameter(ValueFromRemainingArguments = $true)]
     $RemainingArguments
 )
 
+$ErrorActionPreference = 'Stop'
+$PSNativeCommandUseErrorActionPreference = $true
+
 if ($CI) {
-    az login --federated-token $env:OIDC_TOKEN --service-principal -t $TenantId -u $TestApplicationId
-    az account set --subscription $DeploymentOutputs['AZCONTAINERREGISTRY_SUBSCRIPTION_ID']
+    az cloud set -n $Environment
+    az login --federated-token $env:ARM_OIDC_TOKEN --service-principal -t $TenantId -u $TestApplicationId
+    az account set --subscription $SubscriptionId
 }
diff --git a/sdk/containers/azcontainerregistry/utils_test.go b/sdk/containers/azcontainerregistry/utils_test.go
@@ -53,6 +53,7 @@ func getEndpointCredAndClientOptions(t *testing.T) (string, azcore.TokenCredenti
 	transport, err := recording.NewRecordingHTTPClient(t, nil)
 	require.NoError(t, err)
 	options := azcore.ClientOptions{
+		Cloud:     testConfig.cloud,
 		Transport: transport,
 	}
 	return "https://" + testConfig.loginServer, testConfig.credential, options