[Monitor] Enable sovereign cloud testing (#35209)

This refactors some of the monitor ingestion/query library tests to allow for sovereign cloud testing. China cloud and Government cloud testing are also enabled on the weekly pipelines. Signed-off-by: Paul Van Eck <paulvaneck@microsoft.com>
Azure · May 3, 2024 · 62d3e2d · 62d3e2d
1 parent b64cfdd
commit 62d3e2d
Show file tree

Hide file tree

Showing 6 changed files with 95 additions and 27 deletions.
diff --git a/sdk/monitor/azure-monitor-ingestion/tests/base_testcase.py b/sdk/monitor/azure-monitor-ingestion/tests/base_testcase.py
@@ -0,0 +1,28 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See LICENSE.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+import os
+
+from devtools_testutils import AzureRecordedTestCase
+
+
+ENV_MONITOR_ENVIRONMENT = "MONITOR_ENVIRONMENT"
+AUDIENCE_MAP = {
+    "AzureCloud": "https://monitor.azure.com",
+    "AzureChinaCloud": "https://monitor.azure.cn",
+    "AzureUSGovernment": "https://monitor.azure.us"
+}
+
+
+class LogsIngestionClientTestCase(AzureRecordedTestCase):
+
+    def get_client(self, client_class, credential, **kwargs):
+
+        environment = os.getenv(ENV_MONITOR_ENVIRONMENT)
+        if environment:
+            audience = AUDIENCE_MAP.get(environment)
+            kwargs["credential_scopes"] = [f"{audience}/.default"]
+
+        return self.create_client_from_credential(client_class, credential, **kwargs)
diff --git a/sdk/monitor/azure-monitor-ingestion/tests/test_logs_ingestion.py b/sdk/monitor/azure-monitor-ingestion/tests/test_logs_ingestion.py
@@ -13,7 +13,8 @@
 
 from azure.core.exceptions import HttpResponseError
 from azure.monitor.ingestion import LogsIngestionClient
-from devtools_testutils import AzureRecordedTestCase
+
+from base_testcase import LogsIngestionClientTestCase
 
 
 LOGS_BODY = [
@@ -36,30 +37,30 @@
 ]
 
 
-class TestLogsIngestionClient(AzureRecordedTestCase):
+class TestLogsIngestionClient(LogsIngestionClientTestCase):
 
     def test_send_logs(self, recorded_test, monitor_info):
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
         with client:
             client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=LOGS_BODY)
 
     def test_send_logs_large(self, recorded_test, monitor_info, large_data):
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
         with client:
             client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=large_data)
 
     def test_send_logs_error(self, recorded_test, monitor_info):
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
         body = [{"foo": "bar"}]
 
         with pytest.raises(HttpResponseError) as ex:
             client.upload(rule_id='bad-rule', stream_name=monitor_info['stream_name'], logs=body)
 
     def test_send_logs_error_custom(self, recorded_test, monitor_info):
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
         body = [{"foo": "bar"}]
 
@@ -75,7 +76,7 @@ def on_error(e):
         assert on_error.called
 
     def test_send_logs_json_file(self, recorded_test, monitor_info):
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
 
         temp_file = str(uuid.uuid4()) + '.json'
@@ -89,7 +90,7 @@ def test_send_logs_json_file(self, recorded_test, monitor_info):
 
     @pytest.mark.live_test_only("Issues recording binary streams with test-proxy")
     def test_send_logs_gzip_file(self, monitor_info):
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
 
         temp_file = str(uuid.uuid4()) + '.json.gz'
@@ -101,7 +102,7 @@ def test_send_logs_gzip_file(self, monitor_info):
         os.remove(temp_file)
 
     def test_abort_error_handler(self, monitor_info):
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
 
         class TestException(Exception):
@@ -142,7 +143,7 @@ def on_error(e):
 
     @pytest.mark.parametrize("logs", ['[{"foo": "bar"}]', "foo", {"foo": "bar"}, None])
     def test_invalid_logs_format(self, monitor_info, logs):
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, self.get_credential(LogsIngestionClient), endpoint=monitor_info['dce'])
 
         with pytest.raises(ValueError):

diff --git a/sdk/monitor/azure-monitor-ingestion/tests/test_logs_ingestion_async.py b/sdk/monitor/azure-monitor-ingestion/tests/test_logs_ingestion_async.py
@@ -13,7 +13,8 @@
 
 from azure.core.exceptions import HttpResponseError
 from azure.monitor.ingestion.aio import LogsIngestionClient
-from devtools_testutils import AzureRecordedTestCase
+
+from base_testcase import LogsIngestionClientTestCase
 
 
 LOGS_BODY = [
@@ -36,12 +37,12 @@
 ]
 
 
