Merge remote-tracking branch 'origin/nunit-headless-intergration' into unify-test-mocks

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/AvaloniaTestAssemblyRunner.cs

@@ -0,0 +1,126 @@
+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 AvaloniaTestAssemblyRunner : XunitTestAssemblyRunner
+{
+    private HeadlessUnitTestSession? _session;
+
+    public AvaloniaTestAssemblyRunner(ITestAssembly testAssembly, IEnumerable<IXunitTestCase> testCases,
+        IMessageSink diagnosticMessageSink, IMessageSink executionMessageSink,
+        ITestFrameworkExecutionOptions executionOptions) : base(testAssembly, testCases, diagnosticMessageSink,
+        executionMessageSink, executionOptions)
+    {
+    }
+
+    protected override void SetupSyncContext(int maxParallelThreads)
+    {
+        _session = HeadlessUnitTestSession.GetOrStartForAssembly(
+            Assembly.Load(new AssemblyName(TestAssembly.Assembly.Name)));
+        base.SetupSyncContext(1);
+    }
+
+    public override void Dispose()
+    {
+        _session?.Dispose();
+        base.Dispose();
+    }
+
+    protected override Task<RunSummary> RunTestCollectionAsync(
+        IMessageBus messageBus,
+        ITestCollection testCollection,
+        IEnumerable<IXunitTestCase> testCases,
+        CancellationTokenSource cancellationTokenSource)
+    {
+        return new AvaloniaTestCollectionRunner(_session!, testCollection, testCases, DiagnosticMessageSink, messageBus,
+            TestCaseOrderer, new ExceptionAggregator(Aggregator), cancellationTokenSource).RunAsync();
+    }
+
+    private class AvaloniaTestCollectionRunner : XunitTestCollectionRunner
+    {
+        private readonly HeadlessUnitTestSession _session;
+
+        public AvaloniaTestCollectionRunner(HeadlessUnitTestSession session,
+            ITestCollection testCollection, IEnumerable<IXunitTestCase> testCases,
+            IMessageSink diagnosticMessageSink, IMessageBus messageBus, ITestCaseOrderer testCaseOrderer,
+            ExceptionAggregator aggregator, CancellationTokenSource cancellationTokenSource) : base(testCollection,
+            testCases, diagnosticMessageSink, messageBus, testCaseOrderer, aggregator, cancellationTokenSource)
+        {
+            _session = session;
+        }
+
+        protected override Task<RunSummary> RunTestClassAsync(
+            ITestClass testClass,
+            IReflectionTypeInfo @class,
+            IEnumerable<IXunitTestCase> testCases)
+        {
+            return new AvaloniaTestClassRunner(_session, testClass, @class, testCases, DiagnosticMessageSink, MessageBus,
+                TestCaseOrderer, new ExceptionAggregator(Aggregator), CancellationTokenSource,
+                CollectionFixtureMappings).RunAsync();
+        }
+    }
+
+    private class AvaloniaTestClassRunner : XunitTestClassRunner
+    {
+        private readonly HeadlessUnitTestSession _session;
+
+        public AvaloniaTestClassRunner(HeadlessUnitTestSession session, ITestClass testClass,
+            IReflectionTypeInfo @class,
+            IEnumerable<IXunitTestCase> testCases, IMessageSink diagnosticMessageSink, IMessageBus messageBus,
+            ITestCaseOrderer testCaseOrderer, ExceptionAggregator aggregator,
+            CancellationTokenSource cancellationTokenSource, IDictionary<Type, object> collectionFixtureMappings) :
+            base(testClass, @class, testCases, diagnosticMessageSink, messageBus, testCaseOrderer, aggregator,
+                cancellationTokenSource, collectionFixtureMappings)
+        {
+            _session = session;
+        }
+
+        protected override Task<RunSummary> RunTestMethodAsync(
+            ITestMethod testMethod,
+            IReflectionMethodInfo method,
+            IEnumerable<IXunitTestCase> testCases,
+            object[] constructorArguments)
+        {
+            return new AvaloniaTestMethodRunner(_session, testMethod, Class, method, testCases, DiagnosticMessageSink,
+                MessageBus, new ExceptionAggregator(Aggregator), CancellationTokenSource,
+                constructorArguments).RunAsync();
+        }
+    }
+
+    private class AvaloniaTestMethodRunner : XunitTestMethodRunner
+    {
+        private readonly HeadlessUnitTestSession _session;
+        private readonly IMessageBus _messageBus;
+        private readonly ExceptionAggregator _aggregator;
+        private readonly CancellationTokenSource _cancellationTokenSource;
+        private readonly object[] _constructorArguments;
+
+        public AvaloniaTestMethodRunner(HeadlessUnitTestSession session, ITestMethod testMethod,
+            IReflectionTypeInfo @class,
+            IReflectionMethodInfo method, IEnumerable<IXunitTestCase> testCases, IMessageSink diagnosticMessageSink,
+            IMessageBus messageBus, ExceptionAggregator aggregator, CancellationTokenSource cancellationTokenSource,
+            object[] constructorArguments) : base(testMethod, @class, method, testCases, diagnosticMessageSink,
+            messageBus, aggregator, cancellationTokenSource, constructorArguments)
+        {
+            _session = session;
+            _messageBus = messageBus;
+            _aggregator = aggregator;
+            _cancellationTokenSource = cancellationTokenSource;
+            _constructorArguments = constructorArguments;
+        }
+
+        protected override Task<RunSummary> RunTestCaseAsync(IXunitTestCase testCase)
+        {
+            return AvaloniaTestCaseRunner.RunTest(_session, testCase, testCase.DisplayName, testCase.SkipReason,
+                _constructorArguments, testCase.TestMethodArguments, _messageBus, _aggregator,
+                _cancellationTokenSource);
+        }
+    }
+}

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 runSummary = AvaloniaTestCaseRunner
-        {
+            .RunTest(session, this, DisplayName, SkipReason, constructorArguments,
-            var runner = new XunitTestCaseRunner(this, DisplayName, SkipReason, constructorArguments,
+                TestMethodArguments, messageBus, aggregator, cancellationTokenSource)
-                TestMethodArguments, messageBus, aggregator, cancellationTokenSource);
+            .GetAwaiter().GetResult();
-            return await runner.RunAsync();
-        }, cancellationTokenSource.Token);
 
         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/AvaloniaTestFramework.cs

