From 8cafe21c8bfb72560a25aec146bd618d8db4c6b3 Mon Sep 17 00:00:00 2001
From: Azure SDK Bot <53356347+azure-sdk@users.noreply.github.com>
Date: Fri, 13 Aug 2021 15:22:39 -0700
Subject: [PATCH] Add Rest Method checks to Prepare-Release (#15290)

Co-authored-by: Wes Haggard <Wes.Haggard@microsoft.com>
---
 .../Helpers/DevOps-WorkItem-Helpers.ps1       | 62 +++++++++++--------
 .../Update-DevOps-Release-WorkItem.ps1        |  2 +
 2 files changed, 37 insertions(+), 27 deletions(-)

diff --git a/eng/common/scripts/Helpers/DevOps-WorkItem-Helpers.ps1 b/eng/common/scripts/Helpers/DevOps-WorkItem-Helpers.ps1
index aca34080da46..24420cef2b66 100644
--- a/eng/common/scripts/Helpers/DevOps-WorkItem-Helpers.ps1
+++ b/eng/common/scripts/Helpers/DevOps-WorkItem-Helpers.ps1
@@ -3,6 +3,38 @@ $ReleaseDevOpsOrgParameters =  @("--organization", "https://dev.azure.com/azure-
 $ReleaseDevOpsCommonParameters =  $ReleaseDevOpsOrgParameters + @("--output", "json")
 $ReleaseDevOpsCommonParametersWithProject = $ReleaseDevOpsCommonParameters + @("--project", "Release")
 
+function Get-DevOpsRestHeaders()
+{
+  $headers = $null
+  if (Get-Variable -Name "devops_pat" -ValueOnly -ErrorAction "Ignore")
+  {
+    $encodedToken = [System.Convert]::ToBase64String([System.Text.Encoding]::ASCII.GetBytes([string]::Format("{0}:{1}", "", $devops_pat)))
+    $headers = @{ Authorization = "Basic $encodedToken" }
+  }
+  else
+  {
+    # Get a temp access token from the logged in az cli user for azure devops resource
+    $jwt_accessToken = (az account get-access-token --resource "499b84ac-1321-427f-aa17-267ca6975798" --query "accessToken" --output tsv)
+    $headers = @{ Authorization = "Bearer $jwt_accessToken" }
+  }
+
+  return $headers
+}
+
+function CheckDevOpsAccess()
+{
+  # Dummy test query to validate permissions
+  $query = "SELECT [System.ID] FROM WorkItems WHERE [Work Item Type] = 'Package' AND [Package] = 'azure-sdk-template'"
+
+  $response = Invoke-RestMethod -Method POST `
+    -Uri "https://dev.azure.com/azure-sdk/Release/_apis/wit/wiql/?api-version=6.0" `
+    -Headers (Get-DevOpsRestHeaders) -Body "{ ""query"": ""$query"" }" -ContentType "application/json" | ConvertTo-Json -Depth 10 | ConvertFrom-Json -AsHashTable
+
+  if ($response -isnot [HashTable] -or !$response.ContainsKey("workItems")) {
+    throw "Failed to run test query against Azure DevOps. Please ensure you are logged into the public azure cloud. Consider running 'az logout' and then 'az login'."
+  }
+}
+
 function Invoke-AzBoardsCmd($subCmd, $parameters, $output = $true)
 {
   $azCmdStr = "az boards ${subCmd} $($parameters -join ' ')"
@@ -26,23 +58,11 @@ function Invoke-Query($fields, $wiql, $output = $true)
     Write-Host "Executing query $wiql"
   }
 
-  $headers = $null
-  if (Get-Variable -Name "devops_pat" -ValueOnly -ErrorAction "Ignore")
-  {
-    $encodedToken = [System.Convert]::ToBase64String([System.Text.Encoding]::ASCII.GetBytes([string]::Format("{0}:{1}", "", $devops_pat)))
-    $headers = @{ Authorization = "Basic $encodedToken" }
-  }
-  else
-  {
-    # Get a temp access token from the logged in az cli user for azure devops resource
-    $jwt_accessToken = (az account get-access-token --resource "499b84ac-1321-427f-aa17-267ca6975798" --query "accessToken" --output tsv)
-    $headers = @{ Authorization = "Bearer $jwt_accessToken" }
-  }
   $response = Invoke-RestMethod -Method POST `
     -Uri "https://dev.azure.com/azure-sdk/Release/_apis/wit/wiql/?`$top=10000&api-version=6.0" `
-    -Headers $headers -Body $body -ContentType "application/json" | ConvertTo-Json -Depth 10 | ConvertFrom-Json -AsHashTable
+    -Headers (Get-DevOpsRestHeaders) -Body $body -ContentType "application/json" | ConvertTo-Json -Depth 10 | ConvertFrom-Json -AsHashTable
 
-  if (!$response.workItems) {
+  if ($response -isnot [HashTable] -or !$response.ContainsKey("workItems")) {
     Write-Verbose "Query returned no items. $wiql"
     return ,@()
   }
@@ -891,20 +911,8 @@ function UpdatePackageVersions($pkgWorkItem, $plannedVersions, $shippedVersions)
 
   $body = "[" + ($fieldUpdates -join ',') + "]"
 
-  $headers = $null
-  if (Get-Variable -Name "devops_pat" -ValueOnly -ErrorAction "Ignore")
-  {
-    $encodedToken = [System.Convert]::ToBase64String([System.Text.Encoding]::ASCII.GetBytes([string]::Format("{0}:{1}", "", $devops_pat)))
-    $headers = @{ Authorization = "Basic $encodedToken" }
-  }
-  else
-  {
-    # Get a temp access token from the logged in az cli user for azure devops resource
-    $jwt_accessToken = (az account get-access-token --resource "499b84ac-1321-427f-aa17-267ca6975798" --query "accessToken" --output tsv)
-    $headers = @{ Authorization = "Bearer $jwt_accessToken" }
-  }
   $response = Invoke-RestMethod -Method PATCH `
     -Uri "https://dev.azure.com/azure-sdk/_apis/wit/workitems/${id}?api-version=6.0" `
-    -Headers $headers -Body $body -ContentType "application/json-patch+json" | ConvertTo-Json -Depth 10 | ConvertFrom-Json -AsHashTable
+    -Headers (Get-DevOpsRestHeaders) -Body $body -ContentType "application/json-patch+json" | ConvertTo-Json -Depth 10 | ConvertFrom-Json -AsHashTable
   return $response
 }
\ No newline at end of file
diff --git a/eng/common/scripts/Update-DevOps-Release-WorkItem.ps1 b/eng/common/scripts/Update-DevOps-Release-WorkItem.ps1
index b1a45abffa2e..34c754006276 100644
--- a/eng/common/scripts/Update-DevOps-Release-WorkItem.ps1
+++ b/eng/common/scripts/Update-DevOps-Release-WorkItem.ps1
@@ -38,6 +38,8 @@ if (!$?){
 . (Join-Path $PSScriptRoot SemVer.ps1)
 . (Join-Path $PSScriptRoot Helpers DevOps-WorkItem-Helpers.ps1)
 
+CheckDevOpsAccess
+
 $parsedNewVersion = [AzureEngSemanticVersion]::new($version)
 $state = "In Release"
 $releaseType = $parsedNewVersion.VersionType