-class TestLogsIngestionClientAsync(AzureRecordedTestCase):
+class TestLogsIngestionClientAsync(LogsIngestionClientTestCase):
 
     @pytest.mark.asyncio
     async def test_send_logs_async(self, recorded_test, monitor_info):
         credential = self.get_credential(LogsIngestionClient, is_async=True)
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, credential, endpoint=monitor_info['dce'])
         async with client:
             await client.upload(rule_id=monitor_info['dcr_id'], stream_name=monitor_info['stream_name'], logs=LOGS_BODY)
@@ -50,7 +51,7 @@ async def test_send_logs_async(self, recorded_test, monitor_info):
     @pytest.mark.asyncio
     async def test_send_logs_large(self, recorded_test, monitor_info, large_data):
         credential = self.get_credential(LogsIngestionClient, is_async=True)
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, credential, endpoint=monitor_info['dce'])
         async with client:
             await client.upload(
@@ -60,7 +61,7 @@ async def test_send_logs_large(self, recorded_test, monitor_info, large_data):
     @pytest.mark.asyncio
     async def test_send_logs_error(self, recorded_test, monitor_info):
         credential = self.get_credential(LogsIngestionClient, is_async=True)
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, credential, endpoint=monitor_info['dce'])
         body = [{"foo": "bar"}]
 
@@ -72,7 +73,7 @@ async def test_send_logs_error(self, recorded_test, monitor_info):
     @pytest.mark.asyncio
     async def test_send_logs_error_custom(self, recorded_test, monitor_info):
         credential = self.get_credential(LogsIngestionClient, is_async=True)
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, credential, endpoint=monitor_info['dce'])
         body = [{"foo": "bar"}]
 
@@ -92,7 +93,7 @@ async def on_error(e):
     @pytest.mark.asyncio
     async def test_send_logs_json_file(self, recorded_test, monitor_info):
         credential = self.get_credential(LogsIngestionClient, is_async=True)
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, credential, endpoint=monitor_info['dce'])
 
         temp_file = str(uuid.uuid4()) + '.json'
@@ -109,7 +110,7 @@ async def test_send_logs_json_file(self, recorded_test, monitor_info):
     @pytest.mark.live_test_only("Issues recording binary streams with test-proxy")
     async def test_send_logs_gzip_file(self, monitor_info):
         credential = self.get_credential(LogsIngestionClient, is_async=True)
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, credential, endpoint=monitor_info['dce'])
 
         temp_file = str(uuid.uuid4()) + '.json.gz'
@@ -125,7 +126,7 @@ async def test_send_logs_gzip_file(self, monitor_info):
     @pytest.mark.asyncio
     async def test_abort_error_handler(self, monitor_info):
         credential = self.get_credential(LogsIngestionClient, is_async=True)
-        client = self.create_client_from_credential(
+        client = self.get_client(
             LogsIngestionClient, credential, endpoint=monitor_info['dce'])
 
         class TestException(Exception):
@@ -169,7 +170,7 @@ async def on_error(e):
     @pytest.mark.parametrize("logs", ['[{"foo": "bar"}]', "foo", {"foo": "bar"}, None])
     async def test_invalid_logs_format(self, monitor_info, logs):
         credential = self.get_credential(LogsIngestionClient, is_async=True)
-        client = self.create_client_from_credential(LogsIngestionClient, credential, endpoint=monitor_info['dce'])
+        client = self.get_client(LogsIngestionClient, credential, endpoint=monitor_info['dce'])
 
         async with client:
             with pytest.raises(ValueError):

diff --git a/sdk/monitor/azure-monitor-query/tests/base_testcase.py b/sdk/monitor/azure-monitor-query/tests/base_testcase.py
@@ -18,6 +18,18 @@
     "AzureUSGovernment": "https://api.loganalytics.us/v1"
 }
 
+METRICS_CLIENT_ENVIRONMENT_AUDIENCE_MAP = {
+    "AzureCloud": "https://metrics.monitor.azure.com",
+    "AzureChinaCloud": "https://metrics.monitor.azure.cn",
+    "AzureUSGovernment": "https:/metrics.monitor.azure.us"
+}
+
+TLD_MAP = {
+    "AzureCloud": "com",
+    "AzureChinaCloud": "cn",
+    "AzureUSGovernment": "us"
+}
+
 
 class AzureMonitorQueryLogsTestCase(AzureRecordedTestCase):
 
@@ -47,9 +59,15 @@ class MetricsClientTestCase(AzureRecordedTestCase):
 
     def get_client(self, client_class, credential, endpoint = None):
 
+        environment = os.getenv(ENV_MONITOR_ENVIRONMENT)
         kwargs = {}
+        tld = "com"
+        if environment:
+            kwargs["audience"] = METRICS_CLIENT_ENVIRONMENT_AUDIENCE_MAP.get(environment)
+            tld = TLD_MAP.get(environment, "com")
+
         if not endpoint:
             region = os.getenv(ENV_MONITOR_LOCATION) or "westus2"