@@ -27,7 +27,7 @@ internal class AvaloniaTestFramework : XunitTestFramework
             IMessageSink executionMessageSink,
             ITestFrameworkExecutionOptions executionOptions)
         {
-            using (var assemblyRunner = new AvaloniaTestRunner(
+            using (var assemblyRunner = new AvaloniaTestAssemblyRunner(
                        TestAssembly, testCases, DiagnosticMessageSink, executionMessageSink,
                        executionOptions)) await assemblyRunner.RunAsync();
         }

src/Headless/Avalonia.Headless.XUnit/AvaloniaTestRunner.cs

@@ -1,34 +0,0 @@
-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 AvaloniaTestRunner : XunitTestAssemblyRunner
-{
-    private HeadlessUnitTestSession? _session;
-
-    public AvaloniaTestRunner(ITestAssembly testAssembly, IEnumerable<IXunitTestCase> testCases,
-        IMessageSink diagnosticMessageSink, IMessageSink executionMessageSink,
-        ITestFrameworkExecutionOptions executionOptions) : base(testAssembly, testCases, diagnosticMessageSink,
-        executionMessageSink, executionOptions)
-    {
-    }
-
-    protected override void SetupSyncContext(int maxParallelThreads)
-    {
-        _session = HeadlessUnitTestSession.GetOrStartForAssembly(Assembly.Load(new AssemblyName(TestAssembly.Assembly.Name)));
-        SynchronizationContext.SetSynchronizationContext(_session.SynchronizationContext);
-    }
-
-    public override void Dispose()
-    {
-        _session?.Dispose();
-        base.Dispose();
-    }
-}

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 () =>
+        return AvaloniaTestCaseRunner
-        {
+            .RunTest(session, this, DisplayName, SkipReason, constructorArguments,
-            var runner = new XunitTestCaseRunner(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,34 +20,84 @@ 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 |
         DynamicallyAccessedMemberTypes.NonPublicMethods |
         DynamicallyAccessedMemberTypes.PublicParameterlessConstructor;
     
-    private HeadlessUnitTestSession(Type entryPointType, Application application,
+    private HeadlessUnitTestSession(Type entryPointType, CancellationTokenSource cancellationTokenSource, BlockingCollection<Action> queue, Task _dispatchTask)
-        SynchronizationContext synchronizationContext,
-        Dispatcher dispatcher, CancellationTokenSource cancellationToken)
     {
-        _cancellationToken = cancellationToken;
+        _cancellationTokenSource = cancellationTokenSource;
+        _queue = queue;
+        this._dispatchTask = _dispatchTask;
         EntryPointType = entryPointType;
-        Dispatcher = dispatcher;
-        Application = application;
-        SynchronizationContext = synchronizationContext;
     }
-    
+
-    public Application Application { get; }
-    public SynchronizationContext SynchronizationContext { get; }
-    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.UIThread.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 +125,10 @@ public sealed class HeadlessUnitTestSession : IDisposable
     {
         var tcs = new TaskCompletionSource<HeadlessUnitTestSession>();
         var cancellationTokenSource = new CancellationTokenSource();
+        var queue = new BlockingCollection<Action>();
 
-        Thread? thread = null;
+        Task? task = null;
-        thread = new Thread(() =>
+        task = Task.Run(() =>
         {
             try
             {
@@ -93,10 +146,9 @@ 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!,
+                tcs.SetResult(new HeadlessUnitTestSession(entryPointType, cancellationTokenSource, queue, task!));
-                    SynchronizationContext.Current!, Dispatcher.UIThread, cancellationTokenSource));
             }
             catch (Exception e)
             {
@@ -104,13 +156,23 @@ public sealed class HeadlessUnitTestSession : IDisposable
                 return;
             }
 
-            Dispatcher.UIThread.MainLoop(cancellationTokenSource.Token);
+            while (!cancellationTokenSource.IsCancellationRequested)
-        }) { IsBackground = true };
+            {
-        thread.Start();
+                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. 
@@ -128,7 +190,8 @@ public sealed class HeadlessUnitTestSession : IDisposable
 
             if (s_session is not null)
             {
-                if (appBuilderEntryPointType != s_session.EntryPointType)
+                var hasNoAttribute = appBuilderEntryPointType == null && s_session.EntryPointType == typeof(Application);
+                if (!hasNoAttribute && appBuilderEntryPointType != s_session.EntryPointType)
                 {
                     // Avalonia doesn't support multiple Application instances. At least at the moment.
                     throw new System.InvalidOperationException(

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)]