Add a few tests for Task.Yield in async methods (#28887)

1 files changed, 30 insertions, 0 deletions

diff --git a/src/System.Threading.Tasks/tests/System.Runtime.CompilerServices/YieldAwaitableTests.cs b/src/System.Threading.Tasks/tests/System.Runtime.CompilerServices/YieldAwaitableTests.cs
index 852de99779..92ffdc4ab9 100644
--- a/src/System.Threading.Tasks/tests/System.Runtime.CompilerServices/YieldAwaitableTests.cs
+++ b/src/System.Threading.Tasks/tests/System.Runtime.CompilerServices/YieldAwaitableTests.cs
@@ -143,8 +143,38 @@ namespace System.Threading.Tasks.Tests
             SynchronizationContext.SetSynchronizationContext(new ValidateCorrectContextSynchronizationContext());
             var ya = Task.Yield().GetAwaiter();
             Assert.Throws<ArgumentNullException>(() => { ya.OnCompleted(null); });
+            SynchronizationContext.SetSynchronizationContext(null);
         }
 
+        [Fact]
+        public static async Task AsyncMethod_Yields_ReturnsToDefaultTaskScheduler()
+        {
+            await Task.Yield();
+            Assert.Same(TaskScheduler.Default, TaskScheduler.Current);
+        }
+
+        [Fact]
+        public static async Task AsyncMethod_Yields_ReturnsToCorrectTaskScheduler()
+        {
+            QUWITaskScheduler ts = new QUWITaskScheduler();
+            Assert.NotSame(ts, TaskScheduler.Current);
+            await Task.Factory.StartNew(async delegate
+            {
+                Assert.Same(ts, TaskScheduler.Current);
+                await Task.Yield();
+                Assert.Same(ts, TaskScheduler.Current);
+            }, CancellationToken.None, TaskCreationOptions.None, ts).Unwrap();
+            Assert.NotSame(ts, TaskScheduler.Current);
+        }
+
+        [Fact]
+        public static async Task AsyncMethod_Yields_ReturnsToCorrectSynchronizationContext()
+        {
+            var sc = new ValidateCorrectContextSynchronizationContext ();
+            SynchronizationContext.SetSynchronizationContext(sc);
+            await Task.Yield();
+            Assert.Equal(1, sc.PostCount);
+        }
         #region Helper Methods / Classes
 
         private class ValidateCorrectContextSynchronizationContext : SynchronizationContext