Make HeadlessUnitTestSession run individual frames per test

src/Headless/Avalonia.Headless.NUnit/AvaloniaTestMethodCommand.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Reflection;
 using System.Threading.Tasks;
+using Avalonia.Threading;
 using NUnit.Framework.Interfaces;
 using NUnit.Framework.Internal;
 using NUnit.Framework.Internal.Commands;
@@ -30,12 +31,12 @@ internal class AvaloniaTestMethodCommand : DelegatingTestCommand
         {
             if (s_beforeTest.GetValue(beforeAndAfterTestCommand) is Action<TestExecutionContext> beforeTest)
             {
-                Action<TestExecutionContext> beforeAction = c => session.Dispatcher.Invoke(() => beforeTest(c));
+                Action<TestExecutionContext> beforeAction = c => session.Dispatch(() => beforeTest(c), default);
                 s_beforeTest.SetValue(beforeAndAfterTestCommand, beforeAction);
             }
             if (s_afterTest.GetValue(beforeAndAfterTestCommand) is Action<TestExecutionContext> afterTest)
             {
-                Action<TestExecutionContext> afterAction = c => session.Dispatcher.Invoke(() => afterTest(c));
+                Action<TestExecutionContext> afterAction = c => session.Dispatch(() => afterTest(c), default);
                 s_afterTest.SetValue(beforeAndAfterTestCommand, afterAction);
             }
         }
@@ -52,10 +53,10 @@ internal class AvaloniaTestMethodCommand : DelegatingTestCommand
 
         return command;
     }
-    
+
     public override TestResult Execute(TestExecutionContext context)
     {
-        return _session.Dispatcher.InvokeOnQueue(() => ExecuteTestMethod(context));
+        return _session.Dispatch(() => ExecuteTestMethod(context), default).GetAwaiter().GetResult();
     }
 
     // Unfortunately, NUnit has issues with custom synchronization contexts, which means we need to add some hacks to make it work.

src/Headless/Avalonia.Headless.NUnit/ExecutionQueue.cs

@@ -1,71 +0,0 @@
-using System;
-using System.Collections.Generic;
-using System.Threading;
-using System.Threading.Tasks;
-using Avalonia.Threading;
-
-namespace Avalonia.Headless.NUnit;
-
-// To force all tests (which might run in parallel) to be executed in a queue.
-internal static class ExecutionQueue
-{
-    private static bool s_running;
-    private static Queue<Func<Task>> s_queue = new();
-
-    private static async void TryExecuteNext()
-    {
-        if (s_running || s_queue.Count == 0) return;
-        try
-        {
-            s_running = true;
-            await s_queue.Dequeue()();
-        }
-        finally
-        {
-            s_running = false;
-        }
-        TryExecuteNext();
-    }
-
-    private static void PostToTheQueue(this Dispatcher dispatcher, Func<Task> cb)
-    {
-        dispatcher.Post(() =>
-        {
-            s_queue.Enqueue(cb);
-            TryExecuteNext();
-        });
-    }
-
-    internal static Task<TResult> ExecuteOnQueue<TResult>(this Dispatcher dispatcher, Func<Task<TResult>> cb)
-    {
-        var tcs = new TaskCompletionSource<TResult>();
-        PostToTheQueue(dispatcher, async () =>
-        {
-            try
-            {
-                var result = await cb();
-                tcs.TrySetResult(result);
-            }
-            catch (Exception ex)
-            {
-                tcs.TrySetException(ex);
-            }
-        });
-        return tcs.Task;
-    }
-
-    public static TResult InvokeOnQueue<TResult>(this Dispatcher dispatcher, Func<Task<TResult>> cb, CancellationToken cancellationToken = default)
-    {
-        return dispatcher
-            .InvokeAsync(() => ExecuteOnQueue(dispatcher, cb), DispatcherPriority.Normal, cancellationToken)
-            .GetTask().Unwrap()
-            .GetAwaiter().GetResult();
-    }
-    
-    public static Task<TResult> InvokeOnQueueAsync<TResult>(this Dispatcher dispatcher, Func<Task<TResult>> cb, CancellationToken cancellationToken = default)
-    {
-        return dispatcher
-            .InvokeAsync(() => ExecuteOnQueue(dispatcher, cb), DispatcherPriority.Normal, cancellationToken)
-            .GetTask().Unwrap();
-    }
-}

src/Headless/Avalonia.Headless.XUnit/Avalonia.Headless.XUnit.csproj

@@ -13,12 +13,6 @@
     <ProjectReference Include="..\Avalonia.Headless\Avalonia.Headless.csproj" />
   </ItemGroup>
 
