Make TestMethodReflectionInfo.Invoke async-aware (Task-returning) in MSTest.AotReflection.SourceGeneration

MSTest.slnf

@@ -11,6 +11,7 @@
       "src\\Analyzers\\MSTest.Analyzers.CodeFixes\\MSTest.Analyzers.CodeFixes.csproj",
       "src\\Analyzers\\MSTest.Analyzers.Package\\MSTest.Analyzers.Package.csproj",
       "src\\Analyzers\\MSTest.Analyzers\\MSTest.Analyzers.csproj",
+      "src\\Analyzers\\MSTest.AotReflection.SourceGeneration\\MSTest.AotReflection.SourceGeneration.csproj",
       "src\\Analyzers\\MSTest.GlobalConfigsGenerator\\MSTest.GlobalConfigsGenerator.csproj",
       "src\\Analyzers\\MSTest.SourceGeneration\\MSTest.SourceGeneration.csproj",
       "src\\Package\\MSTest.Sdk\\MSTest.Sdk.csproj",

TestFx.slnx

@@ -69,6 +69,7 @@
       <BuildDependency Project="src/Analyzers/MSTest.GlobalConfigsGenerator/MSTest.GlobalConfigsGenerator.csproj" />
     </Project>
     <Project Path="src/Analyzers/MSTest.Analyzers/MSTest.Analyzers.csproj" />
+    <Project Path="src/Analyzers/MSTest.AotReflection.SourceGeneration/MSTest.AotReflection.SourceGeneration.csproj" />
     <Project Path="src/Analyzers/MSTest.GlobalConfigsGenerator/MSTest.GlobalConfigsGenerator.csproj" />
     <Project Path="src/Analyzers/MSTest.SourceGeneration/MSTest.SourceGeneration.csproj" />
   </Folder>
@@ -133,6 +134,7 @@
     <Project Path="test/UnitTests/Microsoft.Testing.Platform.MSBuild.UnitTests/Microsoft.Testing.Platform.MSBuild.UnitTests.csproj" />
     <Project Path="test/UnitTests/Microsoft.Testing.Platform.UnitTests/Microsoft.Testing.Platform.UnitTests.csproj" />
     <Project Path="test/UnitTests/MSTest.Analyzers.UnitTests/MSTest.Analyzers.UnitTests.csproj" />
+    <Project Path="test/UnitTests/MSTest.AotReflection.SourceGeneration.UnitTests/MSTest.AotReflection.SourceGeneration.UnitTests.csproj" />
     <Project Path="test/UnitTests/MSTest.SelfRealExamples.UnitTests/MSTest.SelfRealExamples.UnitTests.csproj" />
     <Project Path="test/UnitTests/MSTest.SourceGeneration.UnitTests/MSTest.SourceGeneration.UnitTests.csproj" />
     <Project Path="test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/MSTestAdapter.PlatformServices.UnitTests.csproj" />

src/Analyzers/MSTest.AotReflection.SourceGeneration/Generators/MSTestReflectionMetadataGenerator.cs

@@ -45,14 +45,27 @@ node is TypeDeclarationSyntax type
             .Where(static model => model is not null)
             .Select(static (model, _) => model!);
 
-        IncrementalValueProvider<(string? AssemblyName, ImmutableArray<TestClassModel> Classes)> combined =
+        // Pull assembly-level attributes from the compilation (one value per run) and
+        // wrap them in an equatable model so this branch of the pipeline can stay cached
+        // when only test-class code changes.
+        IncrementalValueProvider<AssemblyMetadataModel> assemblyMetadata =
+            context.CompilationProvider.Select(static (c, ct) =>
+            {
+                ct.ThrowIfCancellationRequested();
+                return new AssemblyMetadataModel(
+                    TestClassModelBuilder.BuildAttributes(c.Assembly.GetAttributes()));
+            });
+
+        IncrementalValueProvider<(string? AssemblyName, AssemblyMetadataModel Metadata, ImmutableArray<TestClassModel> Classes)> combined =
             context.CompilationProvider.Select(static (c, _) => c.AssemblyName)
-                .Combine(testClasses.Collect());
+                .Combine(assemblyMetadata)
+                .Combine(testClasses.Collect())
+                .Select(static (tuple, _) => (tuple.Left.Left, tuple.Left.Right, tuple.Right));
 
         context.RegisterImplementationSourceOutput(combined, static (ctx, payload) =>
         {
             string assemblyName = payload.AssemblyName ?? "Unknown";
-            string source = MetadataRegistryEmitter.EmitRegistry(assemblyName, payload.Classes);
+            string source = MetadataRegistryEmitter.EmitRegistry(assemblyName, payload.Metadata, payload.Classes);
             ctx.AddSource("MSTestReflectionMetadata.Registry.g.cs", SourceText.From(source, Encoding.UTF8));
         });
     }

