[Internal] Client Telemetry: Fixes tests leaking environment variables

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientTelemetryTests.cs

@@ -30,7 +30,6 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
     [TestClass]
     public class ClientTelemetryTests : BaseCosmosClientHelper
     {
-        private const string telemetryEndpointUrl = "http://dummy.telemetry.endpoint/";
         private const int scheduledInSeconds = 1;
         private static readonly object jsonObject = JsonConvert.DeserializeObject("{\"compute\":{\"azEnvironment\":\"AzurePublicCloud\",\"customData\":\"\",\"isHostCompatibilityLayerVm\":\"false\",\"licenseType\":\"\",\"location\":\"eastus\",\"name\":\"sourabh-testing\",\"offer\":\"UbuntuServer\",\"osProfile\":{\"adminUsername\":\"azureuser\",\"computerName\":\"sourabh-testing\"},\"osType\":\"Linux\",\"placementGroupId\":\"\",\"plan\":{\"name\":\"\",\"product\":\"\",\"publisher\":\"\"},\"platformFaultDomain\":\"0\",\"platformUpdateDomain\":\"0\",\"provider\":\"Microsoft.Compute\",\"publicKeys\":[{\"keyData\":\"ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQC5uCeOAm3ehmhI+2PbMoMl17Eo\r\nqfHKCycSaBJsv9qxlmBOuFheSJc1XknJleXUSsuTO016/d1PyWpevnqOZNRksWoa\r\nJvQ23sDTxcK+X2OP3QlCUeX4cMjPXqlL8z1UYzU4Bx3fFvf8fs67G3N72sxWBw5P\r\nZyuXyhBm0NCe/2NYMKgEDT4ma8XszO0ikbhoPKbMbgHAQk/ktWQHNcqYOPQKEWqp\r\nEK1R0rjS2nmtovfScP/ZGXcvOpJ1/NDBo4dh1K+OxOGM/4PSH/F448J5Zy4eAyEk\r\nscys+IpeIOTOlRUy/703SNIX0LEWlnYqbyL9c1ypcYLQqF76fKkDfzzFI/OWVlGw\r\nhj/S9uP8iMsR+fhGIbn6MAa7O4DWPWLuedSp7KDYyjY09gqNJsfuaAJN4LiC6bPy\r\nhknm0PVLK3ux7EUOt+cZrHCdIFWbdOtxiPNIl1tkv9kV5aE5Aj2gJm4MeB9uXYhS\r\nOuksboBc0wyUGrl9+XZJ1+NlZOf7IjVi86CieK8= generated-by-azure\r\n\",\"path\":\"/home/azureuser/.ssh/authorized_keys\"}],\"publisher\":\"Canonical\",\"resourceGroupName\":\"sourabh-telemetry-sdk\",\"resourceId\":\"/subscriptions/8fba6d4f-7c37-4d13-9063-fd58ad2b86e2/resourceGroups/sourabh-telemetry-sdk/providers/Microsoft.Compute/virtualMachines/sourabh-testing\",\"securityProfile\":{\"secureBootEnabled\":\"false\",\"virtualTpmEnabled\":\"false\"},\"sku\":\"18.04-LTS\",\"storageProfile\":{\"dataDisks\":[],\"imageReference\":{\"id\":\"\",\"offer\":\"UbuntuServer\",\"publisher\":\"Canonical\",\"sku\":\"18.04-LTS\",\"version\":\"latest\"},\"osDisk\":{\"caching\":\"ReadWrite\",\"createOption\":\"FromImage\",\"diffDiskSettings\":{\"option\":\"\"},\"diskSizeGB\":\"30\",\"encryptionSettings\":{\"enabled\":\"false\"},\"image\":{\"uri\":\"\"},\"managedDisk\":{\"id\":\"/subscriptions/8fba6d4f-7c37-4d13-9063-fd58ad2b86e2/resourceGroups/sourabh-telemetry-sdk/providers/Microsoft.Compute/disks/sourabh-testing_OsDisk_1_9a54abfc5ba149c6a106bd9e5b558c2a\",\"storageAccountType\":\"Premium_LRS\"},\"name\":\"sourabh-testing_OsDisk_1_9a54abfc5ba149c6a106bd9e5b558c2a\",\"osType\":\"Linux\",\"vhd\":{\"uri\":\"\"},\"writeAcceleratorEnabled\":\"false\"}},\"subscriptionId\":\"8fba6d4f-7c37-4d13-9063-fd58ad2b86e2\",\"tags\":\"azsecpack:nonprod;platformsettings.host_environment.service.platform_optedin_for_rootcerts:true\",\"tagsList\":[{\"name\":\"azsecpack\",\"value\":\"nonprod\"},{\"name\":\"platformsettings.host_environment.service.platform_optedin_for_rootcerts\",\"value\":\"true\"}],\"version\":\"18.04.202103250\",\"vmId\":\"d0cb93eb-214b-4c2b-bd3d-cc93e90d9efd\",\"vmScaleSetName\":\"\",\"vmSize\":\"Standard_D2s_v3\",\"zone\":\"1\"},\"network\":{\"interface\":[{\"ipv4\":{\"ipAddress\":[{\"privateIpAddress\":\"10.0.7.5\",\"publicIpAddress\":\"\"}],\"subnet\":[{\"address\":\"10.0.7.0\",\"prefix\":\"24\"}]},\"ipv6\":{\"ipAddress\":[]},\"macAddress\":\"000D3A8F8BA0\"}]}}");
 
@@ -46,12 +45,8 @@ public class ClientTelemetryTests : BaseCosmosClientHelper
         private HttpClientHandlerHelper httpHandlerForNonAzureInstance;
 
         [ClassInitialize]
-        public static void ClassInitialize(TestContext context)
+        public static void ClassInitialize(TestContext _)
         {
-            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetryEnabled, "true");
-            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetrySchedulingInSeconds, "1");
-            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetryEndpoint, telemetryEndpointUrl);
-
             SystemUsageMonitor oldSystemUsageMonitor = (SystemUsageMonitor)typeof(DiagnosticsHandlerHelper)
                 .GetField("systemUsageMonitor", BindingFlags.Instance | BindingFlags.NonPublic).GetValue(DiagnosticsHandlerHelper.Instance);
             oldSystemUsageMonitor.Stop();