-  <ItemGroup>
-    <Compile Include="..\Avalonia.Headless.NUnit\ExecutionQueue.cs">
-      <Link>ExecutionQueue.cs</Link>
-    </Compile>
-  </ItemGroup>
-
   <Import Project="..\..\..\build\ApiDiff.props" />
   <Import Project="..\..\..\build\DevAnalyzers.props" />
   <Import Project="..\..\..\build\NullableEnable.props" />

src/Headless/Avalonia.Headless.XUnit/AvaloniaTestCase.cs

@@ -1,8 +1,8 @@
 using System;
 using System.ComponentModel;
+using System.Runtime.ExceptionServices;
 using System.Threading;
 using System.Threading.Tasks;
-using Avalonia.Headless.NUnit;
 using Avalonia.Threading;
 using Xunit.Abstractions;
 using Xunit.Sdk;
@@ -37,12 +37,10 @@ internal class AvaloniaTestCase : XunitTestCase
 
         // We need to block the XUnit thread to ensure its concurrency throttle is effective.
         // See https://github.com/AArnott/Xunit.StaFact/pull/55#issuecomment-826187354 for details.
-        var runSummary = session.Dispatcher.InvokeOnQueue(async () =>
-        {
-            var runner = new XunitTestCaseRunner(this, DisplayName, SkipReason, constructorArguments,
-                TestMethodArguments, messageBus, aggregator, cancellationTokenSource);
-            return await runner.RunAsync();
-        }, cancellationTokenSource.Token);
+        var runSummary = AvaloniaTestCaseRunner
+            .RunTest(session, this, DisplayName, SkipReason, constructorArguments,
+                TestMethodArguments, messageBus, aggregator, cancellationTokenSource)
+            .GetAwaiter().GetResult();
 
         return Task.FromResult(runSummary);
     }

src/Headless/Avalonia.Headless.XUnit/AvaloniaTestCaseRunner.cs

@@ -0,0 +1,98 @@
+using System;
+using System.Collections.Generic;
+using System.Reflection;
+using System.Threading;
+using System.Threading.Tasks;
+using Avalonia.Threading;
+using Xunit.Abstractions;
+using Xunit.Sdk;
+
+namespace Avalonia.Headless.XUnit;
+
+internal class AvaloniaTestCaseRunner : XunitTestCaseRunner
+{
+    private readonly Action? _onAfterTestInvoked;
+
+    public AvaloniaTestCaseRunner(
+        Action? onAfterTestInvoked,
+        IXunitTestCase testCase, string displayName, string skipReason, object[] constructorArguments,
+        object[] testMethodArguments, IMessageBus messageBus, ExceptionAggregator aggregator,
+        CancellationTokenSource cancellationTokenSource) : base(testCase, displayName, skipReason, constructorArguments,
+        testMethodArguments, messageBus, aggregator, cancellationTokenSource)
+    {
+        _onAfterTestInvoked = onAfterTestInvoked;
+    }
+
+    public static Task<RunSummary> RunTest(HeadlessUnitTestSession session,
+        IXunitTestCase testCase, string displayName, string skipReason, object[] constructorArguments,
+        object[] testMethodArguments, IMessageBus messageBus, ExceptionAggregator aggregator,
+        CancellationTokenSource cancellationTokenSource)
+    {
+        var afterTest = () => Dispatcher.UIThread.RunJobs();
+        return session.Dispatch(async () =>
+        {
+            var runner = new AvaloniaTestCaseRunner(afterTest, testCase, displayName,
+                skipReason, constructorArguments, testMethodArguments, messageBus, aggregator, cancellationTokenSource);
+            return await runner.RunAsync();
+        }, cancellationTokenSource.Token);
+    }
+    
+    protected override XunitTestRunner CreateTestRunner(ITest test, IMessageBus messageBus, Type testClass,
+        object[] constructorArguments,
+        MethodInfo testMethod, object[] testMethodArguments, string skipReason,
+        IReadOnlyList<BeforeAfterTestAttribute> beforeAfterAttributes,
+        ExceptionAggregator aggregator, CancellationTokenSource cancellationTokenSource)
+    {
+        return new AvaloniaTestRunner(_onAfterTestInvoked, test, messageBus, testClass, constructorArguments,
+            testMethod, testMethodArguments, skipReason, beforeAfterAttributes, aggregator, cancellationTokenSource);
+    }
+
+    private class AvaloniaTestRunner : XunitTestRunner
+    {
+        private readonly Action? _onAfterTestInvoked;
+
+        public AvaloniaTestRunner(
+            Action? onAfterTestInvoked,
+            ITest test, IMessageBus messageBus, Type testClass, object[] constructorArguments, MethodInfo testMethod,
+            object[] testMethodArguments, string skipReason,
+            IReadOnlyList<BeforeAfterTestAttribute> beforeAfterAttributes, ExceptionAggregator aggregator,
+            CancellationTokenSource cancellationTokenSource) : base(test, messageBus, testClass, constructorArguments,
+            testMethod, testMethodArguments, skipReason, beforeAfterAttributes, aggregator, cancellationTokenSource)
+        {
+            _onAfterTestInvoked = onAfterTestInvoked;
+        }
+
+        protected override Task<decimal> InvokeTestMethodAsync(ExceptionAggregator aggregator)
+        {
+            return new AvaloniaTestInvoker(_onAfterTestInvoked, Test, MessageBus, TestClass, ConstructorArguments,
+                TestMethod, TestMethodArguments, BeforeAfterAttributes, aggregator, CancellationTokenSource).RunAsync();
+        }
+    }
+
+    private class AvaloniaTestInvoker : XunitTestInvoker
+    {
+        private readonly Action? _onAfterTestInvoked;
+
+        public AvaloniaTestInvoker(
+            Action? onAfterTestInvoked,
+            ITest test, IMessageBus messageBus, Type testClass, object[] constructorArguments, MethodInfo testMethod,
+            object[] testMethodArguments, IReadOnlyList<BeforeAfterTestAttribute> beforeAfterAttributes,
+            ExceptionAggregator aggregator, CancellationTokenSource cancellationTokenSource) : base(test, messageBus,
+            testClass, constructorArguments, testMethod, testMethodArguments, beforeAfterAttributes, aggregator,
+            cancellationTokenSource)
+        {
+            _onAfterTestInvoked = onAfterTestInvoked;
+        }
+
+        protected override async Task AfterTestMethodInvokedAsync()
+        {
+            await base.AfterTestMethodInvokedAsync();
+
+            // Only here we can execute random code after the test, where exception will be properly handled by the XUnit.
+            if (_onAfterTestInvoked is not null)
+            {
+                Aggregator.Run(_onAfterTestInvoked);
+            }
+        }
+    }
+}