src/Analyzers/MSTest.AotReflection.SourceGeneration/Generators/MetadataRegistryEmitter.cs

@@ -32,6 +32,7 @@ public static string EmitSupportTypes()
 
         sb.AppendLine("using System;");
         sb.AppendLine("using System.Collections.Generic;");
+        sb.AppendLine("using System.Threading.Tasks;");
         sb.AppendLine();
 
         using (sb.Block($"namespace {GeneratedNamespace}"))
@@ -57,8 +58,10 @@ public static string EmitSupportTypes()
                 sb.AppendLine("public Type[] ParameterTypes { get; init; } = Array.Empty<Type>();");
                 sb.AppendLine("public string[] ParameterNames { get; init; } = Array.Empty<string>();");
                 sb.AppendLine("public Attribute[] Attributes { get; init; } = Array.Empty<Attribute>();");
-                sb.AppendLine("/// <summary>Direct invoker — replaces <see cref=\"System.Reflection.MethodInfo.Invoke(object, object[])\" />.</summary>");
-                sb.AppendLine("public Func<object?, object?[]?, object?> Invoke { get; init; } = static (_, _) => null;");
+                sb.AppendLine("/// <summary>Materialized argument tuples from <c>[DataRow]</c> attributes (empty for non-data-driven tests). Each <c>object?[]</c> corresponds to one <c>[DataRow]</c> application.</summary>");
+                sb.AppendLine("public IReadOnlyList<object?[]> DataRows { get; init; } = Array.Empty<object?[]>();");
+                sb.AppendLine("/// <summary>Direct invoker — replaces <see cref=\"System.Reflection.MethodInfo.Invoke(object, object[])\" />. Always returns a non-null <see cref=\"Task\" /> so the caller can <c>await</c> regardless of whether the underlying test method is <c>void</c>, <c>Task</c>, <c>Task&lt;T&gt;</c>, <c>ValueTask</c>, or <c>ValueTask&lt;T&gt;</c>; the result value (if any) is discarded.</summary>");
+                sb.AppendLine("public Func<object?, object?[]?, Task> Invoke { get; init; } = static (_, _) => Task.CompletedTask;");
             }
 
             sb.AppendLine();
@@ -83,13 +86,14 @@ public static string EmitSupportTypes()
         return sb.ToString();
     }
 
