From a9083d489b34a364472fe809d472f015ff76ade6 Mon Sep 17 00:00:00 2001
From: Azure SDK Bot <53356347+azure-sdk@users.noreply.github.com>
Date: Tue, 28 Jan 2025 11:44:31 -0800
Subject: [PATCH] Add fallback user id lookup for test resource deployment
 (#6375)

Co-authored-by: Ben Broderick Phillips <bebroder@microsoft.com>
---
 eng/common/TestResources/New-TestResources.ps1 | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/eng/common/TestResources/New-TestResources.ps1 b/eng/common/TestResources/New-TestResources.ps1
index 8a61b68c6..2c1308505 100644
--- a/eng/common/TestResources/New-TestResources.ps1
+++ b/eng/common/TestResources/New-TestResources.ps1
@@ -320,8 +320,14 @@ try {
     # Make sure the provisioner OID is set so we can pass it through to the deployment.
     if (!$ProvisionerApplicationId -and !$ProvisionerApplicationOid) {
         if ($context.Account.Type -eq 'User') {
-            # Use -Mail as the lookup works in both corp and TME tenants
+            # Support corp tenant and TME tenant user id lookups
             $user = Get-AzADUser -Mail $context.Account.Id
+            if ($user -eq $null -or !$user.Id) {
+                $user = Get-AzADUser -UserPrincipalName $context.Account.Id
+            }
+            if ($user -eq $null -or !$user.Id) {
+                throw "Failed to find entra object ID for the current user"
+            }
             $ProvisionerApplicationOid = $user.Id
         } elseif ($context.Account.Type -eq 'ServicePrincipal') {
             $sp = Get-AzADServicePrincipal -ApplicationId $context.Account.Id
@@ -391,8 +397,14 @@ try {
             Write-Warning "The specified TestApplicationId '$TestApplicationId' will be ignored when -ServicePrincipalAutth is not set."
         }
 
-        # Use -Mail as the lookup works in both corp and TME tenants
+        # Support corp tenant and TME tenant user id lookups
         $userAccount = (Get-AzADUser -Mail (Get-AzContext).Account.Id)
+        if ($userAccount -eq $null -or !$userAccount.Id) {
+            $userAccount = (Get-AzADUser -UserPrincipalName (Get-AzContext).Account)
+        }
+        if ($userAccount -eq $null -or !$userAccount.Id) {
+            throw "Failed to find entra object ID for the current user"
+        }
         $TestApplicationOid = $userAccount.Id
         $TestApplicationId = $testApplicationOid
         $userAccountName = $userAccount.UserPrincipalName