src/Headless/Avalonia.Headless.XUnit/AvaloniaTheoryTestCase.cs

@@ -2,7 +2,6 @@
 using System.ComponentModel;
 using System.Threading;
 using System.Threading.Tasks;
-using Avalonia.Headless.NUnit;
 using Xunit.Abstractions;
 using Xunit.Sdk;
 
@@ -25,11 +24,8 @@ internal class AvaloniaTheoryTestCase : XunitTheoryTestCase
     {
         var session = HeadlessUnitTestSession.GetOrStartForAssembly(Method.ToRuntimeMethod().DeclaringType?.Assembly);
 
-        return session.Dispatcher.InvokeOnQueueAsync(async () =>
-        {
-            var runner = new XunitTestCaseRunner(this, DisplayName, SkipReason, constructorArguments,
+        return AvaloniaTestCaseRunner
+            .RunTest(session, this, DisplayName, SkipReason, constructorArguments,
                 TestMethodArguments, messageBus, aggregator, cancellationTokenSource);
-            return await runner.RunAsync();
-        }, cancellationTokenSource.Token);
     }
 }

src/Headless/Avalonia.Headless/HeadlessUnitTestSession.cs

@@ -1,6 +1,8 @@
 using System;
+using System.Collections.Concurrent;
 using System.Diagnostics.CodeAnalysis;
 using System.Reflection;
