Address review: respect AttributeUsage and align stub with real MSTest

Evangelink · Copilot · Evangelink · commit 29e18a6aeff5 · 2026-06-10T17:48:45.000+02:00
- IsAccessibleFromConsumer now excludes Protected / ProtectedAndInternal.
- BuildMethodSignatureKey encodes the method's generic arity.
- CollectInheritedAttributes respects AttributeUsage(Inherited = false) on
  inherited levels and preserves every instance of AllowMultiple = true
  attributes (e.g. [TestCategory], [DataRow]).
- MinimalMSTestStub TestMethodAttribute and DataRowAttribute now declare
  AttributeUsage(Inherited = false), matching the real MSTest attribute
  definitions so inheritance behavior is exercised correctly.
- Adds regression tests covering protected exclusion, generic arity,
  AllowMultiple preservation, and Inherited = false propagation rules.

Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;
diff --git a/src/Analyzers/MSTest.AotReflection.SourceGeneration/Generators/TestClassModelBuilder.cs b/src/Analyzers/MSTest.AotReflection.SourceGeneration/Generators/TestClassModelBuilder.cs
@@ -94,19 +94,28 @@ when IsAccessibleFromConsumer(property):
             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.Protected
-            or Accessibility.ProtectedOrInternal
-            or Accessibility.ProtectedAndInternal;
+            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)
@@ -228,7 +237,9 @@ private static TestPropertyModel BuildProperty(IPropertySymbol 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.
+    // 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();
@@ -239,10 +250,10 @@ private static ImmutableArray<AttributeData> CollectInheritedAttributes(IMethodS
 
         var seen = new HashSet<string>(StringComparer.Ordinal);
         ImmutableArray<AttributeData>.Builder builder = ImmutableArray.CreateBuilder<AttributeData>();
-        AppendUnique(builder, seen, own);
+        AppendUnique(builder, seen, own, isInheritedLevel: false);
         for (IMethodSymbol? baseMethod = method.OverriddenMethod; baseMethod is not null; baseMethod = baseMethod.OverriddenMethod)
         {
-            AppendUnique(builder, seen, baseMethod.GetAttributes());
+            AppendUnique(builder, seen, baseMethod.GetAttributes(), isInheritedLevel: true);
         }
 
         return builder.ToImmutable();
@@ -258,10 +269,10 @@ private static ImmutableArray<AttributeData> CollectInheritedAttributes(IPropert
 
         var seen = new HashSet<string>(StringComparer.Ordinal);
         ImmutableArray<AttributeData>.Builder builder = ImmutableArray.CreateBuilder<AttributeData>();
-        AppendUnique(builder, seen, own);
+        AppendUnique(builder, seen, own, isInheritedLevel: false);
         for (IPropertySymbol? baseProperty = property.OverriddenProperty; baseProperty is not null; baseProperty = baseProperty.OverriddenProperty)
         {
-            AppendUnique(builder, seen, baseProperty.GetAttributes());
+            AppendUnique(builder, seen, baseProperty.GetAttributes(), isInheritedLevel: true);
         }
 
         return builder.ToImmutable();
@@ -270,7 +281,8 @@ private static ImmutableArray<AttributeData> CollectInheritedAttributes(IPropert
     private static void AppendUnique(
         ImmutableArray<AttributeData>.Builder builder,
         HashSet<string> seen,
-        ImmutableArray<AttributeData> attributes)
+        ImmutableArray<AttributeData> attributes,
+        bool isInheritedLevel)
     {
         foreach (AttributeData attribute in attributes)
         {
@@ -279,6 +291,24 @@ private static void AppendUnique(
                 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], [DataRow])
+            // — keep every instance instead of collapsing them to one.
+            if (allowMultiple)
+            {
+                builder.Add(attribute);
+                continue;
+            }
+
             string key = attributeClass.ToDisplayString(FullyQualifiedFormat);
             if (seen.Add(key))
             {
@@ -287,6 +317,40 @@ private static void AppendUnique(
         }
     }
 
+    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)
     {
         if (method.Parameters.IsDefaultOrEmpty)
diff --git a/test/UnitTests/MSTest.AotReflection.SourceGeneration.UnitTests/MSTestReflectionMetadataGeneratorTests.cs b/test/UnitTests/MSTest.AotReflection.SourceGeneration.UnitTests/MSTestReflectionMetadataGeneratorTests.cs
@@ -28,7 +28,7 @@ namespace Microsoft.VisualStudio.TestTools.UnitTesting
             [System.AttributeUsage(System.AttributeTargets.Class)]
             public class TestClassAttribute : System.Attribute { }
 
-            [System.AttributeUsage(System.AttributeTargets.Method)]
+            [System.AttributeUsage(System.AttributeTargets.Method, Inherited = false)]
             public class TestMethodAttribute : System.Attribute
             {
                 public TestMethodAttribute() { }
@@ -59,7 +59,7 @@ public class ParallelizeAttribute : System.Attribute
                 public string? Scope { get; set; }
             }
 
-            [System.AttributeUsage(System.AttributeTargets.Method, AllowMultiple = true)]
+            [System.AttributeUsage(System.AttributeTargets.Method, AllowMultiple = true, Inherited = false)]
             public class DataRowAttribute : System.Attribute
             {
                 public DataRowAttribute(object? data1) { }
@@ -571,6 +571,9 @@ public virtual void Run() { }
                 [TestClass]
                 public class DerivedTests : BaseTests
                 {
+                    // [TestMethod] is re-applied here because the real attribute is declared
+                    // with AttributeUsage(Inherited = false) and would not be inherited.
+                    [TestMethod]
                     public override void Run() { }
                 }
             }
@@ -583,9 +586,6 @@ public override void Run() { }
         runEntries.Should().Be(1, "the derived override must replace the base entry (not duplicate it)");
         registry.Should().Contain("((global::Sample.DerivedTests)instance!).Run();");
         registry.Should().NotContain("((global::Sample.BaseTests)instance!).Run();");
-
-        // The override does NOT re-apply [TestMethod] but the attribute must still be visible
-        // via the override chain — matching the runtime semantics of GetCustomAttributes(inherit: true).
         registry.Should().Contain("global::Microsoft.VisualStudio.TestTools.UnitTesting.TestMethodAttribute");
     }
 
@@ -1032,6 +1032,162 @@ public void Test(string? value) { }
         registry.Should().Contain("new object?[] { (object)null! }");
     }
 
+    [TestMethod]
+    public void Generator_SkipsProtectedAndPrivateProtectedMembers()
+    {
+        // Generated invokers live in a separate static helper class (not a derived type),
+        // so 'protected' and 'private protected' members are not callable from the emitted
+        // code. They MUST be excluded from the registry to keep the generated source compiling.
+        const string userCode = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            namespace Sample
+            {
+                public class BaseTests
+                {
+                    [TestMethod]
+                    protected void ProtectedTest() { }
+
+                    [TestMethod]
+                    private protected void PrivateProtectedTest() { }
+                }
+
+                [TestClass]
+                public class DerivedTests : BaseTests
+                {
+                    [TestMethod]
+                    public void PublicTest() { }
+                }
+            }
+            """;
+
+        Compilation outputCompilation = RunGeneratorAndGetCompilation(MinimalMSTestStub, userCode);
+        IEnumerable<Diagnostic> errors = outputCompilation
+            .GetDiagnostics()
+            .Where(d => d.Severity == DiagnosticSeverity.Error);
+        errors.Should().BeEmpty("the generated registry MUST NOT reference protected / private protected members");
+
+        string registry = outputCompilation
+            .SyntaxTrees
+            .Single(t => t.FilePath.EndsWith("MSTestReflectionMetadata.Registry.g.cs", System.StringComparison.Ordinal))
+            .ToString();
+
+        registry.Should().Contain("Name = \"PublicTest\"");
+        registry.Should().NotContain("Name = \"ProtectedTest\"");
+        registry.Should().NotContain("Name = \"PrivateProtectedTest\"");
+    }
+
+    [TestMethod]
+    public void Generator_KeepsAllowMultipleAttributes_AcrossOverrideChain()
+    {
+        // [TestCategory] is AllowMultiple=true. Collecting attributes across the override chain
+        // MUST keep every instance instead of collapsing them by type, otherwise the inherited
+        // categories disappear from the registry.
+        const string userCode = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            namespace Sample
+            {
+                public class BaseTests
+                {
+                    [TestMethod]
+                    [TestCategory("BaseCat")]
+                    public virtual void Run() { }
+                }
+
+                [TestClass]
+                public class DerivedTests : BaseTests
+                {
+                    [TestCategory("DerivedCat")]
+                    public override void Run() { }
+                }
+            }
+            """;
+
+        string registry = GetRegistry(RunGenerator(MinimalMSTestStub, userCode));
+
+        registry.Should().Contain("\"BaseCat\"");
+        registry.Should().Contain("\"DerivedCat\"");
+    }
+
+    [TestMethod]
+    public void Generator_DistinguishesGenericArity_BetweenSameNamedMethods()
+    {
+        // Methods that differ only in generic arity (e.g. M() vs M<T>()) MUST be treated as
+        // distinct in the per-class dedup key, otherwise the generator drops one of them.
+        const string userCode = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            namespace Sample
+            {
+                [TestClass]
+                public class GenericArityTests
+                {
+                    [TestMethod]
+                    public void Run() { }
+
+                    [TestMethod]
+                    public void Run<T>() { }
+
+                    [TestMethod]
+                    public void Run<T, U>() { }
+                }
+            }
+            """;
+
+        Compilation outputCompilation = RunGeneratorAndGetCompilation(MinimalMSTestStub, userCode);
+        IEnumerable<Diagnostic> errors = outputCompilation
+            .GetDiagnostics()
+            .Where(d => d.Severity == DiagnosticSeverity.Error);
+        errors.Should().BeEmpty();
+
+        string registry = outputCompilation
+            .SyntaxTrees
+            .Single(t => t.FilePath.EndsWith("MSTestReflectionMetadata.Registry.g.cs", System.StringComparison.Ordinal))
+            .ToString();
+
+        int occurrences = System.Text.RegularExpressions.Regex
+            .Matches(registry, "Name = \"Run\"")
+            .Count;
+        occurrences.Should().Be(3);
+    }
+
+    [TestMethod]
+    public void Generator_DoesNotInherit_TestMethodOrDataRow_OntoOverride()
+    {
+        // [TestMethod] and [DataRow] are declared with AttributeUsage(Inherited = false) in
+        // the real MSTest assembly. An override that does NOT re-apply [TestMethod] is not a
+        // test, and inherited [DataRow]s must not leak from the base method onto the override.
+        const string userCode = """
+            using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+            namespace Sample
+            {
+                public class BaseTests
+                {
+                    [TestMethod]
+                    [DataRow(1)]
+                    [DataRow(2)]
+                    public virtual void Run(int x) { }
+                }
+
+                [TestClass]
+                public class DerivedTests : BaseTests
+                {
+                    // Override deliberately does not re-apply [TestMethod] or any [DataRow].
+                    public override void Run(int x) { }
+                }
+            }
+            """;
+
+        string registry = GetRegistry(RunGenerator(MinimalMSTestStub, userCode));
+
+        // Neither the [TestMethod] nor the inherited [DataRow]s should propagate onto the override.
+        registry.Should().NotContain("Microsoft.VisualStudio.TestTools.UnitTesting.TestMethodAttribute()");
+        registry.Should().NotContain("Microsoft.VisualStudio.TestTools.UnitTesting.DataRowAttribute(");
+        registry.Should().NotContain("new DataRowModel(");
+    }
+
     private static string GetRegistry(GeneratorRunResult result)
         => result.GeneratedSources
             .Single(s => s.HintName == "MSTestReflectionMetadata.Registry.g.cs")
