From cfcd3caac14a04de9eee0258e126896bab91cfa6 Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:11:15 +0100 Subject: [PATCH 1/4] perf(engine): skip absent hook pipelines --- src/TUnit.Engine/TestExecutor.cs | 133 +++++++++++++++++++++++++------ 1 file changed, 109 insertions(+), 24 deletions(-) diff --git a/src/TUnit.Engine/TestExecutor.cs b/src/TUnit.Engine/TestExecutor.cs index 9f922e80d3..968a3badb3 100644 --- a/src/TUnit.Engine/TestExecutor.cs +++ b/src/TUnit.Engine/TestExecutor.cs @@ -1,9 +1,11 @@ +using System.Collections.Concurrent; using System.Diagnostics.CodeAnalysis; using System.Reflection; using System.Runtime.ExceptionServices; using TUnit.Core; using TUnit.Core.Enums; using TUnit.Core.Exceptions; +using TUnit.Core.Hooks; using TUnit.Core.Interfaces; using TUnit.Core.Services; using TUnit.Engine.Helpers; @@ -20,6 +22,9 @@ namespace TUnit.Engine; /// internal class TestExecutor { + private static readonly ConcurrentDictionary ClassHookPresenceCache = new(); + private static readonly ConcurrentDictionary TestHookPresenceCache = new(); + private readonly HookExecutor _hookExecutor; private readonly TestLifecycleCoordinator _lifecycleCoordinator; private readonly BeforeHookTaskCache _beforeHookTaskCache; @@ -59,6 +64,11 @@ public TestExecutor( /// public async ValueTask EnsureTestSessionHooksExecutedAsync(CancellationToken cancellationToken) { + if (!HasTestSessionHooks()) + { + return; + } + // Get or create and cache Before hooks - these run only once await _beforeHookTaskCache.GetOrCreateBeforeTestSessionTask( _beforeTestSessionHookFactory, @@ -82,9 +92,12 @@ public async ValueTask EnsureClassAndAssemblyHooksExecutedAsync(AbstractExecutab { var testClass = test.Metadata.TestClassType; - await _beforeHookTaskCache.GetOrCreateBeforeTestSessionTask( - _beforeTestSessionHookFactory, - cancellationToken).ConfigureAwait(false); + if (HasTestSessionHooks()) + { + await _beforeHookTaskCache.GetOrCreateBeforeTestSessionTask( + _beforeTestSessionHookFactory, + cancellationToken).ConfigureAwait(false); + } // Flow AsyncLocals captured by BeforeTestSession into the BeforeAssembly hook, and likewise // BeforeAssembly into BeforeClass. This mirrors the RestoreExecutionContext chain in @@ -92,14 +105,20 @@ await _beforeHookTaskCache.GetOrCreateBeforeTestSessionTask( // it re-applies the same captured contexts). test.Context.ClassContext.AssemblyContext.TestSessionContext.RestoreExecutionContext(); - await _beforeHookTaskCache.GetOrCreateBeforeAssemblyTask( - testClass.Assembly, - _beforeAssemblyHookFactory, - cancellationToken).ConfigureAwait(false); + if (HasAssemblyHooks(testClass.Assembly)) + { + await _beforeHookTaskCache.GetOrCreateBeforeAssemblyTask( + testClass.Assembly, + _beforeAssemblyHookFactory, + cancellationToken).ConfigureAwait(false); + } test.Context.ClassContext.AssemblyContext.RestoreExecutionContext(); - await _beforeHookTaskCache.GetOrCreateBeforeClassTask(testClass, _hookExecutor, cancellationToken).ConfigureAwait(false); + if (HasClassHooks(testClass)) + { + await _beforeHookTaskCache.GetOrCreateBeforeClassTask(testClass, _hookExecutor, cancellationToken).ConfigureAwait(false); + } // Note: the caller (TestCoordinator) restores ClassContext.RestoreExecutionContext() right // before constructing the instance so AsyncLocals captured by BeforeAssembly/BeforeClass flow @@ -117,13 +136,20 @@ public async ValueTask ExecuteAsync(AbstractExecutableTest executableTest, TestI var testClass = executableTest.Metadata.TestClassType; var testAssembly = testClass.Assembly; + var hasSessionHooks = HasTestSessionHooks(); + var hasAssemblyHooks = HasAssemblyHooks(testAssembly); + var hasClassHooks = HasClassHooks(testClass); + var hasTestHooks = HasTestHooks(testClass); Exception? capturedException = null; Exception? hookException = null; try { - await EnsureTestSessionHooksExecutedAsync(cancellationToken).ConfigureAwait(false); + if (hasSessionHooks) + { + await EnsureTestSessionHooksExecutedAsync(cancellationToken).ConfigureAwait(false); + } await _eventReceiverOrchestrator.InvokeFirstTestInSessionEventReceiversAsync( executableTest.Context, @@ -132,16 +158,19 @@ await _eventReceiverOrchestrator.InvokeFirstTestInSessionEventReceiversAsync( executableTest.Context.ClassContext.AssemblyContext.TestSessionContext.RestoreExecutionContext(); - await _beforeHookTaskCache.GetOrCreateBeforeAssemblyTask( - testAssembly, - _beforeAssemblyHookFactory, - cancellationToken).ConfigureAwait(false); - - // Register After Assembly hook to run on cancellation (guarantees cleanup) - _afterHookPairTracker.RegisterAfterAssemblyHook( - testAssembly, - cancellationToken, - (assembly) => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, CancellationToken.None)); + if (hasAssemblyHooks) + { + await _beforeHookTaskCache.GetOrCreateBeforeAssemblyTask( + testAssembly, + _beforeAssemblyHookFactory, + cancellationToken).ConfigureAwait(false); + + // Register After Assembly hook to run on cancellation (guarantees cleanup) + _afterHookPairTracker.RegisterAfterAssemblyHook( + testAssembly, + cancellationToken, + (assembly) => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, CancellationToken.None)); + } await _eventReceiverOrchestrator.InvokeFirstTestInAssemblyEventReceiversAsync( executableTest.Context, @@ -150,10 +179,13 @@ await _eventReceiverOrchestrator.InvokeFirstTestInAssemblyEventReceiversAsync( executableTest.Context.ClassContext.AssemblyContext.RestoreExecutionContext(); - await _beforeHookTaskCache.GetOrCreateBeforeClassTask(testClass, _hookExecutor, cancellationToken).ConfigureAwait(false); + if (hasClassHooks) + { + await _beforeHookTaskCache.GetOrCreateBeforeClassTask(testClass, _hookExecutor, cancellationToken).ConfigureAwait(false); - // Register After Class hook to run on cancellation (guarantees cleanup) - _afterHookPairTracker.RegisterAfterClassHook(testClass, _hookExecutor, cancellationToken); + // Register After Class hook to run on cancellation (guarantees cleanup) + _afterHookPairTracker.RegisterAfterClassHook(testClass, _hookExecutor, cancellationToken); + } await _eventReceiverOrchestrator.InvokeFirstTestInClassEventReceiversAsync( executableTest.Context, @@ -224,7 +256,10 @@ await _eventReceiverOrchestrator.InvokeFirstTestInClassEventReceiversAsync( executableTest.Context.RestoreExecutionContext(); - await _hookExecutor.ExecuteBeforeTestHooksAsync(executableTest, cancellationToken).ConfigureAwait(false); + if (hasTestHooks) + { + await _hookExecutor.ExecuteBeforeTestHooksAsync(executableTest, cancellationToken).ConfigureAwait(false); + } // Late stage test start receivers run after instance-level hooks (default behavior) await _eventReceiverOrchestrator.InvokeTestStartEventReceiversAsync(executableTest.Context, cancellationToken, EventReceiverStage.Late).ConfigureAwait(false); @@ -343,7 +378,9 @@ await ExecuteTestAsync( // Early stage test end receivers run before instance-level hooks var earlyStageExceptions = await _eventReceiverOrchestrator.InvokeTestEndEventReceiversAsync(executableTest.Context, CancellationToken.None, EventReceiverStage.Early).ConfigureAwait(false); - var hookExceptions = await _hookExecutor.ExecuteAfterTestHooksAsync(executableTest, CancellationToken.None).ConfigureAwait(false); + var hookExceptions = hasTestHooks + ? await _hookExecutor.ExecuteAfterTestHooksAsync(executableTest, CancellationToken.None).ConfigureAwait(false) + : []; // Late stage test end receivers run after instance-level hooks (default behavior) var lateStageExceptions = await _eventReceiverOrchestrator.InvokeTestEndEventReceiversAsync(executableTest.Context, CancellationToken.None, EventReceiverStage.Late).ConfigureAwait(false); @@ -405,6 +442,54 @@ await ExecuteTestAsync( } } + private static bool HasTestSessionHooks() + => !Sources.BeforeTestSessionHooks.IsEmpty || !Sources.AfterTestSessionHooks.IsEmpty; + + private static bool HasAssemblyHooks(Assembly assembly) + => !Sources.BeforeEveryAssemblyHooks.IsEmpty || + !Sources.AfterEveryAssemblyHooks.IsEmpty || + Sources.BeforeAssemblyHooks.ContainsKey(assembly) || + Sources.AfterAssemblyHooks.ContainsKey(assembly); + + private static bool HasClassHooks(Type testClass) + => !Sources.BeforeEveryClassHooks.IsEmpty || + !Sources.AfterEveryClassHooks.IsEmpty || + ClassHookPresenceCache.GetOrAdd(testClass, static type => + HasHooksInHierarchy(type, Sources.BeforeClassHooks, Sources.AfterClassHooks)); + + private static bool HasTestHooks(Type testClass) + => !Sources.BeforeEveryTestHooks.IsEmpty || + !Sources.AfterEveryTestHooks.IsEmpty || + TestHookPresenceCache.GetOrAdd(testClass, static type => + HasHooksInHierarchy(type, Sources.BeforeTestHooks, Sources.AfterTestHooks)); + + private static bool HasHooksInHierarchy( + Type type, + ConcurrentDictionary>> beforeHooks, + ConcurrentDictionary>> afterHooks) + where TBeforeHook : HookMethod + where TAfterHook : HookMethod + { + for (var current = type; current is not null; current = current.BaseType) + { + if (beforeHooks.ContainsKey(current) || afterHooks.ContainsKey(current)) + { + return true; + } + + if (current is { IsGenericType: true, IsGenericTypeDefinition: false }) + { + var genericDefinition = current.GetGenericTypeDefinition(); + if (beforeHooks.ContainsKey(genericDefinition) || afterHooks.ContainsKey(genericDefinition)) + { + return true; + } + } + } + + return false; + } + #if NET internal static void FinishTestActivity(AbstractExecutableTest executableTest) { From 8684d1fc45172a4c4afe2768f8a743ef0d27605e Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:47:18 +0100 Subject: [PATCH 2/4] fix(engine): preserve hookless lifecycle spans --- .../Services/AfterHookPairTracker.cs | 19 ++- src/TUnit.Engine/Services/HookExecutor.cs | 158 +++++++++++++----- src/TUnit.Engine/TestExecutor.cs | 139 +++++++++++++-- .../SessionActivityLifecycleTests.cs | 55 +++++- 4 files changed, 300 insertions(+), 71 deletions(-) diff --git a/src/TUnit.Engine/Services/AfterHookPairTracker.cs b/src/TUnit.Engine/Services/AfterHookPairTracker.cs index a7d9b5ae71..69ee5ece79 100644 --- a/src/TUnit.Engine/Services/AfterHookPairTracker.cs +++ b/src/TUnit.Engine/Services/AfterHookPairTracker.cs @@ -5,6 +5,10 @@ namespace TUnit.Engine.Services; +internal delegate ValueTask> AfterClassExecutor( + [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] + Type testClass); + /// /// Responsible for ensuring After hooks run even when tests are cancelled. /// When a Before hook completes, this tracker registers the corresponding After hook @@ -109,8 +113,8 @@ public void RegisterAfterAssemblyHook( public void RegisterAfterClassHook( [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] Type testClass, - HookExecutor hookExecutor, - CancellationToken sessionCancellationToken) + CancellationToken sessionCancellationToken, + AfterClassExecutor afterHookExecutor) { if (!_classHookRegistered.Add(testClass)) { @@ -119,9 +123,9 @@ public void RegisterAfterClassHook( var registration = sessionCancellationToken.Register(static state => { - var (pairTracker, testClass, hookExecutor) = ((AfterHookPairTracker, Type, HookExecutor))state!; - _ = pairTracker.GetOrCreateAfterClassTask(testClass, hookExecutor, CancellationToken.None); - }, (this, testClass, hookExecutor)); + var (pairTracker, testClass, afterHookExecutor) = ((AfterHookPairTracker, Type, AfterClassExecutor))state!; + _ = pairTracker.GetOrCreateAfterClassTask(testClass, afterHookExecutor); + }, (this, testClass, afterHookExecutor)); _registrations.Add(registration); } @@ -176,8 +180,7 @@ public ValueTask> GetOrCreateAfterAssemblyTask(Assembly assembly public ValueTask> GetOrCreateAfterClassTask( [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] Type testClass, - HookExecutor hookExecutor, - CancellationToken cancellationToken) + AfterClassExecutor taskFactory) { // Lock-free fast path avoids allocating a closure on the common cache-hit case. if (_afterClassTasks.TryGetValue(testClass, out var existingTask)) @@ -190,7 +193,7 @@ public ValueTask> GetOrCreateAfterClassTask( // behind a shared lock. var task = _afterClassTasks.GetOrAdd( testClass, - _ => hookExecutor.ExecuteAfterClassHooksAsync(testClass, cancellationToken).AsTask()); + _ => taskFactory(testClass).AsTask()); return new ValueTask>(task); } diff --git a/src/TUnit.Engine/Services/HookExecutor.cs b/src/TUnit.Engine/Services/HookExecutor.cs index ea0b86a808..5e6167290c 100644 --- a/src/TUnit.Engine/Services/HookExecutor.cs +++ b/src/TUnit.Engine/Services/HookExecutor.cs @@ -162,17 +162,7 @@ public async ValueTask ExecuteBeforeAssemblyHooksAsync(Assembly assembly, Cancel var assemblyContext = _contextProvider.GetOrCreateAssemblyContext(assembly); #if NET - if (TUnitActivitySource.LifecycleSource.HasListeners()) - { - var sessionActivity = _contextProvider.TestSessionContext.Activity; - assemblyContext.Activity = TUnitActivitySource.StartLifecycleActivity( - TUnitActivitySource.SpanTestAssembly, - System.Diagnostics.ActivityKind.Internal, - sessionActivity?.Context ?? default, - [ - new(TUnitActivitySource.TagAssemblyName, assembly.GetName().Name) - ]); - } + TryStartAssemblyActivity(assembly); #endif // Execute BeforeEvery(Assembly) hooks first (global hooks run before specific hooks) @@ -291,25 +281,67 @@ public async ValueTask> ExecuteAfterAssemblyHooksAsync(Assembly } #if NET - private void FinishAssemblyActivity(Assembly assembly, bool hasErrors) + internal void TryStartAssemblyActivity(Assembly assembly) { - var assemblyContext = _contextProvider.GetOrCreateAssemblyContext(assembly); - var activity = assemblyContext.Activity; - - if (activity is null) + if (!TUnitActivitySource.LifecycleSource.HasListeners()) { return; } - activity.SetTag(TUnitActivitySource.TagTestCount, assemblyContext.TestCount); + var assemblyContext = _contextProvider.GetOrCreateAssemblyContext(assembly); + if (assemblyContext.Activity is not null) + { + return; + } - if (hasErrors) + lock (assemblyContext) { - activity.SetStatus(System.Diagnostics.ActivityStatusCode.Error); + if (assemblyContext.Activity is not null) + { + return; + } + + var sessionActivity = _contextProvider.TestSessionContext.Activity; + assemblyContext.Activity = TUnitActivitySource.StartLifecycleActivity( + TUnitActivitySource.SpanTestAssembly, + System.Diagnostics.ActivityKind.Internal, + sessionActivity?.Context ?? default, + [ + new(TUnitActivitySource.TagAssemblyName, assembly.GetName().Name) + ]); } + } - TUnitActivitySource.StopActivity(activity); - assemblyContext.Activity = null; + internal bool HasAssemblyActivity(Assembly assembly) + => _contextProvider.GetOrCreateAssemblyContext(assembly).Activity is not null; + + internal ValueTask> FinishAssemblyActivityAsync(Assembly assembly) + { + FinishAssemblyActivity(assembly, hasErrors: false); + return new ValueTask>([]); + } + + private void FinishAssemblyActivity(Assembly assembly, bool hasErrors) + { + var assemblyContext = _contextProvider.GetOrCreateAssemblyContext(assembly); + lock (assemblyContext) + { + var activity = assemblyContext.Activity; + if (activity is null) + { + return; + } + + activity.SetTag(TUnitActivitySource.TagTestCount, assemblyContext.TestCount); + + if (hasErrors) + { + activity.SetStatus(System.Diagnostics.ActivityStatusCode.Error); + } + + TUnitActivitySource.StopActivity(activity); + assemblyContext.Activity = null; + } } #endif @@ -320,18 +352,7 @@ public async ValueTask ExecuteBeforeClassHooksAsync( var classContext = _contextProvider.GetOrCreateClassContext(testClass); #if NET - if (TUnitActivitySource.LifecycleSource.HasListeners()) - { - var assemblyActivity = classContext.AssemblyContext.Activity; - classContext.Activity = TUnitActivitySource.StartLifecycleActivity( - TUnitActivitySource.SpanTestSuite, - System.Diagnostics.ActivityKind.Internal, - assemblyActivity?.Context ?? default, - [ - new(TUnitActivitySource.TagTestSuiteName, testClass.Name), - new(TUnitActivitySource.TagClassNamespace, testClass.Namespace) - ]); - } + TryStartClassActivity(testClass); #endif // Execute BeforeEvery(Class) hooks first (global hooks run before specific hooks) @@ -452,27 +473,76 @@ public async ValueTask> ExecuteAfterClassHooksAsync( } #if NET - private void FinishClassActivity( + internal void TryStartClassActivity( [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] - Type testClass, bool hasErrors) + Type testClass) { - var classContext = _contextProvider.GetOrCreateClassContext(testClass); - var activity = classContext.Activity; - - if (activity is null) + if (!TUnitActivitySource.LifecycleSource.HasListeners()) { return; } - activity.SetTag(TUnitActivitySource.TagTestCount, classContext.TestCount); + var classContext = _contextProvider.GetOrCreateClassContext(testClass); + if (classContext.Activity is not null) + { + return; + } - if (hasErrors) + lock (classContext) { - activity.SetStatus(System.Diagnostics.ActivityStatusCode.Error); + if (classContext.Activity is not null) + { + return; + } + + var assemblyActivity = classContext.AssemblyContext.Activity; + classContext.Activity = TUnitActivitySource.StartLifecycleActivity( + TUnitActivitySource.SpanTestSuite, + System.Diagnostics.ActivityKind.Internal, + assemblyActivity?.Context ?? default, + [ + new(TUnitActivitySource.TagTestSuiteName, testClass.Name), + new(TUnitActivitySource.TagClassNamespace, testClass.Namespace) + ]); } + } - TUnitActivitySource.StopActivity(activity); - classContext.Activity = null; + internal bool HasClassActivity( + [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] + Type testClass) + => _contextProvider.GetOrCreateClassContext(testClass).Activity is not null; + + internal ValueTask> FinishClassActivityAsync( + [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] + Type testClass) + { + FinishClassActivity(testClass, hasErrors: false); + return new ValueTask>([]); + } + + private void FinishClassActivity( + [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] + Type testClass, bool hasErrors) + { + var classContext = _contextProvider.GetOrCreateClassContext(testClass); + lock (classContext) + { + var activity = classContext.Activity; + if (activity is null) + { + return; + } + + activity.SetTag(TUnitActivitySource.TagTestCount, classContext.TestCount); + + if (hasErrors) + { + activity.SetStatus(System.Diagnostics.ActivityStatusCode.Error); + } + + TUnitActivitySource.StopActivity(activity); + classContext.Activity = null; + } } #endif diff --git a/src/TUnit.Engine/TestExecutor.cs b/src/TUnit.Engine/TestExecutor.cs index 968a3badb3..f3dd9c5c69 100644 --- a/src/TUnit.Engine/TestExecutor.cs +++ b/src/TUnit.Engine/TestExecutor.cs @@ -36,7 +36,17 @@ internal class TestExecutor // each time (these run on the hot path and the factory is only invoked on the first cache miss). private readonly Func _beforeTestSessionHookFactory; private readonly Func _beforeAssemblyHookFactory; + private readonly Func>> _cancelledAfterAssemblyHookFactory; + private readonly AfterClassExecutor _cancelledAfterClassHookFactory; +#if NET + private readonly Func>> _finishAssemblyActivityFactory; + private readonly AfterClassExecutor _finishClassActivityFactory; +#endif + [UnconditionalSuppressMessage("Trimming", "IL2067", + Justification = "Class cleanup delegates receive only test-class types annotated at the execution boundary.")] + [UnconditionalSuppressMessage("Trimming", "IL2111", + Justification = "The annotated class cleanup method is captured as a delegate, not accessed through reflection.")] public TestExecutor( HookExecutor hookExecutor, TestLifecycleCoordinator lifecycleCoordinator, @@ -54,6 +64,12 @@ public TestExecutor( _beforeTestSessionHookFactory = ct => _hookExecutor.ExecuteBeforeTestSessionHooksAsync(ct); _beforeAssemblyHookFactory = (assembly, ct) => _hookExecutor.ExecuteBeforeAssemblyHooksAsync(assembly, ct); + _cancelledAfterAssemblyHookFactory = assembly => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, CancellationToken.None); + _cancelledAfterClassHookFactory = testClass => _hookExecutor.ExecuteAfterClassHooksAsync(testClass, CancellationToken.None); +#if NET + _finishAssemblyActivityFactory = _hookExecutor.FinishAssemblyActivityAsync; + _finishClassActivityFactory = _hookExecutor.FinishClassActivityAsync; +#endif } @@ -105,20 +121,34 @@ await _beforeHookTaskCache.GetOrCreateBeforeTestSessionTask( // it re-applies the same captured contexts). test.Context.ClassContext.AssemblyContext.TestSessionContext.RestoreExecutionContext(); - if (HasAssemblyHooks(testClass.Assembly)) + var hasAssemblyHooks = HasAssemblyHooks(testClass.Assembly); + if (hasAssemblyHooks) { await _beforeHookTaskCache.GetOrCreateBeforeAssemblyTask( testClass.Assembly, _beforeAssemblyHookFactory, cancellationToken).ConfigureAwait(false); } +#if NET + else + { + _hookExecutor.TryStartAssemblyActivity(testClass.Assembly); + } +#endif test.Context.ClassContext.AssemblyContext.RestoreExecutionContext(); - if (HasClassHooks(testClass)) + var hasClassHooks = HasClassHooks(testClass); + if (hasClassHooks) { await _beforeHookTaskCache.GetOrCreateBeforeClassTask(testClass, _hookExecutor, cancellationToken).ConfigureAwait(false); } +#if NET + else + { + _hookExecutor.TryStartClassActivity(testClass); + } +#endif // Note: the caller (TestCoordinator) restores ClassContext.RestoreExecutionContext() right // before constructing the instance so AsyncLocals captured by BeforeAssembly/BeforeClass flow @@ -164,12 +194,33 @@ await _beforeHookTaskCache.GetOrCreateBeforeAssemblyTask( testAssembly, _beforeAssemblyHookFactory, cancellationToken).ConfigureAwait(false); + } +#if NET + else + { + _hookExecutor.TryStartAssemblyActivity(testAssembly); + } +#endif + + Func>>? cancelledAfterAssemblyFactory = null; + if (hasAssemblyHooks) + { + cancelledAfterAssemblyFactory = _cancelledAfterAssemblyHookFactory; + } +#if NET + else if (_hookExecutor.HasAssemblyActivity(testAssembly)) + { + cancelledAfterAssemblyFactory = _finishAssemblyActivityFactory; + } +#endif - // Register After Assembly hook to run on cancellation (guarantees cleanup) + if (cancelledAfterAssemblyFactory is not null) + { + // Register lifecycle cleanup on cancellation. _afterHookPairTracker.RegisterAfterAssemblyHook( testAssembly, cancellationToken, - (assembly) => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, CancellationToken.None)); + cancelledAfterAssemblyFactory); } await _eventReceiverOrchestrator.InvokeFirstTestInAssemblyEventReceiversAsync( @@ -182,9 +233,33 @@ await _eventReceiverOrchestrator.InvokeFirstTestInAssemblyEventReceiversAsync( if (hasClassHooks) { await _beforeHookTaskCache.GetOrCreateBeforeClassTask(testClass, _hookExecutor, cancellationToken).ConfigureAwait(false); + } +#if NET + else + { + _hookExecutor.TryStartClassActivity(testClass); + } +#endif + + AfterClassExecutor? cancelledAfterClassFactory = null; + if (hasClassHooks) + { + cancelledAfterClassFactory = _cancelledAfterClassHookFactory; + } +#if NET + else if (_hookExecutor.HasClassActivity(testClass)) + { + cancelledAfterClassFactory = _finishClassActivityFactory; + } +#endif - // Register After Class hook to run on cancellation (guarantees cleanup) - _afterHookPairTracker.RegisterAfterClassHook(testClass, _hookExecutor, cancellationToken); + if (cancelledAfterClassFactory is not null) + { + // Register lifecycle cleanup on cancellation. + _afterHookPairTracker.RegisterAfterClassHook( + testClass, + cancellationToken, + cancelledAfterClassFactory); } await _eventReceiverOrchestrator.InvokeFirstTestInClassEventReceiversAsync( @@ -577,6 +652,8 @@ await executableTest.InvokeTestAsync( } } + [UnconditionalSuppressMessage("Trimming", "IL2067", + Justification = "The class cleanup delegate is invoked with the annotated testClass parameter.")] internal async Task?> ExecuteAfterClassAssemblyHooks(AbstractExecutableTest executableTest, [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] @@ -593,23 +670,53 @@ await executableTest.InvokeTestAsync( if (flags.ShouldExecuteAfterClass) { - // Use AfterHookPairTracker to prevent double execution if already triggered by cancellation - var classExceptions = await _afterHookPairTracker.GetOrCreateAfterClassTask(testClass, _hookExecutor, cancellationToken).ConfigureAwait(false); - if (classExceptions.Count > 0) + AfterClassExecutor? afterClassFactory = null; + if (HasClassHooks(testClass)) + { + afterClassFactory = type => _hookExecutor.ExecuteAfterClassHooksAsync(type, cancellationToken); + } +#if NET + else if (_hookExecutor.HasClassActivity(testClass)) { - (exceptions ??= []).AddRange(classExceptions); + afterClassFactory = _finishClassActivityFactory; + } +#endif + + if (afterClassFactory is not null) + { + // Use AfterHookPairTracker to prevent double execution if already triggered by cancellation + var classExceptions = await _afterHookPairTracker.GetOrCreateAfterClassTask(testClass, afterClassFactory).ConfigureAwait(false); + if (classExceptions.Count > 0) + { + (exceptions ??= []).AddRange(classExceptions); + } } } if (flags.ShouldExecuteAfterAssembly) { - // Use AfterHookPairTracker to prevent double execution if already triggered by cancellation - var assemblyExceptions = await _afterHookPairTracker.GetOrCreateAfterAssemblyTask( - testAssembly, - (assembly) => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, cancellationToken)).ConfigureAwait(false); - if (assemblyExceptions.Count > 0) + Func>>? afterAssemblyFactory = null; + if (HasAssemblyHooks(testAssembly)) + { + afterAssemblyFactory = assembly => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, cancellationToken); + } +#if NET + else if (_hookExecutor.HasAssemblyActivity(testAssembly)) + { + afterAssemblyFactory = _finishAssemblyActivityFactory; + } +#endif + + if (afterAssemblyFactory is not null) { - (exceptions ??= []).AddRange(assemblyExceptions); + // Use AfterHookPairTracker to prevent double execution if already triggered by cancellation + var assemblyExceptions = await _afterHookPairTracker.GetOrCreateAfterAssemblyTask( + testAssembly, + afterAssemblyFactory).ConfigureAwait(false); + if (assemblyExceptions.Count > 0) + { + (exceptions ??= []).AddRange(assemblyExceptions); + } } } diff --git a/tests/TUnit.UnitTests/SessionActivityLifecycleTests.cs b/tests/TUnit.UnitTests/SessionActivityLifecycleTests.cs index 6de4255305..5e06b11cca 100644 --- a/tests/TUnit.UnitTests/SessionActivityLifecycleTests.cs +++ b/tests/TUnit.UnitTests/SessionActivityLifecycleTests.cs @@ -29,6 +29,13 @@ public class SessionActivityLifecycleTests /// private static (HookExecutor Executor, TestSessionContext SessionContext) CreateHookExecutor( string? testFilter = null) + { + var (executor, sessionContext, _) = CreateLifecycleHookExecutor(testFilter); + return (executor, sessionContext); + } + + private static (HookExecutor Executor, TestSessionContext SessionContext, StubContextProvider ContextProvider) + CreateLifecycleHookExecutor(string? testFilter = null) { var beforeDiscovery = new BeforeTestDiscoveryContext { TestFilter = testFilter }; var discoveryContext = new TestDiscoveryContext(beforeDiscovery) { TestFilter = testFilter }; @@ -45,7 +52,7 @@ private static (HookExecutor Executor, TestSessionContext SessionContext) Create // so we pass null — the executor will never call it in these tests. var executor = new HookExecutor(hookDelegateBuilder, contextProvider, null!); - return (executor, sessionContext); + return (executor, sessionContext, contextProvider); } [Test] @@ -219,6 +226,38 @@ public async Task LifecycleHierarchy_UsesSessionTrace_AndTestStartsSeparateTrace TUnitActivitySource.StopActivity(assemblyActivity); } + [Test] + public async Task HooklessAssemblyAndClassActivities_PreserveLifecycleHierarchy() + { + var (executor, sessionContext, contextProvider) = CreateLifecycleHookExecutor(); + var testClass = typeof(SessionActivityLifecycleTests); + var assembly = testClass.Assembly; + + using var scope = new ActivityListenerScope(); + + executor.TryStartSessionActivity(); + executor.TryStartAssemblyActivity(assembly); + executor.TryStartClassActivity(testClass); + + var assemblyContext = contextProvider.GetOrCreateAssemblyContext(assembly); + var classContext = contextProvider.GetOrCreateClassContext(testClass); + var assemblyActivity = assemblyContext.Activity; + var classActivity = classContext.Activity; + + await Assert.That(assemblyActivity).IsNotNull(); + await Assert.That(assemblyActivity!.ParentId).IsEqualTo(sessionContext.Activity!.Id); + await Assert.That(classActivity).IsNotNull(); + await Assert.That(classActivity!.ParentId).IsEqualTo(assemblyActivity.Id); + + await executor.FinishClassActivityAsync(testClass); + await executor.FinishAssemblyActivityAsync(assembly); + + await Assert.That(classActivity.IsStopped).IsTrue(); + await Assert.That(assemblyActivity.IsStopped).IsTrue(); + await Assert.That(classContext.Activity).IsNull(); + await Assert.That(assemblyContext.Activity).IsNull(); + } + [Test] public async Task DiscoverySpan_InSameTrace_WhenParentedUnderSession() { @@ -290,6 +329,9 @@ public async Task HookSpan_ParentsUnderSession_WhenSessionActivityExists() /// private sealed class StubContextProvider(TestSessionContext sessionContext) : IContextProvider { + private readonly ConcurrentDictionary _assemblyContexts = new(); + private readonly ConcurrentDictionary _classContexts = new(); + public BeforeTestDiscoveryContext BeforeTestDiscoveryContext => throw new NotSupportedException(); @@ -299,14 +341,21 @@ private sealed class StubContextProvider(TestSessionContext sessionContext) : IC public TestSessionContext TestSessionContext => sessionContext; public AssemblyHookContext GetOrCreateAssemblyContext(Assembly assembly) => - throw new NotSupportedException(); + _assemblyContexts.GetOrAdd(assembly, static (key, context) => new AssemblyHookContext(context) + { + Assembly = key + }, sessionContext); public ClassHookContext GetOrCreateClassContext( [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] Type classType) => - throw new NotSupportedException(); + _classContexts.GetOrAdd(classType, static (key, provider) => + new ClassHookContext(provider.GetOrCreateAssemblyContext(key.Assembly)) + { + ClassType = key + }, this); public TestContext CreateTestContext( [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | From 0d8e1cfd1741dfb69193200fe232017ad02ba007 Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:17:51 +0100 Subject: [PATCH 3/4] fix(engine): synchronize lifecycle activities --- src/TUnit.Core/Context.cs | 8 +- src/TUnit.Core/Models/AssemblyHookContext.cs | 1 + src/TUnit.Core/Models/ClassHookContext.cs | 1 + src/TUnit.Engine/Services/HookExecutor.cs | 8 +- src/TUnit.Engine/TestExecutor.cs | 101 +++++++++++-------- 5 files changed, 70 insertions(+), 49 deletions(-) diff --git a/src/TUnit.Core/Context.cs b/src/TUnit.Core/Context.cs index 6e2c9a677e..236aa923df 100644 --- a/src/TUnit.Core/Context.cs +++ b/src/TUnit.Core/Context.cs @@ -78,7 +78,13 @@ internal Context(Context? parent) } #if NET - internal System.Diagnostics.Activity? Activity { get; set; } + private System.Diagnostics.Activity? _activity; + + internal System.Diagnostics.Activity? Activity + { + get => Volatile.Read(ref _activity); + set => Volatile.Write(ref _activity, value); + } internal ExecutionContext? ExecutionContext { get; private set; } #endif diff --git a/src/TUnit.Core/Models/AssemblyHookContext.cs b/src/TUnit.Core/Models/AssemblyHookContext.cs index 650f1217cb..2f6096334c 100644 --- a/src/TUnit.Core/Models/AssemblyHookContext.cs +++ b/src/TUnit.Core/Models/AssemblyHookContext.cs @@ -27,6 +27,7 @@ internal AssemblyHookContext(TestSessionContext testSessionContext) : base(testS public required Assembly Assembly { get; init; } private readonly Lock _lock = new(); + internal Lock SynchronizationLock => _lock; private readonly List _testClasses = []; private TestContext[]? _cachedAllTests; diff --git a/src/TUnit.Core/Models/ClassHookContext.cs b/src/TUnit.Core/Models/ClassHookContext.cs index 452a9cb5f1..7d4571e0ee 100644 --- a/src/TUnit.Core/Models/ClassHookContext.cs +++ b/src/TUnit.Core/Models/ClassHookContext.cs @@ -29,6 +29,7 @@ internal ClassHookContext(AssemblyHookContext assemblyHookContext) : base(assemb public required Type ClassType { get; init; } private readonly Lock _lock = new(); + internal Lock SynchronizationLock => _lock; private readonly HashSet _testSet = new(ReferenceEqualityComparer.Instance); private readonly List _tests = []; diff --git a/src/TUnit.Engine/Services/HookExecutor.cs b/src/TUnit.Engine/Services/HookExecutor.cs index 5e6167290c..56200da59e 100644 --- a/src/TUnit.Engine/Services/HookExecutor.cs +++ b/src/TUnit.Engine/Services/HookExecutor.cs @@ -294,7 +294,7 @@ internal void TryStartAssemblyActivity(Assembly assembly) return; } - lock (assemblyContext) + lock (assemblyContext.SynchronizationLock) { if (assemblyContext.Activity is not null) { @@ -324,7 +324,7 @@ internal ValueTask> FinishAssemblyActivityAsync(Assembly assembl private void FinishAssemblyActivity(Assembly assembly, bool hasErrors) { var assemblyContext = _contextProvider.GetOrCreateAssemblyContext(assembly); - lock (assemblyContext) + lock (assemblyContext.SynchronizationLock) { var activity = assemblyContext.Activity; if (activity is null) @@ -488,7 +488,7 @@ internal void TryStartClassActivity( return; } - lock (classContext) + lock (classContext.SynchronizationLock) { if (classContext.Activity is not null) { @@ -525,7 +525,7 @@ private void FinishClassActivity( Type testClass, bool hasErrors) { var classContext = _contextProvider.GetOrCreateClassContext(testClass); - lock (classContext) + lock (classContext.SynchronizationLock) { var activity = classContext.Activity; if (activity is null) diff --git a/src/TUnit.Engine/TestExecutor.cs b/src/TUnit.Engine/TestExecutor.cs index f3dd9c5c69..0cb4c752e2 100644 --- a/src/TUnit.Engine/TestExecutor.cs +++ b/src/TUnit.Engine/TestExecutor.cs @@ -202,17 +202,10 @@ await _beforeHookTaskCache.GetOrCreateBeforeAssemblyTask( } #endif - Func>>? cancelledAfterAssemblyFactory = null; - if (hasAssemblyHooks) - { - cancelledAfterAssemblyFactory = _cancelledAfterAssemblyHookFactory; - } -#if NET - else if (_hookExecutor.HasAssemblyActivity(testAssembly)) - { - cancelledAfterAssemblyFactory = _finishAssemblyActivityFactory; - } -#endif + var cancelledAfterAssemblyFactory = ResolveAssemblyCleanup( + testAssembly, + hasAssemblyHooks, + _cancelledAfterAssemblyHookFactory); if (cancelledAfterAssemblyFactory is not null) { @@ -241,17 +234,10 @@ await _eventReceiverOrchestrator.InvokeFirstTestInAssemblyEventReceiversAsync( } #endif - AfterClassExecutor? cancelledAfterClassFactory = null; - if (hasClassHooks) - { - cancelledAfterClassFactory = _cancelledAfterClassHookFactory; - } -#if NET - else if (_hookExecutor.HasClassActivity(testClass)) - { - cancelledAfterClassFactory = _finishClassActivityFactory; - } -#endif + var cancelledAfterClassFactory = ResolveClassCleanup( + testClass, + hasClassHooks, + _cancelledAfterClassHookFactory); if (cancelledAfterClassFactory is not null) { @@ -652,6 +638,47 @@ await executableTest.InvokeTestAsync( } } + private Func>>? ResolveAssemblyCleanup( + Assembly assembly, + bool hasHooks, + Func>> hookFactory) + { + if (hasHooks) + { + return hookFactory; + } + +#if NET + if (_hookExecutor.HasAssemblyActivity(assembly)) + { + return _finishAssemblyActivityFactory; + } +#endif + + return null; + } + + private AfterClassExecutor? ResolveClassCleanup( + [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] + Type testClass, + bool hasHooks, + AfterClassExecutor hookFactory) + { + if (hasHooks) + { + return hookFactory; + } + +#if NET + if (_hookExecutor.HasClassActivity(testClass)) + { + return _finishClassActivityFactory; + } +#endif + + return null; + } + [UnconditionalSuppressMessage("Trimming", "IL2067", Justification = "The class cleanup delegate is invoked with the annotated testClass parameter.")] internal async Task?> ExecuteAfterClassAssemblyHooks(AbstractExecutableTest executableTest, @@ -670,17 +697,10 @@ await executableTest.InvokeTestAsync( if (flags.ShouldExecuteAfterClass) { - AfterClassExecutor? afterClassFactory = null; - if (HasClassHooks(testClass)) - { - afterClassFactory = type => _hookExecutor.ExecuteAfterClassHooksAsync(type, cancellationToken); - } -#if NET - else if (_hookExecutor.HasClassActivity(testClass)) - { - afterClassFactory = _finishClassActivityFactory; - } -#endif + var afterClassFactory = ResolveClassCleanup( + testClass, + HasClassHooks(testClass), + type => _hookExecutor.ExecuteAfterClassHooksAsync(type, cancellationToken)); if (afterClassFactory is not null) { @@ -695,17 +715,10 @@ await executableTest.InvokeTestAsync( if (flags.ShouldExecuteAfterAssembly) { - Func>>? afterAssemblyFactory = null; - if (HasAssemblyHooks(testAssembly)) - { - afterAssemblyFactory = assembly => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, cancellationToken); - } -#if NET - else if (_hookExecutor.HasAssemblyActivity(testAssembly)) - { - afterAssemblyFactory = _finishAssemblyActivityFactory; - } -#endif + var afterAssemblyFactory = ResolveAssemblyCleanup( + testAssembly, + HasAssemblyHooks(testAssembly), + assembly => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, cancellationToken)); if (afterAssemblyFactory is not null) { From f25e027a18e034649681c29c7aa7fed85f7a60e2 Mon Sep 17 00:00:00 2001 From: Tom Longhurst <30480171+thomhurst@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:51:26 +0100 Subject: [PATCH 4/4] fix(engine): make class cleanup AOT-safe --- .../Services/AfterHookPairTracker.cs | 55 ++++++++++++++++--- src/TUnit.Engine/TestExecutor.cs | 36 ++++++------ 2 files changed, 62 insertions(+), 29 deletions(-) diff --git a/src/TUnit.Engine/Services/AfterHookPairTracker.cs b/src/TUnit.Engine/Services/AfterHookPairTracker.cs index 69ee5ece79..e8fc8d79e5 100644 --- a/src/TUnit.Engine/Services/AfterHookPairTracker.cs +++ b/src/TUnit.Engine/Services/AfterHookPairTracker.cs @@ -5,9 +5,46 @@ namespace TUnit.Engine.Services; -internal delegate ValueTask> AfterClassExecutor( - [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] - Type testClass); +internal readonly struct AfterClassCleanup +{ + private readonly HookExecutor _hookExecutor; + private readonly CancellationToken _cancellationToken; + private readonly bool _finishActivity; + + private AfterClassCleanup( + HookExecutor hookExecutor, + CancellationToken cancellationToken, + bool finishActivity) + { + _hookExecutor = hookExecutor; + _cancellationToken = cancellationToken; + _finishActivity = finishActivity; + } + + internal static AfterClassCleanup ForHooks( + HookExecutor hookExecutor, + CancellationToken cancellationToken) + => new(hookExecutor, cancellationToken, finishActivity: false); + +#if NET + internal static AfterClassCleanup ForActivity(HookExecutor hookExecutor) + => new(hookExecutor, CancellationToken.None, finishActivity: true); +#endif + + internal ValueTask> ExecuteAsync( + [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] + Type testClass) + { +#if NET + if (_finishActivity) + { + return _hookExecutor.FinishClassActivityAsync(testClass); + } +#endif + + return _hookExecutor.ExecuteAfterClassHooksAsync(testClass, _cancellationToken); + } +} /// /// Responsible for ensuring After hooks run even when tests are cancelled. @@ -114,7 +151,7 @@ public void RegisterAfterClassHook( [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] Type testClass, CancellationToken sessionCancellationToken, - AfterClassExecutor afterHookExecutor) + AfterClassCleanup cleanup) { if (!_classHookRegistered.Add(testClass)) { @@ -123,9 +160,9 @@ public void RegisterAfterClassHook( var registration = sessionCancellationToken.Register(static state => { - var (pairTracker, testClass, afterHookExecutor) = ((AfterHookPairTracker, Type, AfterClassExecutor))state!; - _ = pairTracker.GetOrCreateAfterClassTask(testClass, afterHookExecutor); - }, (this, testClass, afterHookExecutor)); + var (pairTracker, testClass, cleanup) = ((AfterHookPairTracker, Type, AfterClassCleanup))state!; + _ = pairTracker.GetOrCreateAfterClassTask(testClass, cleanup); + }, (this, testClass, cleanup)); _registrations.Add(registration); } @@ -180,7 +217,7 @@ public ValueTask> GetOrCreateAfterAssemblyTask(Assembly assembly public ValueTask> GetOrCreateAfterClassTask( [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] Type testClass, - AfterClassExecutor taskFactory) + AfterClassCleanup cleanup) { // Lock-free fast path avoids allocating a closure on the common cache-hit case. if (_afterClassTasks.TryGetValue(testClass, out var existingTask)) @@ -193,7 +230,7 @@ public ValueTask> GetOrCreateAfterClassTask( // behind a shared lock. var task = _afterClassTasks.GetOrAdd( testClass, - _ => taskFactory(testClass).AsTask()); + _ => cleanup.ExecuteAsync(testClass).AsTask()); return new ValueTask>(task); } diff --git a/src/TUnit.Engine/TestExecutor.cs b/src/TUnit.Engine/TestExecutor.cs index 0cb4c752e2..4ca09252de 100644 --- a/src/TUnit.Engine/TestExecutor.cs +++ b/src/TUnit.Engine/TestExecutor.cs @@ -37,16 +37,12 @@ internal class TestExecutor private readonly Func _beforeTestSessionHookFactory; private readonly Func _beforeAssemblyHookFactory; private readonly Func>> _cancelledAfterAssemblyHookFactory; - private readonly AfterClassExecutor _cancelledAfterClassHookFactory; + private readonly AfterClassCleanup _cancelledAfterClassHookCleanup; #if NET private readonly Func>> _finishAssemblyActivityFactory; - private readonly AfterClassExecutor _finishClassActivityFactory; + private readonly AfterClassCleanup _finishClassActivityCleanup; #endif - [UnconditionalSuppressMessage("Trimming", "IL2067", - Justification = "Class cleanup delegates receive only test-class types annotated at the execution boundary.")] - [UnconditionalSuppressMessage("Trimming", "IL2111", - Justification = "The annotated class cleanup method is captured as a delegate, not accessed through reflection.")] public TestExecutor( HookExecutor hookExecutor, TestLifecycleCoordinator lifecycleCoordinator, @@ -65,10 +61,10 @@ public TestExecutor( _beforeTestSessionHookFactory = ct => _hookExecutor.ExecuteBeforeTestSessionHooksAsync(ct); _beforeAssemblyHookFactory = (assembly, ct) => _hookExecutor.ExecuteBeforeAssemblyHooksAsync(assembly, ct); _cancelledAfterAssemblyHookFactory = assembly => _hookExecutor.ExecuteAfterAssemblyHooksAsync(assembly, CancellationToken.None); - _cancelledAfterClassHookFactory = testClass => _hookExecutor.ExecuteAfterClassHooksAsync(testClass, CancellationToken.None); + _cancelledAfterClassHookCleanup = AfterClassCleanup.ForHooks(_hookExecutor, CancellationToken.None); #if NET _finishAssemblyActivityFactory = _hookExecutor.FinishAssemblyActivityAsync; - _finishClassActivityFactory = _hookExecutor.FinishClassActivityAsync; + _finishClassActivityCleanup = AfterClassCleanup.ForActivity(_hookExecutor); #endif } @@ -234,18 +230,18 @@ await _eventReceiverOrchestrator.InvokeFirstTestInAssemblyEventReceiversAsync( } #endif - var cancelledAfterClassFactory = ResolveClassCleanup( + var cancelledAfterClassCleanup = ResolveClassCleanup( testClass, hasClassHooks, - _cancelledAfterClassHookFactory); + _cancelledAfterClassHookCleanup); - if (cancelledAfterClassFactory is not null) + if (cancelledAfterClassCleanup is { } cleanup) { // Register lifecycle cleanup on cancellation. _afterHookPairTracker.RegisterAfterClassHook( testClass, cancellationToken, - cancelledAfterClassFactory); + cleanup); } await _eventReceiverOrchestrator.InvokeFirstTestInClassEventReceiversAsync( @@ -658,21 +654,21 @@ await executableTest.InvokeTestAsync( return null; } - private AfterClassExecutor? ResolveClassCleanup( + private AfterClassCleanup? ResolveClassCleanup( [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicConstructors | DynamicallyAccessedMemberTypes.PublicProperties | DynamicallyAccessedMemberTypes.PublicMethods)] Type testClass, bool hasHooks, - AfterClassExecutor hookFactory) + AfterClassCleanup hookCleanup) { if (hasHooks) { - return hookFactory; + return hookCleanup; } #if NET if (_hookExecutor.HasClassActivity(testClass)) { - return _finishClassActivityFactory; + return _finishClassActivityCleanup; } #endif @@ -697,15 +693,15 @@ await executableTest.InvokeTestAsync( if (flags.ShouldExecuteAfterClass) { - var afterClassFactory = ResolveClassCleanup( + var afterClassCleanup = ResolveClassCleanup( testClass, HasClassHooks(testClass), - type => _hookExecutor.ExecuteAfterClassHooksAsync(type, cancellationToken)); + AfterClassCleanup.ForHooks(_hookExecutor, cancellationToken)); - if (afterClassFactory is not null) + if (afterClassCleanup is { } cleanup) { // Use AfterHookPairTracker to prevent double execution if already triggered by cancellation - var classExceptions = await _afterHookPairTracker.GetOrCreateAfterClassTask(testClass, afterClassFactory).ConfigureAwait(false); + var classExceptions = await _afterHookPairTracker.GetOrCreateAfterClassTask(testClass, cleanup).ConfigureAwait(false); if (classExceptions.Count > 0) { (exceptions ??= []).AddRange(classExceptions);