AotReflection: capture assembly-level attributes

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

@@ -83,7 +83,7 @@ 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);
@@ -99,6 +99,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 +125,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;

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

@@ -1,9 +1,11 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
+using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Linq;
+using System.Text;
 
 using Microsoft.CodeAnalysis;
 
@@ -19,31 +21,61 @@ internal static class TestClassModelBuilder
 {
     private static readonly SymbolDisplayFormat FullyQualifiedFormat =
         SymbolDisplayFormat.FullyQualifiedFormat.WithMiscellaneousOptions(
-            SymbolDisplayMiscellaneousOptions.IncludeNullableReferenceTypeModifier
-            | SymbolDisplayMiscellaneousOptions.UseSpecialTypes);
+            SymbolDisplayMiscellaneousOptions.UseSpecialTypes);
 
     public static TestClassModel Build(INamedTypeSymbol typeSymbol)
     {
+        // Methods / properties are walked across the full inheritance chain (excluding
+        // System.Object) so that MSTest members declared on a base class —
+        // [ClassInitialize], [ClassCleanup], [TestInitialize], [TestCleanup],
+        // [TestMethod], the [TestContext] setter, … — are visible to the consumer
+        // without runtime reflection.
+        //
+        // Iteration order is derived-first so that an override or `new`-shadowed member
+        // on the derived type wins over the base declaration with the same signature.
+        // Constructors are NEVER inherited and are taken only from the leaf type.
+        var methodsByKey = new Dictionary<string, TestMethodModel>(StringComparer.Ordinal);
+        var propertiesByName = new Dictionary<string, TestPropertyModel>(StringComparer.Ordinal);
         ImmutableArray<TestMethodModel>.Builder methods = ImmutableArray.CreateBuilder<TestMethodModel>();
         ImmutableArray<TestPropertyModel>.Builder properties = ImmutableArray.CreateBuilder<TestPropertyModel>();
         ImmutableArray<TestConstructorModel>.Builder ctors = ImmutableArray.CreateBuilder<TestConstructorModel>();
 
-        foreach (ISymbol member in typeSymbol.GetMembers())
+        for (INamedTypeSymbol? current = typeSymbol;
+             current is not null && current.SpecialType != SpecialType.System_Object;
+             current = current.BaseType)
         {
-            switch (member)
+            bool isLeaf = SymbolEqualityComparer.Default.Equals(current, typeSymbol);
+
+            foreach (ISymbol member in current.GetMembers())
             {
-                case IMethodSymbol { MethodKind: MethodKind.Ordinary } method
-                    when method.DeclaredAccessibility is Accessibility.Public or Accessibility.Internal:
-                    methods.Add(BuildMethod(method));
-                    break;
-                case IPropertySymbol property
-                    when property.DeclaredAccessibility is Accessibility.Public or Accessibility.Internal:
-                    properties.Add(BuildProperty(property));
-                    break;
-                case IMethodSymbol { MethodKind: MethodKind.Constructor, IsStatic: false } ctor
-                    when ctor.DeclaredAccessibility is Accessibility.Public or Accessibility.Internal:
-                    ctors.Add(new TestConstructorModel(BuildParameters(ctor)));
-                    break;
+                switch (member)
+                {
+                    case IMethodSymbol { MethodKind: MethodKind.Ordinary } method
+                        when IsAccessibleFromConsumer(method):
+                        string key = BuildMethodSignatureKey(method);
+                        if (!methodsByKey.ContainsKey(key))
+                        {
+                            TestMethodModel model = BuildMethod(method);
+                            methodsByKey[key] = model;
+                            methods.Add(model);
+                        }
+
+                        break;
+                    case IPropertySymbol property
+                        when IsAccessibleFromConsumer(property):
+                        if (!propertiesByName.ContainsKey(property.Name))
+                        {
+                            TestPropertyModel model = BuildProperty(property);
+                            propertiesByName[property.Name] = model;
+                            properties.Add(model);
+                        }
+
+                        break;
+                    case IMethodSymbol { MethodKind: MethodKind.Constructor, IsStatic: false } ctor
+                        when isLeaf && ctor.DeclaredAccessibility is Accessibility.Public or Accessibility.Internal:
+                        ctors.Add(new TestConstructorModel(BuildParameters(ctor)));
+                        break;
+                }
             }
         }
 
@@ -61,6 +93,58 @@ public static TestClassModel Build(INamedTypeSymbol typeSymbol)
             Attributes: BuildAttributes(typeSymbol.GetAttributes()));
     }
 
