[CmdletBinding(SupportsShouldProcess)]
param(
	[Parameter(Mandatory)]
	[string]$CloudEnvironment,
    [Parameter(Mandatory)]
    [string]$SubscriptionId
)

Connect-AzAccount -Identity -Environment $CloudEnvironment | Out-Null

#### Add this line
Set-AzContext -SubscriptionId $SubscriptionId

$AVDHostPools = Get-AzWvdHostPool -SubscriptionId $SubscriptionId

# Foreach storage account
Foreach ($storageAcct in $storageAccountResourceIDs) {


    ##### Context handling - add this block
    $Subscription = (Get-azSubscription -SubscriptionId ($StorageAccountResourceIDs -split '/')[2])
    $SubName = $Subscription.Name
    
    #Get current context
    $Context = Get-AzContext
    if ($Context.Subscription.Id -ne $Subscription.Id) {
        Set-AzContext -SubscriptionId $Subscription.Id
        Write-Output "Changed context to subscription $($Subscription.Name) $($Subscription.Id)"
    } 

    ##### Context handling

    $resourceGroup = ($storageAcct -split '/')[4]
    $storageAcctName = ($storageAcct -split '/')[8]
    #Write-Host "Working on Storage:" $storageAcctName "in" $resourceGroup