@@ -62,6 +57,8 @@ public static void ClassInitialize(TestContext context)
         [TestInitialize]
         public void TestInitialize()
         {
+            Util.EnableClientTelemetryEnvironmentVariables();
+
             this.actualInfo = new List<ClientTelemetryProperties>();
 
             this.httpHandler = new HttpClientHandlerHelper
@@ -134,16 +131,6 @@ public void TestInitialize()
                                         .WithApplicationPreferredRegions(this.preferredRegionList);
         }
 
-        [ClassCleanup]
-        public static void FinalCleanup()
-        {
-            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetryEnabled, null);
-            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetrySchedulingInSeconds, null);
-            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetryEndpoint, null);
-
-            ClientTelemetryTests.ResetSystemUsageMonitor(false);
-        }
-
         private static void ResetSystemUsageMonitor(bool isTelemetryEnabled)
         {
             ClientTelemetryTests.systemUsageMonitor?.Stop();
@@ -186,6 +173,14 @@ public async Task Cleanup()
             azMetadataField.SetValue(null, null);
 
             await base.TestCleanup();
+
+            Util.DisableClientTelemetryEnvironmentVariables();
+        }
+
+        [ClassCleanup]
+        public static void FinalCleanup()
+        {
+            ClientTelemetryTests.ResetSystemUsageMonitor(false);
         }
 
         [TestMethod]

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/PartitionKeyRangeCacheTests.cs