-    public static string EmitRegistry(string assemblyName, IReadOnlyList<TestClassModel> testClasses)
+    public static string EmitRegistry(string assemblyName, AssemblyMetadataModel assemblyMetadata, IReadOnlyList<TestClassModel> testClasses)
     {
         var sb = new IndentedStringBuilder();
         AppendHeader(sb);
 
         sb.AppendLine("using System;");
         sb.AppendLine("using System.Collections.Generic;");
+        sb.AppendLine("using System.Threading.Tasks;");
         sb.AppendLine();
 
         using (sb.Block($"namespace {GeneratedNamespace}"))
@@ -99,6 +103,12 @@ public static string EmitRegistry(string assemblyName, IReadOnlyList<TestClassMo
             {
                 sb.AppendLine($"public const string AssemblyName = \"{Escape(assemblyName)}\";");
                 sb.AppendLine();
+
+                // Emit assembly-level [assembly: ...] attributes so the consumer never has to call
+                // Assembly.GetCustomAttributes for attributes declared in the same compilation.
+                EmitAssemblyAttributesProperty(sb, assemblyMetadata.Attributes);
+                sb.AppendLine();
+
                 sb.AppendLine("public static IReadOnlyList<TestClassReflectionInfo> TestClasses { get; } = new TestClassReflectionInfo[]");
                 using (sb.Block(null))
                 {
@@ -119,6 +129,35 @@ public static string EmitRegistry(string assemblyName, IReadOnlyList<TestClassMo
         return sb.ToString();
     }
 
+    private static void EmitAssemblyAttributesProperty(IndentedStringBuilder sb, EquatableArray<AttributeApplicationModel> attributes)
+    {
+        if (attributes.Length == 0)
+        {
+            sb.AppendLine("public static IReadOnlyList<Attribute> AssemblyAttributes { get; } = Array.Empty<Attribute>();");
+            return;
+        }
+
+        sb.AppendLine("public static IReadOnlyList<Attribute> AssemblyAttributes { get; } = new Attribute[]");
+        using (sb.Block(null))
+        {
+            for (int i = 0; i < attributes.Length; i++)
+            {
+                AttributeApplicationModel attr = attributes[i];
+                sb.Append(BuildAttributeExpression(attr));
+                if (i < attributes.Length - 1)
+                {
+                    sb.AppendLine(",");
+                }
+                else
+                {
+                    sb.AppendLine();
+                }
+            }
+        }
+
+        sb.AppendLine(";");
+    }
+
     private static void EmitTestClass(IndentedStringBuilder sb, TestClassModel model)
     {
         string fqn = model.FullyQualifiedTypeName;
@@ -190,6 +229,8 @@ private static void EmitMethods(IndentedStringBuilder sb, string fqn, TestClassM
                     sb.AppendLine(",");
                     EmitAttributesProperty(sb, "Attributes", method.Attributes);
                     sb.AppendLine(",");
+                    EmitDataRows(sb, method.DataRows);
+                    sb.AppendLine(",");
                     EmitMethodInvoker(sb, fqn, method);
                 }
 
@@ -238,13 +279,76 @@ private static void EmitMethodInvoker(IndentedStringBuilder sb, string classFqn,
         string args = BuildArgumentsFromObjectArray(method.Parameters);
         string call = $"{target}.{method.Name}({args})";
 
-        string body = method.ReturnsVoid
-            ? $"{{ {call}; return null; }}"
-            : $"{{ return {call}; }}";
+        // The contract is: return a non-null Task representing the (async or sync) completion of the
+        // test method, discarding any result value. This lets the caller use a single `await invoker(...)`
+        // path regardless of the underlying return shape.
+        //   - void / non-Task sync: invoke, return Task.CompletedTask.
+        //   - Task / Task<T>: forward the returned Task (treat a `null` return as success).
+        //   - ValueTask / ValueTask<T>: avoid Task allocation for the synchronously-completed fast path
+        //     via IsCompletedSuccessfully, otherwise call AsTask().
+        string body;
+        if (method.ReturnsTask)
+        {
+            // Task<T> derives from Task, so the same forwarding code handles both. A test method that
+            // *declares* a Task return type and then returns `null` is broken at runtime, but mirroring
+            // reflection-Invoke we tolerate it and treat it as already-completed.
+            body = $"{{ Task? __t = {call}; return __t ?? Task.CompletedTask; }}";
+        }
+        else if (method.ReturnsValueTask)
+        {
+            body = $"{{ var __vt = {call}; return __vt.IsCompletedSuccessfully ? Task.CompletedTask : __vt.AsTask(); }}";
+        }
+        else if (method.ReturnsVoid)
+        {
+            body = $"{{ {call}; return Task.CompletedTask; }}";
+        }
+        else
+        {
+            // Non-void, non-Task return (e.g. `int Test()`). The test runner discards the value; we still
+            // execute the call for its side effects and report success.
+            body = $"{{ _ = {call}; return Task.CompletedTask; }}";
+        }
 
         sb.AppendLine($"Invoke = static (instance, args) => {body},");
     }
 
+    private static void EmitDataRows(IndentedStringBuilder sb, EquatableArray<DataRowModel> dataRows)
+    {
+        if (dataRows.Length == 0)
+        {
+            sb.Append("DataRows = Array.Empty<object?[]>()");
+            sb.AppendLine();
+            return;
+        }
+
+        sb.AppendLine("DataRows = new object?[][]");
+        using (sb.Block(null))
+        {
+            for (int i = 0; i < dataRows.Length; i++)
+            {
+                EquatableArray<TypedConstantModel> args = dataRows[i].Arguments;
+                if (args.Length == 0)
+                {
+                    sb.Append("Array.Empty<object?>()");
+                }
+                else
+                {
+                    string literals = string.Join(", ", args.AsImmutableArray().Select(BuildConstantExpression));
+                    sb.Append($"new object?[] {{ {literals} }}");
+                }
+
+                if (i < dataRows.Length - 1)
+                {
+                    sb.AppendLine(",");
+                }
+                else
+                {
+                    sb.AppendLine();
+                }
+            }
+        }
+    }
+
     private static void EmitParameterTypes(IndentedStringBuilder sb, EquatableArray<TestParameterModel> parameters)
     {
         if (parameters.Length == 0)