From fe7a06d41653ddedd30d5740023f8d4519999a6a Mon Sep 17 00:00:00 2001
From: Jeremy Meng <yumeng@microsoft.com>
Date: Wed, 7 Jul 2021 14:56:16 -1000
Subject: [PATCH] [ContainerRegistry] Add cloud configuration support

- Update test resource deployment
---
 sdk/containerregistry/container-registry/tests.yml | 2 ++
 sdk/containerregistry/test-resources-post.ps1      | 8 ++++++++
 sdk/containerregistry/test-resources.json          | 8 ++++++--
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/sdk/containerregistry/container-registry/tests.yml b/sdk/containerregistry/container-registry/tests.yml
index 58e9e20816ab..65d8e7ea2caf 100644
--- a/sdk/containerregistry/container-registry/tests.yml
+++ b/sdk/containerregistry/container-registry/tests.yml
@@ -7,6 +7,8 @@ stages:
       ServiceDirectory: containerregistry
       MatrixFilters:
         - DependencyVersion=^$
+        - Pool=.*mms-ubuntu-2004.*
+      SupportedClouds: 'Public,UsGov,China'
       EnvVars:
         AZURE_CLIENT_ID: $(aad-azure-sdk-test-client-id)
         AZURE_TENANT_ID: $(aad-azure-sdk-test-tenant-id)
diff --git a/sdk/containerregistry/test-resources-post.ps1 b/sdk/containerregistry/test-resources-post.ps1
index 5ab8bd4521b7..b9b7ddf679bd 100644
--- a/sdk/containerregistry/test-resources-post.ps1
+++ b/sdk/containerregistry/test-resources-post.ps1
@@ -5,6 +5,14 @@ param (
     [string] $TestApplicationSecret
 )
 
+Get-AzEnvironment | Write-Host
+
+if ( $DeploymentOutputs['CONTAINER_REGISTRY_ENDPOINT'].EndsWith('.azurecr.us') ) {
+    Connect-AzAccount –Environment AzureUSGovernment
+} elseif ( $DeploymentOutputs['CONTAINER_REGISTRY_ENDPOINT'].EndsWith('.azurecr.cn') ) {
+    Connect-AzAccount –Environment AzureChinaCloud
+}
+
 Import-AzContainerRegistryImage `
     -ResourceGroupName $DeploymentOutputs['CONTAINERREGISTRY_RESOURCE_GROUP'] `
     -RegistryName $DeploymentOutputs['CONTAINER_REGISTRY_NAME'] `
diff --git a/sdk/containerregistry/test-resources.json b/sdk/containerregistry/test-resources.json
index 233bf2ece220..a9872e91defc 100644
--- a/sdk/containerregistry/test-resources.json
+++ b/sdk/containerregistry/test-resources.json
@@ -15,13 +15,17 @@
       "metadata": {
         "description": "The location of the resource. By default, this is the same as the resource group."
       }
+    },
+    "containerRegistryEndpointSuffix": {
+      "defaultValue": ".azurecr.io",
+      "type": "string"
     }
   },
   "variables": {
     "apiVersion": "2020-11-01-preview",
-    "endpointValue": "[format('https://{0}.azurecr.io', parameters('baseName'))]",
+    "endpointValue": "[format('https://{0}{1}', parameters('baseName'), parameters('containerRegistryEndpointSuffix'))]",
     "anonRegistryName": "[format('{0}anon', parameters('baseName'))]",
-    "anonEndpointValue": "[format('https://{0}.azurecr.io', variables('anonRegistryName'))]"
+    "anonEndpointValue": "[format('https://{0}{1}', variables('anonRegistryName'), parameters('containerRegistryEndpointSuffix'))]"
   },
   "resources": [
     {