-            kwargs["endpoint"] = f"https://{region}.metrics.monitor.azure.com"
+            kwargs["endpoint"] = f"https://{region}.metrics.monitor.azure.{tld}"
 
         return self.create_client_from_credential(client_class, credential, **kwargs)
diff --git a/sdk/monitor/test-resources-post.ps1 b/sdk/monitor/test-resources-post.ps1
@@ -14,6 +14,8 @@ $clientSecret = $DeploymentOutputs['MONITOR_CLIENT_SECRET']
 $dcrImmutableId = $DeploymentOutputs['AZURE_MONITOR_DCR_ID']
 $dceEndpoint = $DeploymentOutputs['AZURE_MONITOR_DCE']
 $streamName = $DeploymentOutputs['AZURE_MONITOR_STREAM_NAME']
+$environment = $DeploymentOutputs['MONITOR_ENVIRONMENT']
+$authorityHost = $DeploymentOutputs['AZURE_AUTHORITY_HOST']
 
 ##################
 ### Step 0: Wait for role assignment to propagate
@@ -24,10 +26,19 @@ Start-Sleep -s 180
 ##################
 ### Step 1: Obtain a bearer token used later to authenticate against the DCE.
 ##################
-$scope= [System.Web.HttpUtility]::UrlEncode("https://monitor.azure.com//.default")
+# Audience Mappings
+$audienceMappings = @{
+    "AzureCloud" = "https://monitor.azure.com";
+    "AzureUSGovernment" = "https://monitor.azure.us";
+    "AzureChinaCloud" = "https://monitor.azure.cn";
+}
+
+$audience = $audienceMappings[$environment]
+
+$scope= [System.Web.HttpUtility]::UrlEncode("$audience/.default")
 $body = "client_id=$clientId&scope=$scope&client_secret=$clientSecret&grant_type=client_credentials";
 $headers = @{"Content-Type"="application/x-www-form-urlencoded"};
-$uri = "https://login.microsoftonline.com/$tenantId/oauth2/v2.0/token"
+$uri = "$authorityHost/$tenantId/oauth2/v2.0/token"
 $bearerToken = (Invoke-RestMethod -Uri $uri -Method "Post" -Body $body -Headers $headers).access_token
 
 ##################
@@ -68,12 +79,12 @@ $staticData = @"
 ##################
 $body = $staticData;
 $headers = @{"Authorization"="Bearer $bearerToken";"Content-Type"="application/json"};
-$uri = "$dceEndpoint/dataCollectionRules/$dcrImmutableId/streams/${streamName}?api-version=2021-11-01-preview"
-$uri2 = "$dceEndpoint/dataCollectionRules/$dcrImmutableId/streams/${streamName}2?api-version=2021-11-01-preview"
+$uri = "$dceEndpoint/dataCollectionRules/$dcrImmutableId/streams/${streamName}?api-version=2023-01-01"
+$uri2 = "$dceEndpoint/dataCollectionRules/$dcrImmutableId/streams/${streamName}2?api-version=2023-01-01"
 
 Write-Host "Sending sample data..."
-Invoke-RestMethod -Uri $uri -Method "Post" -Body $body -Headers $headers -TimeoutSec 20 -MaximumRetryCount 3
-Invoke-RestMethod -Uri $uri2 -Method "Post" -Body $body -Headers $headers -TimeoutSec 20 -MaximumRetryCount 3
+Invoke-RestMethod -Uri $uri -Method "Post" -Body $body -Headers $headers -TimeoutSec 40 -MaximumRetryCount 3
+Invoke-RestMethod -Uri $uri2 -Method "Post" -Body $body -Headers $headers -TimeoutSec 40 -MaximumRetryCount 3
 
 ##################
 ### Step 4: Sleep to allow time for data to reflect in the workspace tables.

diff --git a/sdk/monitor/tests.yml b/sdk/monitor/tests.yml
@@ -6,6 +6,15 @@ extends:
       ServiceDirectory: monitor
       TestTimeoutInMinutes: 300
       BuildTargetingString: azure-monitor-*
+      SupportedClouds: 'Public,UsGov,China'
+      CloudConfig:
+        Public:
+          SubscriptionConfiguration: $(sub-config-azure-cloud-test-resources)
+        UsGov:
+          SubscriptionConfiguration: $(sub-config-gov-test-resources)
+        China:
+          SubscriptionConfiguration: $(sub-config-cn-test-resources)
+          Location: chinanorth3
       EnvVars:
         AZURE_SUBSCRIPTION_ID: $(MONITOR_SUBSCRIPTION_ID)
         AZURE_TENANT_ID: $(MONITOR_TENANT_ID)