Headless NUnit: Handle async SetUp/TearDown (#20612)

* Headless NUnit: Fix TearDown not always working

* Headless NUnit: Handle async SetUp/TearDown

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

@@ -1,6 +1,5 @@
 using System;
 using System.Collections.Generic;
-using System.Reflection;
 using System.Threading.Tasks;
 using Avalonia.Threading;
 using NUnit.Framework.Interfaces;
@@ -13,28 +12,14 @@ internal class AvaloniaTestMethodCommand : TestCommand
 {
     private readonly HeadlessUnitTestSession _session;
     private readonly TestCommand _innerCommand;
-    private readonly List<Action> _beforeTest;
-    private readonly List<Action> _afterTest;
-
-    // There are multiple problems with NUnit integration at the moment when we wrote this integration.
-    // NUnit doesn't have extensibility API for running on custom dispatcher/sync-context.
-    // See https://github.com/nunit/nunit/issues/2917 https://github.com/nunit/nunit/issues/2774
-    // To workaround that we had to replace inner TestMethodCommand with our own implementation while keeping original hierarchy of commands.
-    // Which will respect proper async/await awaiting code that works with our session and can be block-awaited to fit in NUnit.
-    // Also, we need to push BeforeTest/AfterTest callbacks to the very same session call.
-    // I hope there will be a better solution without reflection, but for now that's it.
-    private static FieldInfo s_innerCommand = typeof(DelegatingTestCommand)
-        .GetField("innerCommand", BindingFlags.Instance | BindingFlags.NonPublic)!;
-    private static FieldInfo s_beforeTest = typeof(BeforeAndAfterTestCommand)
-        .GetField("BeforeTest", BindingFlags.Instance | BindingFlags.NonPublic)!;
-    private static FieldInfo s_afterTest = typeof(BeforeAndAfterTestCommand)
-        .GetField("AfterTest", BindingFlags.Instance | BindingFlags.NonPublic)!;
+    private readonly List<Func<TestExecutionContext, Task>> _beforeTest;
+    private readonly List<Func<TestExecutionContext, Task>> _afterTest;
 
     private AvaloniaTestMethodCommand(
         HeadlessUnitTestSession session,
         TestCommand innerCommand,
-        List<Action> beforeTest,
-        List<Action> afterTest)
+        List<Func<TestExecutionContext, Task>> beforeTest,
+        List<Func<TestExecutionContext, Task>> afterTest)
         : base(innerCommand.Test)
     {
         _session = session;
@@ -45,61 +30,65 @@ internal class AvaloniaTestMethodCommand : TestCommand
 
     public static TestCommand ProcessCommand(HeadlessUnitTestSession session, TestCommand command)
     {
-        return ProcessCommand(session, command, new List<Action>(), new List<Action>());
+        return ProcessCommand(session, command, [], []);
     }
 
-    private static TestCommand ProcessCommand(HeadlessUnitTestSession session, TestCommand command, List<Action> before, List<Action> after)
+    private static TestCommand ProcessCommand(
+        HeadlessUnitTestSession session,
+        TestCommand command,
+        List<Func<TestExecutionContext, Task>> before,
+        List<Func<TestExecutionContext, Task>> after)
     {
-        if (command is BeforeAndAfterTestCommand beforeAndAfterTestCommand)
+        var beforeAndAfterTestCommand = command as BeforeAndAfterTestCommand;
+        if (beforeAndAfterTestCommand is not null)
         {
-            if (s_beforeTest.GetValue(beforeAndAfterTestCommand) is Action<TestExecutionContext> beforeTest)
-            {
-                Action<TestExecutionContext> beforeAction = c => before.Add(() => beforeTest(c));
-                s_beforeTest.SetValue(beforeAndAfterTestCommand, beforeAction);
-            }
-            if (s_afterTest.GetValue(beforeAndAfterTestCommand) is Action<TestExecutionContext> afterTest)
+            ref var beforeTest = ref beforeAndAfterTestCommand.BeforeTest();
+            if (beforeTest is not null)
             {
-                Action<TestExecutionContext> afterAction = c => after.Add(() => afterTest(c));
-                s_afterTest.SetValue(beforeAndAfterTestCommand, afterAction);
+                AddBeforeOrAfterAction(beforeTest, before);
+                beforeTest = _ => { };
             }
         }
-        
-        if (command is DelegatingTestCommand delegatingTestCommand
-            && s_innerCommand.GetValue(delegatingTestCommand) is TestCommand inner)
+
+        var delegatingTestCommand = command as DelegatingTestCommand;
+        if (delegatingTestCommand is not null)
         {
-            s_innerCommand.SetValue(delegatingTestCommand, ProcessCommand(session, inner, before, after));
+            ref var innerCommand = ref delegatingTestCommand.InnerCommand();
+            innerCommand = ProcessCommand(session, innerCommand, before, after);
         }
-        else if (command is TestMethodCommand methodCommand)
+
+        if (beforeAndAfterTestCommand is not null)
         {
-            return new AvaloniaTestMethodCommand(session, methodCommand, before, after);
+            ref var afterTest = ref beforeAndAfterTestCommand.AfterTest();
+            if (afterTest is not null)
+            {
+                AddBeforeOrAfterAction(afterTest, after);
+                afterTest = _ => { };
+            }
         }
+        
+        if (delegatingTestCommand is null && command is TestMethodCommand methodCommand)
+            return new AvaloniaTestMethodCommand(session, methodCommand, before, after);
 
         return command;
     }
 
     public override TestResult Execute(TestExecutionContext context)
     {
-        return _session.DispatchCore(() => ExecuteTestMethod(context), true, default).GetAwaiter().GetResult();
+        return _session.DispatchCore(() => ExecuteTestMethod(context), true, context.CancellationToken).GetAwaiter().GetResult();
     }
 
     // Unfortunately, NUnit has issues with custom synchronization contexts, which means we need to add some hacks to make it work.
     private async Task<TestResult> ExecuteTestMethod(TestExecutionContext context)
     {
-        _beforeTest.ForEach(a => a());
+        foreach (var beforeTest in _beforeTest)
+            await beforeTest(context);
         
         var testMethod = _innerCommand.Test.Method;
         var methodInfo = testMethod!.MethodInfo;
 
         var result = methodInfo.Invoke(context.TestObject, _innerCommand.Test.Arguments);
-        // Only Task, non generic ValueTask are supported in async context. No ValueTask<> nor F# tasks.
-        if (result is Task task)
-        {
-            await task;
-        }
-        else if (result is ValueTask valueTask)
-        {
-            await valueTask;
-        }
+        await ToTask(result);
 
         context.CurrentResult.SetResult(ResultState.Success);
 
@@ -108,10 +97,55 @@ internal class AvaloniaTestMethodCommand : TestCommand
 
         if (context.ExecutionStatus != TestExecutionStatus.AbortRequested)
         {
-            _afterTest.ForEach(a => a());
+            foreach (var afterTest in _afterTest)
+                await afterTest(context);
+
             Dispatcher.UIThread.RunJobs();
         }
         
         return context.CurrentResult;
     }
+
+    private static void AddBeforeOrAfterAction(Action<TestExecutionContext> action, List<Func<TestExecutionContext, Task>> targets)
+    {
+        // We need to extract the SetUp and TearDown methods to run them asynchronously on Avalonia's synchronization context.
+        if (action.Target is SetUpTearDownItem setUpTearDownItem)
+        {
+            var methods = action.Method.Name switch
+            {
+                nameof(SetUpTearDownItem.RunSetUp) => setUpTearDownItem.SetUpMethods(),
+                nameof(SetUpTearDownItem.RunTearDown) => setUpTearDownItem.TearDownMethods(),
+                _ => null
+            };
+
+            if (methods is not null)
+            {
+                foreach (var method in methods)
+                {
+                    targets.Add(context =>
+                    {
+                        var result = method.Invoke(method.IsStatic ? null : context.TestObject, null);
+                        return ToTask(result);
+                    });
+                }
+
+                return;
+            }
+        }
+
+        targets.Add(context =>
+        {
+            action(context);
+            return Task.CompletedTask;
+        });
+    }
+
+    private static Task ToTask(object? result)
+        // Only Task, non generic ValueTask are supported in async context. No ValueTask<> nor F# tasks.
+        => result switch
+        {
+            Task task => task,
+            ValueTask valueTask => valueTask.AsTask(),
+            _ => Task.CompletedTask
+        };
 }

src/Headless/Avalonia.Headless.NUnit/NUnitReflectionHelper.cs

@@ -0,0 +1,86 @@
+using System;
+using System.Collections.Generic;
+using System.Runtime.CompilerServices;
+using NUnit.Framework.Interfaces;
+using NUnit.Framework.Internal;
+using NUnit.Framework.Internal.Commands;
+
+namespace Avalonia.Headless.NUnit;
+
+/// <summary>
+/// 2023-05-10, original comment from Max about NUnit 3:
+/// There are multiple problems with NUnit integration at the moment when we wrote this integration.
+/// NUnit doesn't have extensibility API for running on custom dispatcher/sync-context.
+/// See https://github.com/nunit/nunit/issues/2917 https://github.com/nunit/nunit/issues/2774
+/// To workaround that we had to replace inner TestMethodCommand with our own implementation while keeping original hierarchy of commands.
+/// Which will respect proper async/await awaiting code that works with our session and can be block-awaited to fit in NUnit.
+/// Also, we need to push BeforeTest/AfterTest callbacks to the very same session call.
+/// I hope there will be a better solution without reflection, but for now that's it.
+///
+/// 2026-02-04: the situation hasn't changed at all with NUnit 4.
+/// </summary>
+internal static class NUnitReflectionHelper
+{
+    [UnsafeAccessor(UnsafeAccessorKind.Field, Name = ReflectionDelegatingTestCommand.InnerCommandFieldName)]
+    private static extern ref TestCommand DelegatingTestCommand_InnerCommand(DelegatingTestCommand instance);
+
+    [UnsafeAccessor(UnsafeAccessorKind.Field, Name = ReflectionBeforeAndAfterTestCommand.BeforeTestFieldName)]
+    private static extern ref Action<TestExecutionContext>? BeforeAndAfterTestCommand_BeforeTest(BeforeAndAfterTestCommand instance);
+
+    [UnsafeAccessor(UnsafeAccessorKind.Field, Name = ReflectionBeforeAndAfterTestCommand.AfterTestFieldName)]
+    private static extern ref Action<TestExecutionContext>? BeforeAndAfterTestCommand_AfterTest(BeforeAndAfterTestCommand instance);
+
+    [UnsafeAccessor(UnsafeAccessorKind.Field, Name = "_setUpMethods")]
+    private static extern ref IList<IMethodInfo> SetUpTearDownItem_SetUpMethods(SetUpTearDownItem instance);
+
+    [UnsafeAccessor(UnsafeAccessorKind.Field, Name = "_tearDownMethods")]
+    private static extern ref IList<IMethodInfo> SetUpTearDownItem_TearDownMethods(SetUpTearDownItem instance);
+
+    extension(DelegatingTestCommand instance)
+    {
+        public ref TestCommand InnerCommand()
+            => ref DelegatingTestCommand_InnerCommand(instance);
+    }
+
+    extension(BeforeAndAfterTestCommand instance)
+    {
+        public ref Action<TestExecutionContext>? BeforeTest()
+            => ref BeforeAndAfterTestCommand_BeforeTest(instance);
+
+        public ref Action<TestExecutionContext>? AfterTest()
+            => ref BeforeAndAfterTestCommand_AfterTest(instance);
+    }
+
+    extension(SetUpTearDownItem instance)
+    {
+        public ref IList<IMethodInfo> SetUpMethods()
+            => ref SetUpTearDownItem_SetUpMethods(instance);
+
+        public ref IList<IMethodInfo> TearDownMethods()
+            => ref SetUpTearDownItem_TearDownMethods(instance);
+    }
+
+    private sealed class ReflectionDelegatingTestCommand : DelegatingTestCommand
+    {
+        public ReflectionDelegatingTestCommand(TestCommand innerCommand)
+            : base(innerCommand)
+        {
+        }
+
+        public const string InnerCommandFieldName = nameof(innerCommand);
+
+        public override TestResult Execute(TestExecutionContext context)
+            => throw new NotSupportedException("Reflection-only type, this method should never be called");
+    }
+
+    private sealed class ReflectionBeforeAndAfterTestCommand : BeforeAndAfterTestCommand
+    {
+        public ReflectionBeforeAndAfterTestCommand(TestCommand innerCommand)
+            : base(innerCommand)
+        {
+        }
+
+        public const string BeforeTestFieldName = nameof(BeforeTest);
+        public const string AfterTestFieldName = nameof(AfterTest);
+    }
+}

tests/Avalonia.Headless.UnitTests/AsyncSetupTests.cs

@@ -0,0 +1,32 @@
+#if NUNIT
+
+using System.Threading.Tasks;
+
+namespace Avalonia.Headless.UnitTests;
+
+public class AsyncSetupTests
+{
+    private static int s_instanceCount;
+
+    [SetUp]
+    public async Task SetUp()
+    {
+        await Task.Delay(100);
+        ++s_instanceCount;
+    }
+
+    [AvaloniaTest, TestCase(1), TestCase(2)]
+    public void Async_Setup_TearDown_Should_Work(int index)
+    {
+        AssertHelper.Equal(1, s_instanceCount);
+    }
+
+    [TearDown]
+    public async Task TearDown()
+    {
+        await Task.Delay(100);
+        --s_instanceCount;
+    }
+}
+
+#endif

tests/Avalonia.Headless.UnitTests/SetupTests.cs

@@ -0,0 +1,45 @@
+using System;
+
+namespace Avalonia.Headless.UnitTests;
+
+public class SetupTests
+#if XUNIT
+    : IDisposable
+#endif
+{
+    private static int s_instanceCount;
+
+#if NUNIT
+    [SetUp]
+    public void SetUp()
+#elif XUNIT
+    public SetupTests()
+#endif
+    {
+        ++s_instanceCount;
+    }
+
+#if NUNIT
+    [AvaloniaTest, TestCase(1), TestCase(2)]
+#elif XUNIT
+    [AvaloniaTheory, InlineData(1), InlineData(2)]
+    [System.Diagnostics.CodeAnalysis.UnconditionalSuppressMessage(
+        "Usage",
+        "xUnit1026:Theory methods should use all of their parameters",
+        Justification = "Used to run the test several times")]
+#endif
+    public void Setup_TearDown_Should_Work(int index)
+    {
+        AssertHelper.Equal(1, s_instanceCount);
+    }
+
+#if NUNIT
+    [TearDown]
+    public void TearDown()
+#elif XUNIT
+    public void Dispose()
+#endif
+    {
+        --s_instanceCount;
+    }
+}