+    // Restricted to accessibilities the emitted helper class (a separate static type
+    // declared in MSTest.SourceGenerated, not a derived type) can legally call.
+    // 'protected' and 'private protected' members require the caller to be a derived
+    // type, so they are excluded; 'protected internal' is included because the internal
+    // half is satisfied (the generated helper lives in the same assembly).
+    private static bool IsAccessibleFromConsumer(ISymbol symbol)
+        => symbol.DeclaredAccessibility is
+            Accessibility.Public
+            or Accessibility.Internal
+            or Accessibility.ProtectedOrInternal;
+
+    private static string BuildMethodSignatureKey(IMethodSymbol method)
+    {
+        var sb = new StringBuilder();
+        sb.Append(method.IsStatic ? "S:" : "I:");
+        sb.Append(method.Name);
+        if (method.Arity > 0)
+        {
+            sb.Append('`');
+            sb.Append(method.Arity);
+        }
+
+        sb.Append('(');
+        bool first = true;
+        foreach (IParameterSymbol p in method.Parameters)
+        {
+            if (!first)
+            {
+                sb.Append(',');
+            }
+
+            first = false;
+            switch (p.RefKind)
+            {
+                case RefKind.Ref:
+                    sb.Append("ref ");
+                    break;
+                case RefKind.Out:
+                    sb.Append("out ");
+                    break;
+                case RefKind.In:
+                    sb.Append("in ");
+                    break;
+            }
+
+            sb.Append(p.Type.ToDisplayString(FullyQualifiedFormat));
+        }
+
+        sb.Append(')');
+        return sb.ToString();
+    }
+
     private static TestMethodModel BuildMethod(IMethodSymbol method)
     {
         ITypeSymbol returnType = method.ReturnType;
@@ -82,15 +166,131 @@ private static TestMethodModel BuildMethod(IMethodSymbol method)
             ReturnsValueTask: returnsValueTask,
             ReturnsVoid: returnsVoid,
             Parameters: BuildParameters(method),
-            Attributes: BuildAttributes(method.GetAttributes()));
+            Attributes: BuildAttributes(CollectInheritedAttributes(method)));
     }
 
     private static TestPropertyModel BuildProperty(IPropertySymbol property)
         => new(
             Name: property.Name,
             FullyQualifiedType: property.Type.ToDisplayString(FullyQualifiedFormat),
             HasPublicSetter: property.SetMethod is { DeclaredAccessibility: Accessibility.Public },
-            Attributes: BuildAttributes(property.GetAttributes()));
+            Attributes: BuildAttributes(CollectInheritedAttributes(property)));
+
+    // Mirror the runtime behavior of MemberInfo.GetCustomAttributes(inherit: true): walk the
+    // overridden-method chain and union attributes, keeping the most-derived application when
+    // the same attribute type appears on multiple levels — but respect
+    // [AttributeUsage(Inherited = false)] (the attribute is NOT visible past the level it was
+    // declared on) and [AttributeUsage(AllowMultiple = true)] (every occurrence is kept).
+    private static ImmutableArray<AttributeData> CollectInheritedAttributes(IMethodSymbol method)
+    {
+        ImmutableArray<AttributeData> own = method.GetAttributes();
+        if (method.OverriddenMethod is null)
+        {
+            return own;
+        }
+
+        var seen = new HashSet<string>(StringComparer.Ordinal);
+        ImmutableArray<AttributeData>.Builder builder = ImmutableArray.CreateBuilder<AttributeData>();
+        AppendUnique(builder, seen, own, isInheritedLevel: false);
+        for (IMethodSymbol? baseMethod = method.OverriddenMethod; baseMethod is not null; baseMethod = baseMethod.OverriddenMethod)
+        {
+            AppendUnique(builder, seen, baseMethod.GetAttributes(), isInheritedLevel: true);
+        }
+
+        return builder.ToImmutable();
+    }
+
+    private static ImmutableArray<AttributeData> CollectInheritedAttributes(IPropertySymbol property)
+    {
+        ImmutableArray<AttributeData> own = property.GetAttributes();
+        if (property.OverriddenProperty is null)
+        {
+            return own;
+        }
+
+        var seen = new HashSet<string>(StringComparer.Ordinal);
+        ImmutableArray<AttributeData>.Builder builder = ImmutableArray.CreateBuilder<AttributeData>();
+        AppendUnique(builder, seen, own, isInheritedLevel: false);
+        for (IPropertySymbol? baseProperty = property.OverriddenProperty; baseProperty is not null; baseProperty = baseProperty.OverriddenProperty)
+        {
+            AppendUnique(builder, seen, baseProperty.GetAttributes(), isInheritedLevel: true);
+        }
+
+        return builder.ToImmutable();
+    }
+
+    private static void AppendUnique(
+        ImmutableArray<AttributeData>.Builder builder,
+        HashSet<string> seen,
+        ImmutableArray<AttributeData> attributes,
+        bool isInheritedLevel)
+    {
+        foreach (AttributeData attribute in attributes)
+        {
+            if (attribute.AttributeClass is not { } attributeClass)
+            {
+                continue;
+            }
+
+            (bool allowMultiple, bool inherited) = GetAttributeUsage(attributeClass);
+
+            // A base-level attribute declared with AttributeUsage(Inherited = false) must
+            // not leak onto the derived override (matches MemberInfo.GetCustomAttributes(inherit: true)).
+            if (isInheritedLevel && !inherited)
+            {
+                continue;
+            }
+
+            // Attributes declared with AttributeUsage(AllowMultiple = true) may legitimately
+            // appear several times across the override chain (e.g. [TestCategory]) — keep every
+            // instance instead of collapsing them to one.
+            if (allowMultiple)
+            {
+                builder.Add(attribute);
+                continue;
+            }
+
+            string key = attributeClass.ToDisplayString(FullyQualifiedFormat);
+            if (seen.Add(key))
+            {
+                builder.Add(attribute);
+            }
+        }
+    }
+
+    private static (bool AllowMultiple, bool Inherited) GetAttributeUsage(INamedTypeSymbol attributeClass)
+    {
+        for (INamedTypeSymbol? current = attributeClass; current is not null; current = current.BaseType)
+        {
+            foreach (AttributeData attribute in current.GetAttributes())
+            {
+                if (attribute.AttributeClass?.ToDisplayString(FullyQualifiedFormat) != "global::System.AttributeUsageAttribute")
+                {
+                    continue;
+                }
+
+                bool allowMultiple = false;
+                bool inherited = true;
+                foreach (KeyValuePair<string, TypedConstant> named in attribute.NamedArguments)
+                {
+                    if (named.Key == "AllowMultiple" && named.Value.Value is bool am)
+                    {
+                        allowMultiple = am;
+                    }
+                    else if (named.Key == "Inherited" && named.Value.Value is bool inh)
+                    {
+                        inherited = inh;
+                    }
+                }
+
+                // [AttributeUsage] on a derived attribute class shadows the base per CLI rules;
+                // stop at the first level where it is found.
+                return (allowMultiple, inherited);
+            }
+        }
+
+        return (AllowMultiple: false, Inherited: true);
+    }
 
     private static EquatableArray<TestParameterModel> BuildParameters(IMethodSymbol method)
     {