Improve connection cancellation token tests

Azure · Jan 19, 2024 · 9667771 · 9667771
1 parent 8d21715
commit 9667771
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/test/Common/TestRuntimeProvider.cs b/test/Common/TestRuntimeProvider.cs
@@ -85,15 +85,21 @@ public void EndOpenLink(IAsyncResult result)
 
     class TestSaslPlainAuthenticator : ISaslPlainAuthenticator
     {
-        public Task<IPrincipal> AuthenticateAsync(string identity, string password)
+        public int DelayInMilliseconds { get; set; }
+
+        public async Task<IPrincipal> AuthenticateAsync(string identity, string password)
         {
+            if (this.DelayInMilliseconds > 0)
+            {
+                await Task.Delay(this.DelayInMilliseconds);
+            }
+
             if (identity != password)
             {
                 throw new UnauthorizedAccessException();
             }
 
-            IPrincipal principal = new GenericPrincipal(new GenericIdentity(identity), new string[] { "SEND", "RECV" });
-            return Task.FromResult(principal);
+            return new GenericPrincipal(new GenericIdentity(identity), new string[] { "SEND", "RECV" });
         }
     }
 

diff --git a/test/TestCases/CancellationTokenTests.cs b/test/TestCases/CancellationTokenTests.cs
@@ -87,7 +87,7 @@ public Task ConnectionFactoryCanceledTest()
 
         async Task RunConnectionFactoryTest(bool cancelBefore)
         {
-            AmqpConnectionListener listener = OpenListener();
+            AmqpConnectionListener listener = OpenListener(saslDelayMs: 200);
 
             try
             {
@@ -130,7 +130,7 @@ public Task ConnectionOpenCanceledTest()
 
         async Task RunConnectionOpenTest(bool cancelBefore)
         {
-            AmqpConnectionListener listener = OpenListener();
+            AmqpConnectionListener listener = OpenListener(saslDelayMs: 200);
 
             try
             {
@@ -724,11 +724,11 @@ await Assert.ThrowsAnyAsync<TaskCanceledException>(async () =>
             }
         }
 
-        AmqpConnectionListener OpenListener(TestRuntimeProvider runtimeProvider = null)
+        AmqpConnectionListener OpenListener(TestRuntimeProvider runtimeProvider = null, int saslDelayMs = 0)
         {
             AmqpSettings settings = new AmqpSettings { RuntimeProvider = runtimeProvider ?? new TestRuntimeProvider() };
             var saslProvider = new SaslTransportProvider(AmqpVersion.V100);
-            saslProvider.AddHandler(new SaslPlainHandler(new TestSaslPlainAuthenticator()));
+            saslProvider.AddHandler(new SaslPlainHandler(new TestSaslPlainAuthenticator() { DelayInMilliseconds = saslDelayMs }));
             settings.TransportProviders.Add(saslProvider);
             var listener = new AmqpConnectionListener(new[] { addressUri.AbsoluteUri }, settings, new AmqpConnectionSettings());
             listener.Open();