[Internal] Tests: Adds timeouts and polling loops for low-impact test reliability

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

@@ -77,6 +77,7 @@ public override async Task Cleanup()
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct, true)]
         [DataRow(ConnectionMode.Gateway, true)]
         [DataRow(ConnectionMode.Direct, false)]
@@ -87,6 +88,7 @@ public override async Task PointSuccessOperationsTest(ConnectionMode mode, bool
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task PointReadFailureOperationsTest(ConnectionMode mode)
@@ -95,6 +97,7 @@ public override async Task PointReadFailureOperationsTest(ConnectionMode mode)
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task StreamReadFailureOperationsTest(ConnectionMode mode)
@@ -103,6 +106,7 @@ public override async Task StreamReadFailureOperationsTest(ConnectionMode mode)
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task StreamOperationsTest(ConnectionMode mode)
@@ -111,6 +115,7 @@ public override async Task StreamOperationsTest(ConnectionMode mode)
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task BatchOperationsTest(ConnectionMode mode)
@@ -119,6 +124,7 @@ public override async Task BatchOperationsTest(ConnectionMode mode)
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task SingleOperationMultipleTimesTest(ConnectionMode mode)
@@ -127,6 +133,7 @@ public override async Task SingleOperationMultipleTimesTest(ConnectionMode mode)
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task QueryOperationSinglePartitionTest(ConnectionMode mode)
@@ -135,6 +142,7 @@ public override async Task QueryOperationSinglePartitionTest(ConnectionMode mode
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task QueryMultiPageSinglePartitionOperationTest(ConnectionMode mode)
@@ -143,6 +151,7 @@ public override async Task QueryMultiPageSinglePartitionOperationTest(Connection
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task QueryOperationCrossPartitionTest(ConnectionMode mode)
@@ -151,6 +160,7 @@ public override async Task QueryOperationCrossPartitionTest(ConnectionMode mode)
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task QueryOperationMutiplePageCrossPartitionTest(ConnectionMode mode)
@@ -159,6 +169,7 @@ public override async Task QueryOperationMutiplePageCrossPartitionTest(Connectio
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         [DataRow(ConnectionMode.Gateway)]
         public override async Task QueryOperationInvalidContinuationTokenTest(ConnectionMode mode)
@@ -167,6 +178,7 @@ public override async Task QueryOperationInvalidContinuationTokenTest(Connection
         }
 
         [TestMethod]
+        [Timeout(300000)]
         [DataRow(ConnectionMode.Direct)]
         public override async Task CreateItemWithSubStatusCodeTest(ConnectionMode mode)
         {

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Tracing/EndToEndTraceWriterBaselineTests.cs

@@ -494,6 +494,7 @@ public async Task ChangeFeedAsync()
 
         [TestMethod]
         [TestCategory("Flaky")]
+        [Timeout(300000)]
         public async Task QueryAsync()
         {
             List<Input> inputs = new List<Input>();
@@ -818,6 +819,7 @@ public async Task ValidateInvalidCredentialsTraceAsync()
 
         [TestMethod]
         [TestCategory("Flaky")]
+        [Timeout(300000)]
         public async Task TypedPointOperationsAsync()
         {
             List<Input> inputs = new List<Input>();

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchAsyncStreamerTests.cs

@@ -150,9 +150,14 @@ public async Task ValidatesCongestionControlAsync()
             // 300 batch request should atleast sum up to 1000 ms barrier with wait time of 20ms in executor
             await Task.WhenAll(contexts);
 
-            await Task.Delay(2000);
+            // Poll for semaphore count to increase, with a reasonable timeout
+            System.Diagnostics.Stopwatch sw = System.Diagnostics.Stopwatch.StartNew();
+            while (newLimiter.CurrentCount < 2 && sw.Elapsed < TimeSpan.FromSeconds(10))
+            {
+                await Task.Delay(200);
+            }
 
-            Assert.IsTrue(newLimiter.CurrentCount >= 2, "Count of threads that can enter into semaphore should increase atleast by 1");
+            Assert.IsTrue(newLimiter.CurrentCount >= 2, $"Count of threads that can enter into semaphore should increase atleast by 1. Actual: {newLimiter.CurrentCount}");
         }
 
         [TestMethod]

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/ChangeFeed/PartitionControllerTests.cs

@@ -5,6 +5,7 @@
 namespace Microsoft.Azure.Cosmos.ChangeFeed.Tests
 {
     using System;
+    using System.Diagnostics;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.ChangeFeed.Exceptions;
@@ -229,7 +230,22 @@ public async Task Controller_ShouldReleasesLease_IfObserverExits()
                 .Returns(new PartitionSupervisorCore(this.lease, this.observer, this.partitionProcessor, this.leaseRenewer));
 
             await this.sut.AddOrUpdateLeaseAsync(this.lease).ConfigureAwait(false);
-            await Task.Delay(TimeSpan.FromMilliseconds(100)).ConfigureAwait(false);
+
+            // Poll for lease release with a bounded timeout instead of a fixed delay
+            Stopwatch sw = Stopwatch.StartNew();
+            while (sw.Elapsed < TimeSpan.FromSeconds(5))
+            {
+                try
+                {
+                    Mock.Get(this.leaseManager)
+                        .Verify(manager => manager.ReleaseAsync(It.IsAny<DocumentServiceLease>()), Times.Once);
+                    break;
+                }
+                catch (MockException)
+                {
+                    await Task.Delay(50).ConfigureAwait(false);
+                }
+            }
 
             Mock.Get(this.leaseManager)
                 .Verify(manager => manager.ReleaseAsync(It.IsAny<DocumentServiceLease>()), Times.Once);

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosAuthorizationTests.cs

@@ -329,9 +329,11 @@ public async Task TestTokenCredentialBackgroundRefreshAsync()
                 Assert.AreEqual(token1, t2);
 
                 // Wait until the background refresh occurs.
+                Stopwatch sw = Stopwatch.StartNew();
                 while (testTokenCredential.NumTimesInvoked == 1)
                 {
-                    await Task.Delay(500);
+                    Assert.IsTrue(sw.Elapsed < TimeSpan.FromSeconds(20), "Background token refresh did not occur within 20 seconds.");
+                    await Task.Delay(200);
                 }
 
                 string t3 = await tokenCredentialCache.GetTokenAsync(NoOpTrace.Singleton);

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosHttpClientCoreTests.cs

@@ -57,6 +57,7 @@ static Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellati
 
         [TestMethod]
         [TestCategory("Flaky")]
+        [Timeout(120000)]
         public async Task RetryTransientIssuesTestAsync()
         {
             using CancellationTokenSource cancellationTokenSource1 = new CancellationTokenSource();
@@ -68,15 +69,15 @@ public async Task RetryTransientIssuesTestAsync()
             {
                 {HttpTimeoutPolicyControlPlaneRead.Instance,  new List<TimeSpan>()
                 {
-                    TimeSpan.FromSeconds(5.1),
-                    TimeSpan.FromSeconds(10.1),
-                    TimeSpan.FromSeconds(20.1)
+                    TimeSpan.FromSeconds(6),
+                    TimeSpan.FromSeconds(11),
+                    TimeSpan.FromSeconds(21)
                 }},
                 {HttpTimeoutPolicyControlPlaneRetriableHotPath.Instance,  new List<TimeSpan>()
                 {
-                    TimeSpan.FromSeconds(.6),
-                    TimeSpan.FromSeconds(5.1),
-                    TimeSpan.FromSeconds(65.1)
+                    TimeSpan.FromSeconds(1),
+                    TimeSpan.FromSeconds(6),
+                    TimeSpan.FromSeconds(66)
                 }},
             };
 
@@ -388,6 +389,7 @@ Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, CancellationToken
 
         [TestMethod]
         [TestCategory("Flaky")]
+        [Timeout(120000)]
         public async Task RetryTransientIssuesForQueryPlanTestAsync()
         {
             DocumentServiceRequest documentServiceRequest = DocumentServiceRequest.Create(
@@ -411,7 +413,7 @@ async Task<HttpResponseMessage> sendFunc(HttpRequestMessage request, Cancellatio
                 if (count <= 2)
                 {
                     Assert.IsFalse(cancellationToken.IsCancellationRequested);
-                    await Task.Delay(retry.Current.requestTimeout + TimeSpan.FromSeconds(.1));
+                    await Task.Delay(retry.Current.requestTimeout + TimeSpan.FromSeconds(1));
                     cancellationToken.ThrowIfCancellationRequested();
                     Assert.Fail("Cancellation token should be canceled");
                 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GlobalEndpointManagerTest.cs

@@ -29,6 +29,7 @@ public class GlobalEndpointManagerTest
         /// </summary>
         [TestMethod]
         [TestCategory("Flaky")]
+        [Timeout(30000)]
         public async Task EndpointFailureMockTest()
         {
             Environment.SetEnvironmentVariable("MinimumIntervalForNonForceRefreshLocationInMS", "100");
@@ -94,11 +95,20 @@ public async Task EndpointFailureMockTest()
                     Assert.AreEqual(globalEndpointManager.WriteEndpoints[0], globalEndpointManager.ReadEndpoints[0]);
 
                     getAccountInfoCount = 0;
-                    //Sleep 3 seconds for the unavailable endpoint entry to expire and background refresh timer to kick in
-                    await Task.Delay(TimeSpan.FromSeconds(3));
+                    //Poll for the unavailable endpoint entry to expire and background refresh timer to kick in
+                    Stopwatch sw = Stopwatch.StartNew();
+                    while (sw.Elapsed < TimeSpan.FromSeconds(10))
+                    {
+                        await Task.Delay(200);
+                        await globalEndpointManager.RefreshLocationAsync();
+                        if (globalEndpointManager.ReadEndpoints[0].Equals(new Uri(readLocation1.Endpoint)))
+                        {
+                            break;
+                        }
+                    }
+
                     Assert.IsTrue(getAccountInfoCount > 0, "Callback is not working. There should be at least one call in this time frame.");
 
-                    await globalEndpointManager.RefreshLocationAsync();
                     Assert.AreEqual(new Uri(readLocation1.Endpoint), globalEndpointManager.ReadEndpoints[0], "Read endpoint did not switch back to location 1 after the unavailable entry expired.");
                 }