Improve connection cancellation token tests

test/Common/TestRuntimeProvider.cs

@@ -85,15 +85,21 @@ namespace Test.Microsoft.Azure.Amqp
 
     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" });
         }
     }
 

test/TestCases/CancellationTokenTests.cs

@@ -87,7 +87,7 @@ namespace Test.Microsoft.Azure.Amqp
 
         async Task RunConnectionFactoryTest(bool cancelBefore)
         {
-            AmqpConnectionListener listener = OpenListener();
+            AmqpConnectionListener listener = OpenListener(saslDelayMs: 200);
 
             try
             {
@@ -130,7 +130,7 @@ namespace Test.Microsoft.Azure.Amqp
 
         async Task RunConnectionOpenTest(bool cancelBefore)
         {
-            AmqpConnectionListener listener = OpenListener();
+            AmqpConnectionListener listener = OpenListener(saslDelayMs: 200);
 
             try
             {
@@ -724,11 +724,11 @@ namespace Test.Microsoft.Azure.Amqp
             }
         }
 
-        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();