Flow AsyncLocal Values between Hooks + Tests via context.AddAsyncLocalValues()

TUnit.Analyzers.Tests/BeforeHookAsyncLocalAnalyzerTests.cs

@@ -20,49 +20,21 @@ public class MyClass
                 {
                     private static readonly AsyncLocal<int> _asyncLocal = new();
 
-                    [Before(Test)]
-                    public void {|#0:MyTest|}()
+                    {|#0:[Before(Class)]
+                    public void MyTest()
                     {
                         _asyncLocal.Value = 1;
-                    }
-                }
-                """
-            );
-    }
-
-    [Test]
-    public async Task Async_Raises_Error_Setting_AsyncLocal()
-    {
-        await Verifier
-            .VerifyAnalyzerAsync(
-                """
-                using System.Threading;
-                using System.Threading.Tasks;
-                using TUnit.Core;
-                using static TUnit.Core.HookType;
-
-                public class MyClass
-                {
-                    private static readonly AsyncLocal<int> _asyncLocal = new();
-
-                    {|#1:[Before(Test)]
-                    public async Task MyTest()
-                    {
-                        {|#0:_asyncLocal.Value = 1|};
-                        await Task.Yield();
                     }|}
                 }
                 """,
-
                 Verifier
-                    .Diagnostic(Rules.AsyncLocalVoidMethod)
+                    .Diagnostic(Rules.AsyncLocalCallFlowValues)
                     .WithLocation(0)
-                    .WithLocation(1)
             );
     }
 
     [Test]
-    public async Task Async_Raises_Error_Setting_AsyncLocal_Nested_Method()
+    public async Task AddAsyncLocalValues_No_Error()
     {
         await Verifier
             .VerifyAnalyzerAsync(
@@ -76,24 +48,13 @@ public class MyClass
                 {
                     private static readonly AsyncLocal<int> _asyncLocal = new();
 
-                    {|#1:[Before(Test)]
-                    public async Task MyTest()
+                    {|#0:[Before(Class)]
+                    public void MyTest(ClassHookContext context)
                     {
-                        SetAsyncLocal();
-                        await Task.Yield();
+                        _asyncLocal.Value = 1;
+                        context.AddAsyncLocalValues();
                     }|}
-
-                    private void SetAsyncLocal()
-                    {
-                        {|#0:_asyncLocal.Value = 1|};
-                    }
                 }
-                """,
-
-                Verifier
-                    .Diagnostic(Rules.AsyncLocalVoidMethod)
-                    .WithLocation(0)
-                    .WithLocation(1)
-            );
+                """);
     }
 }

TUnit.Analyzers/BeforeHookAsyncLocalAnalyzer.cs

@@ -11,7 +11,7 @@ namespace TUnit.Analyzers;
 public class BeforeHookAsyncLocalAnalyzer : ConcurrentDiagnosticAnalyzer
 {
     public override ImmutableArray<DiagnosticDescriptor> SupportedDiagnostics { get; } =
-        ImmutableArray.Create(Rules.AsyncLocalVoidMethod);
+        ImmutableArray.Create(Rules.AsyncLocalCallFlowValues);
 
     protected override void InitializeInternal(AnalysisContext context)
     { 
@@ -39,16 +39,31 @@ private void AnalyzeOperation(OperationAnalysisContext context)
             return;
         }
 
-        var parent = assignmentOperation.Parent;
-        while (parent != null)
+        var methodBodyOperation = GetParentMethod(assignmentOperation.Parent);
+
+        if (methodBodyOperation is null)
+        {
+            return;
+        }
+
+        CheckMethod(context, methodBodyOperation);
+    }
+
+    private IMethodBodyOperation? GetParentMethod(IOperation? assignmentOperationParent)
+    {
+        var parent = assignmentOperationParent;
+
+        while (parent is not null)
         {
             if (parent is IMethodBodyOperation methodBodyOperation)
             {
-                CheckMethod(context, methodBodyOperation);
+                return methodBodyOperation;
             }
 
             parent = parent.Parent;
         }
+
+        return null;
     }
 
     private void CheckMethod(OperationAnalysisContext context, IMethodBodyOperation methodBodyOperation)
@@ -59,19 +74,20 @@ private void CheckMethod(OperationAnalysisContext context, IMethodBodyOperation
             return;
         }
 
-        if (methodSymbol.IsHookMethod(context.Compilation, out _, out _, out var type)
-            && type is HookType.Before 
-            && !methodSymbol.ReturnsVoid)
+        if (!methodSymbol.IsHookMethod(context.Compilation, out _, out _, out var type)
+            || type is not HookType.Before)
         {
-            context.ReportDiagnostic(Diagnostic.Create(Rules.AsyncLocalVoidMethod,
-                context.Operation.Syntax.GetLocation(),
-                [methodBodyOperation.Syntax.GetLocation()]));
             return;
         }
 
-        var invocations = methodBodyOperation.SemanticModel
-            .SyntaxTree
-            .GetRoot()
+        var syntax = methodSymbol.DeclaringSyntaxReferences.FirstOrDefault()?.GetSyntax();
+
+        if (syntax is null)
+        {
+            return;
+        }
+
+        var invocations = syntax
             .DescendantNodes()
             .OfType<InvocationExpressionSyntax>();
 
@@ -84,36 +100,13 @@ private void CheckMethod(OperationAnalysisContext context, IMethodBodyOperation
                 continue;
             }
 
-            if (!SymbolEqualityComparer.Default.Equals(invocationOperation.TargetMethod, methodSymbol))
+            if (invocationOperation.TargetMethod.Name == "AddAsyncLocalValues")
             {
-                continue;
+                return;
             }
-
-            var parentMethodBody = GetParentMethodBody(invocationOperation);
-
-            if (parentMethodBody == null)
-            {
-                continue;
-            }
-
-            CheckMethod(context, parentMethodBody);
-        }
-    }
-
-    private IMethodBodyOperation? GetParentMethodBody(IInvocationOperation invocationOperation)
-    {
-        var parent = invocationOperation.Parent;
-
-        while (parent != null)
-        {
-            if (parent is IMethodBodyOperation methodBodyOperation)
-            {
-                return methodBodyOperation;
-            }
-
-            parent = parent.Parent;
         }
 
-        return null;
+        context.ReportDiagnostic(Diagnostic.Create(Rules.AsyncLocalCallFlowValues,
+            methodBodyOperation.Syntax.GetLocation()));
     }
 }

TUnit.Analyzers/Resources.resx

@@ -433,13 +433,13 @@
         <value>Return a `Func&lt;T&gt;` rather than a `&lt;T&gt;`</value>
     </data>
     <data name="TUnit0047Description" xml:space="preserve">
-        <value>Before hooks setting AsyncLocal values should be non-async void returning methods.</value>
+        <value>For AsyncLocal values set in before hooks, you must call `context.FlowAsyncLocalValues` to access them within tests.</value>
     </data>
     <data name="TUnit0047MessageFormat" xml:space="preserve">
-        <value>Before hooks setting AsyncLocal values should be non-async void returning methods.</value>
+        <value>For AsyncLocal values set in before hooks, you must call `context.FlowAsyncLocalValues` to access them within tests.</value>
     </data>
     <data name="TUnit0047Title" xml:space="preserve">
-        <value>Before hooks setting AsyncLocal values should be non-async void returning methods</value>
+        <value>Call `context.FlowAsyncLocalValues`</value>
     </data>
     <data name="TUnit0048Description" xml:space="preserve">
         <value>Test methods must not be static.</value>

TUnit.Analyzers/Rules.cs

@@ -108,7 +108,7 @@ public static class Rules
     public static readonly DiagnosticDescriptor ReturnFunc =
         CreateDescriptor("TUnit0046", UsageCategory, DiagnosticSeverity.Warning);
 
-    public static readonly DiagnosticDescriptor AsyncLocalVoidMethod =
+    public static readonly DiagnosticDescriptor AsyncLocalCallFlowValues =
         CreateDescriptor("TUnit0047", UsageCategory, DiagnosticSeverity.Warning);
 
     public static DiagnosticDescriptor InstanceTestMethod =

TUnit.Core.SourceGenerator.Tests/AfterAllTests.Test.verified.txt

@@ -70,7 +70,7 @@ file partial class Hooks_Base1 : global::TUnit.Core.Interfaces.SourceGenerator.I
 						Properties = [],
 					}),
 },
-					AsyncBody = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.Base1.AfterAll1()),
+					Body = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.Base1.AfterAll1()),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					FilePath = @"",
@@ -159,7 +159,7 @@ file partial class Hooks_Base1 : global::TUnit.Core.Interfaces.SourceGenerator.I
 						Properties = [],
 					}),
 },
-					AsyncBody = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.Base1)classInstance).AfterEach1()),
+					Body = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.Base1)classInstance).AfterEach1()),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					MethodAttributes = 
@@ -246,7 +246,7 @@ file partial class Hooks_Base2 : global::TUnit.Core.Interfaces.SourceGenerator.I
 						Properties = [],
 					}),
 },
-					AsyncBody = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.Base2.AfterAll2()),
+					Body = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.Base2.AfterAll2()),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					FilePath = @"",
@@ -335,7 +335,7 @@ file partial class Hooks_Base2 : global::TUnit.Core.Interfaces.SourceGenerator.I
 						Properties = [],
 					}),
 },
-					AsyncBody = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.Base2)classInstance).AfterEach2()),
+					Body = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.Base2)classInstance).AfterEach2()),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					MethodAttributes = 
@@ -422,7 +422,7 @@ file partial class Hooks_Base3 : global::TUnit.Core.Interfaces.SourceGenerator.I
 						Properties = [],
 					}),
 },
-					AsyncBody = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.Base3.AfterAll3()),
+					Body = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.Base3.AfterAll3()),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					FilePath = @"",
@@ -511,7 +511,7 @@ file partial class Hooks_Base3 : global::TUnit.Core.Interfaces.SourceGenerator.I
 						Properties = [],
 					}),
 },
-					AsyncBody = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.Base3)classInstance).AfterEach3()),
+					Body = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.Base3)classInstance).AfterEach3()),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					MethodAttributes = 
@@ -598,7 +598,7 @@ file partial class Hooks_CleanupTests : global::TUnit.Core.Interfaces.SourceGene
 						Properties = [],
 					}),
 },
-					AsyncBody = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.CleanupTests.AfterAllCleanUp()),
+					Body = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.CleanupTests.AfterAllCleanUp()),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					FilePath = @"",
@@ -693,7 +693,7 @@ file partial class Hooks_CleanupTests : global::TUnit.Core.Interfaces.SourceGene
 						Properties = [],
 					}),
 },
-					AsyncBody = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.CleanupTests.AfterAllCleanUpWithContext(context)),
+					Body = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.CleanupTests.AfterAllCleanUpWithContext(context)),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					FilePath = @"",
@@ -788,7 +788,7 @@ file partial class Hooks_CleanupTests : global::TUnit.Core.Interfaces.SourceGene
 						Properties = [],
 					}),
 },
-					AsyncBody = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.CleanupTests.AfterAllCleanUp(cancellationToken)),
+					Body = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.CleanupTests.AfterAllCleanUp(cancellationToken)),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					FilePath = @"",
@@ -888,7 +888,7 @@ file partial class Hooks_CleanupTests : global::TUnit.Core.Interfaces.SourceGene
 						Properties = [],
 					}),
 },
-					AsyncBody = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.CleanupTests.AfterAllCleanUpWithContext(context, cancellationToken)),
+					Body = (context, cancellationToken) => AsyncConvert.Convert(() => global::TUnit.TestProject.AfterTests.CleanupTests.AfterAllCleanUpWithContext(context, cancellationToken)),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					FilePath = @"",
@@ -977,7 +977,7 @@ file partial class Hooks_CleanupTests : global::TUnit.Core.Interfaces.SourceGene
 						Properties = [],
 					}),
 },
-					AsyncBody = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.CleanupTests)classInstance).Cleanup()),
+					Body = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.CleanupTests)classInstance).Cleanup()),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					MethodAttributes = 
@@ -1073,7 +1073,7 @@ file partial class Hooks_CleanupTests : global::TUnit.Core.Interfaces.SourceGene
 						Properties = [],
 					}),
 },
-					AsyncBody = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.CleanupTests)classInstance).Cleanup(cancellationToken)),
+					Body = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.CleanupTests)classInstance).Cleanup(cancellationToken)),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					MethodAttributes = 
@@ -1169,7 +1169,7 @@ file partial class Hooks_CleanupTests : global::TUnit.Core.Interfaces.SourceGene
 						Properties = [],
 					}),
 },
-					AsyncBody = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.CleanupTests)classInstance).CleanupWithContext(context)),
+					Body = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.CleanupTests)classInstance).CleanupWithContext(context)),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					MethodAttributes = 
@@ -1270,7 +1270,7 @@ file partial class Hooks_CleanupTests : global::TUnit.Core.Interfaces.SourceGene
 						Properties = [],
 					}),
 },
-					AsyncBody = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.CleanupTests)classInstance).CleanupWithContext(context, cancellationToken)),
+					Body = (classInstance, context, cancellationToken) => AsyncConvert.Convert(() => ((global::TUnit.TestProject.AfterTests.CleanupTests)classInstance).CleanupWithContext(context, cancellationToken)),
 					HookExecutor = DefaultExecutor.Instance,
 					Order = 0,
 					MethodAttributes =