@@ -19,11 +19,21 @@ public class PartitionKeyRangeCacheTests
     {
         private bool loopBackgroundOperaitons = false;
 
-        [TestMethod]
-        public async Task VerifyPkRangeCacheRefreshOnSplitWithErrorsAsync()
+        [TestInitialize]
+        public void TestInitialize()
+        {
+            this.loopBackgroundOperaitons = false;
+        }
+
+        [TestCleanup]
+        public void TestCleanup()
         {
             this.loopBackgroundOperaitons = false;
+        }
 
+        [TestMethod]
+        public async Task VerifyPkRangeCacheRefreshOnSplitWithErrorsAsync()
+        {
             int throwOnPkRefreshCount = 3;
             int pkRangeCalls = 0;
             bool causeSplitExceptionInRntbdCall = false;

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/SynchronizationContextTests.cs

@@ -13,10 +13,17 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
     [TestClass]
     public class SynchronizationContextTests
     {
-        [TestMethod]
+        [DataTestMethod]
+        [DataRow(false, DisplayName = "SynchronizationContextTests - Client Telemetry disabled")]
+        [DataRow(true, DisplayName = "SynchronizationContextTests - Client Telemetry enabled")]
         [Timeout(30000)]
-        public void VerifySynchronizationContextDoesNotLock()
+        public void VerifySynchronizationContextDoesNotLock(bool withClientTelemetry)
         {
+            if (withClientTelemetry)
+            {
+                Util.EnableClientTelemetryEnvironmentVariables();
+            }
+
             string databaseId = Guid.NewGuid().ToString();
             SynchronizationContext prevContext = SynchronizationContext.Current;
             try
@@ -29,7 +36,6 @@ public void VerifySynchronizationContextDoesNotLock()
                     {
                         Cosmos.Database database = client.CreateDatabaseAsync(databaseId).GetAwaiter().GetResult();
                         database = client.CreateDatabaseIfNotExistsAsync(databaseId).GetAwaiter().GetResult();
-
                         database.ReadStreamAsync().ConfigureAwait(false).GetAwaiter().GetResult();
                         database.ReadAsync().ConfigureAwait(false).GetAwaiter().GetResult();
 
@@ -118,7 +124,7 @@ public void VerifySynchronizationContextDoesNotLock()
 
                         double cost = container.GetItemLinqQueryable<ToDoActivity>(
                             allowSynchronousQueryExecution: true).Select(x => x.cost).Sum();
-
+                        
                         ItemResponse<ToDoActivity> deleteResponse = container.DeleteItemAsync<ToDoActivity>(partitionKey: new Cosmos.PartitionKey(testItem.pk), id: testItem.id).ConfigureAwait(false).GetAwaiter().GetResult();
                         Assert.IsNotNull(deleteResponse);
                     }
@@ -127,13 +133,21 @@ public void VerifySynchronizationContextDoesNotLock()
             finally
             {
                 SynchronizationContext.SetSynchronizationContext(prevContext);
+
                 using (CosmosClient client = TestCommon.CreateCosmosClient())
                 {
                     client.GetDatabase(databaseId).DeleteAsync().GetAwaiter().GetResult();
                 }
+
+                if (withClientTelemetry)
+                {
+                    Util.DisableClientTelemetryEnvironmentVariables();
+                }
             }
         }
 
+
+
         public class TestSynchronizationContext : SynchronizationContext
         {
             private object locker = new object();

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/Util.cs

@@ -5,11 +5,15 @@ namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Globalization;
+    using System.IO;
     using System.Linq;
     using System.Net;
+    using System.Text;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.Services.Management.Tests;
+    using Microsoft.Azure.Cosmos.Telemetry;
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Client;
     using Microsoft.Azure.Documents.Collections;
@@ -514,5 +518,57 @@ internal static void LogRequestOptions(RequestOptions options, bool shouldLogOff
                 options.OfferType,
                 options.OfferThroughput);
         }
+
+        internal static void EnableClientTelemetryEnvironmentVariables()
+        {
+            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetryEnabled, "true");
+            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetrySchedulingInSeconds, "1");
+            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetryEndpoint, "http://dummy.telemetry.endpoint/");
+        }
+
+        internal static void DisableClientTelemetryEnvironmentVariables()
+        {
+            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetryEnabled, null);
+            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetrySchedulingInSeconds, null);
+            Environment.SetEnvironmentVariable(ClientTelemetryOptions.EnvPropsClientTelemetryEndpoint, null);
+        }
+
+        /// <summary>
+        /// Enables traces for local debugging
+        /// </summary>
+        internal static void EnableTracesForDebugging()
+        {
+            Type defaultTrace = Type.GetType("Microsoft.Azure.Cosmos.Core.Trace.DefaultTrace,Microsoft.Azure.Cosmos.Direct");
+            TraceSource traceSource = (TraceSource)defaultTrace.GetProperty("TraceSource").GetValue(null);
+            traceSource.Switch.Level = SourceLevels.All;
+            traceSource.Listeners.Clear();
+            traceSource.Listeners.Add(new DirectToConsoleTraceListener());
+        }
+
+        public class DirectToConsoleTraceListener : TextWriterTraceListener
+        {
+            public DirectToConsoleTraceListener() : base(new DirectToConsoleTextWriter())
+            {
+            }
+
+            public override void Close()
+            {
+            }
+        }
+
+        public class DirectToConsoleTextWriter : TextWriter
+        {
+            public override Encoding Encoding => Console.Out.Encoding;
+
+            public override void Write(string value)
+            {
+                Logger.LogLine(value);
+            }
+
+            public override void WriteLine(string value)
+            {
+                Logger.LogLine(value);
+            }
+        }
     }
 }