+using System.Runtime.ExceptionServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Avalonia.Threading;
@@ -18,9 +20,11 @@ namespace Avalonia.Headless;
 /// </remarks>
 public sealed class HeadlessUnitTestSession : IDisposable
 {
-    private readonly CancellationTokenSource _cancellationToken;
+    private readonly CancellationTokenSource _cancellationTokenSource;
     private static HeadlessUnitTestSession? s_session;
     private static object s_lock = new();
+    private readonly BlockingCollection<Action> _queue;
+    private readonly Task _dispatchTask;
 
     internal const DynamicallyAccessedMemberTypes DynamicallyAccessed =
         DynamicallyAccessedMemberTypes.PublicMethods |
@@ -29,9 +33,11 @@ public sealed class HeadlessUnitTestSession : IDisposable
     
     private HeadlessUnitTestSession(Type entryPointType, Application application,
         SynchronizationContext synchronizationContext,
-        Dispatcher dispatcher, CancellationTokenSource cancellationToken)
+        Dispatcher dispatcher, CancellationTokenSource cancellationTokenSource, BlockingCollection<Action> queue, Task _dispatchTask)
     {
-        _cancellationToken = cancellationToken;
+        _cancellationTokenSource = cancellationTokenSource;
+        _queue = queue;
+        this._dispatchTask = _dispatchTask;
         EntryPointType = entryPointType;
         Dispatcher = dispatcher;
         Application = application;
@@ -43,9 +49,63 @@ public sealed class HeadlessUnitTestSession : IDisposable
     public Dispatcher Dispatcher { get; }
     internal Type EntryPointType { get; }
 
+    public Task Dispatch(Action action, CancellationToken cancellationToken)
+    {
+        return Dispatch(() => { action(); return Task.FromResult(0); }, cancellationToken);
+    }
+    
+    public Task<TResult> Dispatch<TResult>(Func<TResult> action, CancellationToken cancellationToken)
+    {
+        return Dispatch(() => Task.FromResult(action()), cancellationToken);
+    }
+
+    public Task<TResult> Dispatch<TResult>(Func<Task<TResult>> action, CancellationToken cancellationToken)
+    {
+        if (_cancellationTokenSource.IsCancellationRequested)
+        {
+            throw new ObjectDisposedException("Session was already disposed.");
+        }
+
+        var token = _cancellationTokenSource.Token;
+
+        var tcs = new TaskCompletionSource<TResult>();
+        _queue.Add(() =>
+        {
+            var cts = new CancellationTokenSource();
+            using var globalCts = token.Register(s => ((CancellationTokenSource)s!).Cancel(), cts);
+            using var localCts = cancellationToken.Register(s => ((CancellationTokenSource)s!).Cancel(), cts);
+
+            try
+            {
+                var task = action();
+                task.ContinueWith((_, s) => ((CancellationTokenSource)s!).Cancel(), cts);
+
+                if (cts.IsCancellationRequested)
+                {
+                    return;
+                }
+
+                var frame = new DispatcherFrame();
+                using var innerCts = cts.Token.Register(() => frame.Continue = false);
+                Dispatcher.PushFrame(frame);
+
+                var result = task.GetAwaiter().GetResult();
+                tcs.TrySetResult(result);
+            }
+            catch (Exception ex)
+            {
+                tcs.TrySetException(ex);
+            }
+        });
+        return tcs.Task;
+    }
+    
     public void Dispose()
     {
-        _cancellationToken.Cancel();
+        _cancellationTokenSource.Cancel();
+        _queue.CompleteAdding();
+        _dispatchTask.Wait();
+        _cancellationTokenSource.Dispose();
     }
 
     /// <summary>
@@ -73,9 +133,10 @@ public sealed class HeadlessUnitTestSession : IDisposable
     {
         var tcs = new TaskCompletionSource<HeadlessUnitTestSession>();
         var cancellationTokenSource = new CancellationTokenSource();
+        var queue = new BlockingCollection<Action>();
 
-        Thread? thread = null;
-        thread = new Thread(() =>
+        Task? task = null;
+        task = Task.Run(() =>
         {
             try
             {
@@ -93,10 +154,10 @@ public sealed class HeadlessUnitTestSession : IDisposable
                 {
                     throw new InvalidOperationException("Avalonia Headless platform has failed to initialize.");
                 }
-                
+
                 // ReSharper disable once AccessToModifiedClosure
                 tcs.SetResult(new HeadlessUnitTestSession(entryPointType, Application.Current!,
-                    SynchronizationContext.Current!, Dispatcher.UIThread, cancellationTokenSource));
+                    SynchronizationContext.Current!, Dispatcher.UIThread, cancellationTokenSource, queue, task!));
             }
             catch (Exception e)
             {
@@ -104,13 +165,23 @@ public sealed class HeadlessUnitTestSession : IDisposable
                 return;
             }
 
-            Dispatcher.UIThread.MainLoop(cancellationTokenSource.Token);
-        }) { IsBackground = true };
-        thread.Start();
+            while (!cancellationTokenSource.IsCancellationRequested)
+            {
+                try
+                {
+                    var action = queue.Take(cancellationTokenSource.Token);
+                    action();
+                }
+                catch (OperationCanceledException)
+                {
+
+                }
+            }
+        });
 
         return tcs.Task.GetAwaiter().GetResult();
     }
-    
+
     /// <summary>
     /// Creates a session from AvaloniaTestApplicationAttribute attribute or reuses any existing.
     /// If AvaloniaTestApplicationAttribute doesn't exist, empty application is used. 

tests/Avalonia.Headless.UnitTests/ThreadingTests.cs

@@ -13,6 +13,13 @@ public class ThreadingTests
         Dispatcher.UIThread.VerifyAccess();
     }
     
+    // This test should always fail, uncomment to test if it fails.
+    // [AvaloniaFact]
+    // public void Should_Fail_Test_On_Delayed_Post_When_FlushDispatcher()
+    // {
+    //     Dispatcher.UIThread.Post(() => throw new InvalidOperationException(), DispatcherPriority.Default);
+    // }
+
     [AvaloniaTheory]
     [InlineData(1)]
     [InlineData(10)]