Improved continueOnCapturedContext test

AsyncKeyedLock.Tests/AsyncKeyedLocker/OriginalTests.cs

@@ -58,7 +58,7 @@
         {
             var locks = 5000;
             var concurrency = 50;
             var asyncKeyedLocker = new AsyncKeyedLock.AsyncKeyedLocker();
             var concurrentQueue = new ConcurrentQueue<(bool entered, int key)>();
 
             var tasks = Enumerable.Range(1, locks * concurrency)
@@ -462,7 +462,7 @@
         public async Task Test1AtATime()
         {
             var range = 25;
             var asyncKeyedLocker = new AsyncKeyedLock.AsyncKeyedLocker();
             var concurrentQueue = new ConcurrentQueue<int>();
 
             var tasks = Enumerable.Range(1, range * 2)
@@ -497,7 +497,7 @@
         {
             var range = 4;
             var asyncKeyedLocker = new AsyncKeyedLock.AsyncKeyedLocker(o => o.MaxCount = 2);
             var concurrentQueue = new ConcurrentQueue<int>();
 
             var tasks = Enumerable.Range(1, range * 4)
                 .Select(async i =>
@@ -595,26 +595,47 @@
         }
 
         [Fact]
-        public async Task TestContinueOnCapturedContext()
+        public Task TestContinueOnCapturedContextTrue()
+            => TestContinueOnCapturedContext(true);
+
+        [Fact]
+        public Task TestContinueOnCapturedContextFalse()
+            => TestContinueOnCapturedContext(false);
+
+        private async Task TestContinueOnCapturedContext(bool continueOnCapturedContext)
         {
+            const string Key = "test";
+
             var asyncKeyedLocker = new AsyncKeyedLocker<string>();
             var testContext = new TestSynchronizationContext();
-            var currentThreadId = Environment.CurrentManagedThreadId;
 
-            async Task Callback()
+            void Callback()
             {
-                await Task.Yield();
-
-                SynchronizationContext.Current.Should().Be(testContext);
-                Environment.CurrentManagedThreadId.Should().Be(currentThreadId);
+                if(continueOnCapturedContext)
+                {
+                    Environment.CurrentManagedThreadId.Should().Be(testContext.LastPostThreadId);
+                }
+                else
+                {
+                    testContext.LastPostThreadId.Should().Be(default);
+                }
             }
 
             var previousContext = SynchronizationContext.Current;
             SynchronizationContext.SetSynchronizationContext(testContext);
 
             try
             {
-                await asyncKeyedLocker.TryLockAsync("test", Callback, 0, true);
+                // This is just to make WaitAsync in TryLockAsync not finish synchronously
+                var obj = asyncKeyedLocker.Lock(Key);
+
+                _ = Task.Run(async () =>
+                {
+                    await Task.Delay(1000);
+                    obj.Dispose();
+                });
+
+                await asyncKeyedLocker.TryLockAsync(Key, Callback, 5000, continueOnCapturedContext);
             }
             finally
             {

AsyncKeyedLock.Tests/Helpers/TestSynchronizatonContext.cs

@@ -1,10 +1,15 @@
 
+using System.Collections.Concurrent;
+
 namespace AsyncKeyedLock.Tests.Helpers
 {
     public class TestSynchronizationContext : SynchronizationContext
     {
+        public int LastPostThreadId { get; private set; }
+
         public override void Post(SendOrPostCallback d, object? state)
         {
+            LastPostThreadId = Environment.CurrentManagedThreadId;
             d(state);
         }
     }

AsyncKeyedLock.Tests/StripedAsyncKeyedLocker/OriginalTests.cs

@@ -391,26 +391,47 @@ public async Task Test2AtATimeGenerics()
         }
 
         [Fact]
-        public async Task TestContinueOnCapturedContext()
+        public Task TestContinueOnCapturedContextTrue()
+            => TestContinueOnCapturedContext(true);
+
+        [Fact]
+        public Task TestContinueOnCapturedContextFalse()
+            => TestContinueOnCapturedContext(false);
+
+        private async Task TestContinueOnCapturedContext(bool continueOnCapturedContext)
         {
-            var stripedAsyncKeyedLocker = new StripedAsyncKeyedLocker<string>();
+            const string Key = "test";
+
+            var asyncKeyedLocker = new AsyncKeyedLocker<string>();
             var testContext = new TestSynchronizationContext();
-            var currentThreadId = Environment.CurrentManagedThreadId;
 
-            async Task Callback()
+            void Callback()
             {
-                await Task.Yield();
-
-                SynchronizationContext.Current.Should().Be(testContext);
-                Environment.CurrentManagedThreadId.Should().Be(currentThreadId);
+                if (continueOnCapturedContext)
+                {
+                    Environment.CurrentManagedThreadId.Should().Be(testContext.LastPostThreadId);
+                }
+                else
+                {
+                    testContext.LastPostThreadId.Should().Be(default);
+                }
             }
 
             var previousContext = SynchronizationContext.Current;
             SynchronizationContext.SetSynchronizationContext(testContext);
 
             try
             {
-                await stripedAsyncKeyedLocker.TryLockAsync("test", Callback, 0, true);
+                // This is just to make WaitAsync in TryLockAsync not finish synchronously
+                var obj = asyncKeyedLocker.Lock(Key);
+
+                _ = Task.Run(async () =>
+                {
+                    await Task.Delay(1000);
+                    obj.Dispose();
+                });
+
+                await asyncKeyedLocker.TryLockAsync(Key, Callback, 5000, continueOnCapturedContext);
             }
             finally
             {