From 6680d8ca9447daff5f8790e0a713193899717338 Mon Sep 17 00:00:00 2001
From: Ben Broderick Phillips <bebroder@microsoft.com>
Date: Thu, 12 Dec 2024 15:48:06 -0500
Subject: [PATCH 1/2] Use -Mail for test resources local user lookup to support
 TME

---
 eng/common/TestResources/New-TestResources.ps1 | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/eng/common/TestResources/New-TestResources.ps1 b/eng/common/TestResources/New-TestResources.ps1
index 2958ccc6628f..20d880071452 100644
--- a/eng/common/TestResources/New-TestResources.ps1
+++ b/eng/common/TestResources/New-TestResources.ps1
@@ -313,7 +313,8 @@ 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') {
-            $user = Get-AzADUser -UserPrincipalName $context.Account.Id
+            # Use -Mail as the lookup works in both corp and TME tenants
+            $user = Get-AzADUser -Mail $context.Account
             $ProvisionerApplicationOid = $user.Id
         } elseif ($context.Account.Type -eq 'ServicePrincipal') {
             $sp = Get-AzADServicePrincipal -ApplicationId $context.Account.Id
@@ -383,7 +384,8 @@ try {
             Write-Warning "The specified TestApplicationId '$TestApplicationId' will be ignored when -ServicePrincipalAutth is not set."
         }
 
-        $userAccount = (Get-AzADUser -UserPrincipalName (Get-AzContext).Account)
+        # Use -Mail as the lookup works in both corp and TME tenants
+        $userAccount = (Get-AzADUser -Mail (Get-AzContext).Account)
         $TestApplicationOid = $userAccount.Id
         $TestApplicationId = $testApplicationOid
         $userAccountName = $userAccount.UserPrincipalName

From 655d393e8df570a561caf10cb1e7bb5b743b8228 Mon Sep 17 00:00:00 2001
From: Ben Broderick Phillips <bebroder@microsoft.com>
Date: Thu, 12 Dec 2024 17:47:02 -0500
Subject: [PATCH 2/2] Auto-select TME subscription based on user context

---
 eng/common/TestResources/New-TestResources.ps1 | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/eng/common/TestResources/New-TestResources.ps1 b/eng/common/TestResources/New-TestResources.ps1
index 20d880071452..8a61b68c6b2e 100644
--- a/eng/common/TestResources/New-TestResources.ps1
+++ b/eng/common/TestResources/New-TestResources.ps1
@@ -244,10 +244,16 @@ try {
                 $context = Get-AzContext
             }
         } else {
-            if ($currentSubcriptionId -ne 'faa080af-c1d8-40ad-9cce-e1a450ca5b57') {
+            if ($context.Tenant.Name -like '*TME*') {
+                if ($currentSubscriptionId -ne '4d042dc6-fe17-4698-a23f-ec6a8d1e98f4') {
+                    Log "Attempting to select subscription 'Azure SDK Test Resources - TME (4d042dc6-fe17-4698-a23f-ec6a8d1e98f4)'"
+                    $null = Select-AzSubscription -Subscription '4d042dc6-fe17-4698-a23f-ec6a8d1e98f4' -ErrorAction Ignore
+                    # Update the context.
+                    $context = Get-AzContext
+                }
+            } elseif ($currentSubcriptionId -ne 'faa080af-c1d8-40ad-9cce-e1a450ca5b57') {
                 Log "Attempting to select subscription 'Azure SDK Developer Playground (faa080af-c1d8-40ad-9cce-e1a450ca5b57)'"
                 $null = Select-AzSubscription -Subscription 'faa080af-c1d8-40ad-9cce-e1a450ca5b57' -ErrorAction Ignore
-
                 # Update the context.
                 $context = Get-AzContext
             }
@@ -261,6 +267,7 @@ try {
             'faa080af-c1d8-40ad-9cce-e1a450ca5b57' = 'Azure SDK Developer Playground'
             'a18897a6-7e44-457d-9260-f2854c0aca42' = 'Azure SDK Engineering System'
             '2cd617ea-1866-46b1-90e3-fffb087ebf9b' = 'Azure SDK Test Resources'
+            '4d042dc6-fe17-4698-a23f-ec6a8d1e98f4' = 'Azure SDK Test Resources - TME '
         }
 
         # Print which subscription is currently selected.
@@ -314,7 +321,7 @@ try {
     if (!$ProvisionerApplicationId -and !$ProvisionerApplicationOid) {
         if ($context.Account.Type -eq 'User') {
             # Use -Mail as the lookup works in both corp and TME tenants
-            $user = Get-AzADUser -Mail $context.Account
+            $user = Get-AzADUser -Mail $context.Account.Id
             $ProvisionerApplicationOid = $user.Id
         } elseif ($context.Account.Type -eq 'ServicePrincipal') {
             $sp = Get-AzADServicePrincipal -ApplicationId $context.Account.Id
@@ -385,7 +392,7 @@ try {
         }
 
         # Use -Mail as the lookup works in both corp and TME tenants
-        $userAccount = (Get-AzADUser -Mail (Get-AzContext).Account)
+        $userAccount = (Get-AzADUser -Mail (Get-AzContext).Account.Id)
         $TestApplicationOid = $userAccount.Id
         $TestApplicationId = $testApplicationOid
         $userAccountName = $userAccount.UserPrincipalName