Users/chshrikh/azure ps module load fix m130 (#6581)

* updated the logic to load the module as per auth scheme
microsoft · Mar 5, 2018 · 55959a0 · 55959a0
1 parent ffafabd
commit 55959a0
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 8 deletions.
diff --git a/Tasks/AzurePowerShell/AzurePowerShell.ps1 b/Tasks/AzurePowerShell/AzurePowerShell.ps1
@@ -9,6 +9,7 @@ $scriptArguments = Get-VstsInput -Name ScriptArguments
 $targetAzurePs = Get-VstsInput -Name TargetAzurePs
 $customTargetAzurePs = Get-VstsInput -Name CustomTargetAzurePs
 
+
 # Validate the script path and args do not contains new-lines. Otherwise, it will
 # break invoking the script via Invoke-Expression.
 if ($scriptType -eq "FilePath") {
@@ -45,7 +46,32 @@ if ($targetAzurePs -eq $latestVersion) {
 . "$PSScriptRoot\Utility.ps1"
 $targetAzurePs = Get-RollForwardVersion -azurePowerShellVersion $targetAzurePs
 
-Update-PSModulePathForHostedAgent -targetAzurePs $targetAzurePs
+$authScheme = ''
+try
+{
+    $serviceNameInput = Get-VstsInput -Name ConnectedServiceNameSelector -Default 'ConnectedServiceName'
+    $serviceName = Get-VstsInput -Name $serviceNameInput -Default (Get-VstsInput -Name DeploymentEnvironmentName)
+    if (!$serviceName)
+    {
+            Get-VstsInput -Name $serviceNameInput -Require
+    }
+
+    $endpoint = Get-VstsEndpoint -Name $serviceName -Require
+
+    if($endpoint)
+    {
+        $authScheme = $endpoint.Auth.Scheme 
+    }
+
+     Write-Verbose "AuthScheme $authScheme" 
+}
+catch
+{
+   $error = $_.Exception.Message
+   Write-Verbose "Unable to get the authScheme $error" 
+}
+
+Update-PSModulePathForHostedAgent -targetAzurePs $targetAzurePs -authScheme $authScheme
 
 try {
     # Initialize Azure.

diff --git a/Tasks/AzurePowerShell/Tests/PerformsBasicFlow.ps1 b/Tasks/AzurePowerShell/Tests/PerformsBasicFlow.ps1
@@ -10,6 +10,7 @@ Register-Mock Get-VstsInput { $targetAzurePs } -- -Name TargetAzurePs
 Register-Mock Get-VstsInput { 'arg1 arg2' } -- -Name ScriptArguments
 Register-Mock Update-PSModulePathForHostedAgent
 Register-Mock Initialize-Azure
+Register-Mock Get-VstsEndpoint { @{auth = @{ scheme = "ServicePrincipal" }} }
 
 # Act.
 $actual = & $PSScriptRoot\..\AzurePowerShell.ps1

diff --git a/Tasks/AzurePowerShell/Utility.ps1 b/Tasks/AzurePowerShell/Utility.ps1
@@ -17,7 +17,8 @@ function Get-SavedModulePath {
 
 function Update-PSModulePathForHostedAgent {
     [CmdletBinding()]
-    param([string] $targetAzurePs)
+    param([string] $targetAzurePs,
+          [string] $authScheme)
     Trace-VstsEnteringInvocation $MyInvocation
     try {
         if ($targetAzurePs) {
@@ -28,10 +29,22 @@ function Update-PSModulePathForHostedAgent {
             $hostedAgentAzureRmModulePath = Get-LatestModule -patternToMatch "^azurerm_[0-9]+\.[0-9]+\.[0-9]+$" -patternToExtract "[0-9]+\.[0-9]+\.[0-9]+$" -Classic:$false
             $hostedAgentAzureModulePath  =  Get-LatestModule -patternToMatch "^azure_[0-9]+\.[0-9]+\.[0-9]+$"   -patternToExtract "[0-9]+\.[0-9]+\.[0-9]+$" -Classic:$true
         }
-        $env:PSModulePath = $hostedAgentAzureRmModulePath + ";" + $env:PSModulePath
-        $env:PSModulePath = $env:PSModulePath.TrimStart(';')
-        $env:PSModulePath = $hostedAgentAzureModulePath + ";" + $env:PSModulePath
-        $env:PSModulePath = $env:PSModulePath.TrimStart(';')
+
+        if($authScheme -eq 'ServicePrincipal' -or $authScheme -eq '')
+        {
+            $env:PSModulePath = $hostedAgentAzureModulePath + ";" + $env:PSModulePath
+            $env:PSModulePath = $env:PSModulePath.TrimStart(';')
+            $env:PSModulePath = $hostedAgentAzureRmModulePath + ";" + $env:PSModulePath
+            $env:PSModulePath = $env:PSModulePath.TrimStart(';')
+        }
+        else
+        {
+            $env:PSModulePath = $hostedAgentAzureRmModulePath + ";" + $env:PSModulePath
+            $env:PSModulePath = $env:PSModulePath.TrimStart(';')
+            $env:PSModulePath = $hostedAgentAzureModulePath + ";" + $env:PSModulePath
+            $env:PSModulePath = $env:PSModulePath.TrimStart(';')
+        }
+
     } finally {
         Write-Verbose "The updated value of the PSModulePath is: $($env:PSModulePath)"
         Trace-VstsLeavingInvocation $MyInvocation

diff --git a/Tasks/AzurePowerShell/task.json b/Tasks/AzurePowerShell/task.json
@@ -17,7 +17,7 @@
     "version": {
         "Major": 2,
         "Minor": 0,
-        "Patch": 5
+        "Patch": 6
     },
     "demands": [
         "azureps"

diff --git a/Tasks/AzurePowerShell/task.loc.json b/Tasks/AzurePowerShell/task.loc.json
@@ -17,7 +17,7 @@
   "version": {
     "Major": 2,
     "Minor": 0,
-    "Patch": 5
+    "Patch": 6
   },
   "demands": [
     "azureps"