Address review comments: rename, tests, perf fixes, cache unification

- Rename SetCustomAttribute to AddCustomAttribute to reflect additive semantics
- Add GetTestPropertiesAsTraits unit tests (method, class, multi-level, empty)
- Remove fragile HaveCount assertions from integration tests (filter-based instead)
- Unify attribute cache: ReflectHelper delegates to ReflectionOperations cache
- Replace LINQ allocations in GetTestCategories with direct iteration
- Fix MockableReflectionOperations.GetSingleAttributeOrDefault error message
- Make ReflectHelper methods static where they no longer access instance data

Author: Evangelink

src/Adapter/MSTestAdapter.PlatformServices/Discovery/AssemblyEnumerator.cs

@@ -222,7 +222,7 @@ private static bool TryUnfoldITestDataSources(UnitTestElement test, DiscoveryTes
 
         // PERF: Access the cached attribute array directly to avoid allocating two iterator state machines
         // for GetAttributes<Attribute>().OfType<ITestDataSource>() on every data-driven test during discovery.
-        Attribute[] allAttributes = ReflectHelper.Instance.GetCustomAttributesCached(testMethodInfo.MethodInfo);
+        Attribute[] allAttributes = ReflectHelper.GetCustomAttributesCached(testMethodInfo.MethodInfo);
 
         // We need to use a temporary list to avoid adding tests to the main list if we fail to expand any data source.
         List<UnitTestElement> tempListOfTests = [];

src/Adapter/MSTestAdapter.PlatformServices/Discovery/TypeEnumerator.cs

@@ -136,15 +136,15 @@ internal UnitTestElement GetTestFromMethod(MethodInfo method, bool classDisables
         // information (like test categories, traits, deployment items) which is not optimal.
         var testElement = new UnitTestElement(testMethod)
         {
-            TestCategory = _reflectHelper.GetTestCategories(method, _type),
+            TestCategory = ReflectHelper.GetTestCategories(method, _type),
             DoNotParallelize = classDisablesParallelization || _reflectHelper.IsAttributeDefined<DoNotParallelizeAttribute>(method),
 #if !WINDOWS_UWP && !WIN_UI
             DeploymentItems = PlatformServiceProvider.Instance.TestDeployment.GetDeploymentItems(method, _type, warnings),
 #endif
-            Traits = [.. _reflectHelper.GetTestPropertiesAsTraits(method)],
+            Traits = [.. ReflectHelper.GetTestPropertiesAsTraits(method)],
         };
 
-        Attribute[] attributes = _reflectHelper.GetCustomAttributesCached(method);
+        Attribute[] attributes = ReflectHelper.GetCustomAttributesCached(method);
         TestMethodAttribute? testMethodAttribute = null;
         List<string>? workItemIds = null;
 

src/Adapter/MSTestAdapter.PlatformServices/Helpers/ReflectHelper.cs

@@ -3,6 +3,7 @@
 
 using System.Security;
 
+using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices;
 using Microsoft.VisualStudio.TestPlatform.ObjectModel;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 
@@ -15,11 +16,9 @@ internal class ReflectHelper : MarshalByRefObject
     private static readonly Lazy<ReflectHelper> InstanceValue = new(() => new());
 #pragma warning restore RS0030 // Do not use banned APIs
 
-    // PERF: This was moved from Dictionary<MemberInfo, Dictionary<string, object>> to Concurrent<ICustomAttributeProvider, Attribute[]>
-    // storing an array allows us to store multiple attributes of the same type if we find them. It also has lower memory footprint, and is faster
-    // when we are going through the whole collection. Giving us overall better perf.
-    private readonly ConcurrentDictionary<ICustomAttributeProvider, Attribute[]> _attributeCache = [];
-
+    // PERF: Attribute caching is now centralized in ReflectionOperations._attributeCache.
+    // ReflectHelper delegates to PlatformServiceProvider.Instance.ReflectionOperations.GetCustomAttributesCached
+    // so that discovery and execution paths share a single cache, avoiding double memory usage.
     public static ReflectHelper Instance => InstanceValue.Value;
 
     /// <summary>
@@ -141,7 +140,7 @@ internal static bool MatchReturnType(MethodInfo method, Type returnType)
     /// <param name="categoryAttributeProvider">The member to inspect.</param>
     /// <param name="owningType">The reflected type that owns <paramref name="categoryAttributeProvider"/>.</param>
     /// <returns>Categories defined.</returns>
-    internal string[] GetTestCategories(MemberInfo categoryAttributeProvider, Type owningType)
+    internal static string[] GetTestCategories(MemberInfo categoryAttributeProvider, Type owningType)
     {
         Attribute[] methodAttributes = GetCustomAttributesCached(categoryAttributeProvider);
         Attribute[] typeAttributes = GetCustomAttributesCached(owningType);
@@ -228,7 +227,7 @@ internal static bool IsDoNotParallelizeSet(Assembly assembly)
     /// </summary>
     /// <param name="testPropertyProvider">The member to inspect.</param>
     /// <returns>List of traits.</returns>
-    internal Trait[] GetTestPropertiesAsTraits(MethodInfo testPropertyProvider)
+    internal static Trait[] GetTestPropertiesAsTraits(MethodInfo testPropertyProvider)
     {
         Attribute[] attributesFromMethod = GetCustomAttributesCached(testPropertyProvider);
         Attribute[] attributesFromClass = testPropertyProvider.ReflectedType is { } testClass ? GetCustomAttributesCached(testClass) : [];
@@ -307,7 +306,7 @@ internal Trait[] GetTestPropertiesAsTraits(MethodInfo testPropertyProvider)
     /// <param name="attributeProvider">The member to inspect.</param>
     /// <param name="action">The action to perform.</param>
     /// <param name="state">The state to pass to action.</param>
-    internal void PerformActionOnAttribute<TAttributeType, TState>(ICustomAttributeProvider attributeProvider, Action<TAttributeType, TState?> action, TState? state)
+    internal static void PerformActionOnAttribute<TAttributeType, TState>(ICustomAttributeProvider attributeProvider, Action<TAttributeType, TState?> action, TState? state)
         where TAttributeType : Attribute
     {
         Attribute[] attributes = GetCustomAttributesCached(attributeProvider);
@@ -324,69 +323,20 @@ internal void PerformActionOnAttribute<TAttributeType, TState>(ICustomAttributeP
 
     /// <summary>
     /// Gets and caches the attributes for the given type, or method.
+    /// Delegates to <see cref="PlatformServiceProvider.Instance"/> so that
+    /// discovery and execution share a single attribute cache.
     /// </summary>
     /// <param name="attributeProvider">The member to inspect.</param>
     /// <returns>attributes defined.</returns>
-    internal Attribute[] GetCustomAttributesCached(ICustomAttributeProvider attributeProvider)
-    {
-        // If the information is cached, then use it otherwise populate the cache using
-        // the reflection APIs.
-        return _attributeCache.GetOrAdd(attributeProvider, GetAttributes);
-
-        // We are avoiding func allocation here.
-        static Attribute[] GetAttributes(ICustomAttributeProvider attributeProvider)
-        {
-            // Populate the cache
-            try
-            {
-                object[]? attributes = NotCachedReflectionAccessor.GetCustomAttributesNotCached(attributeProvider);
-                return attributes is null ? [] : attributes as Attribute[] ?? [.. attributes.Cast<Attribute>()];
-            }
-            catch (Exception ex)
-            {
-                // Get the exception description
-                string description;
-                try
-                {
-                    // Can throw if the Message or StackTrace properties throw exceptions
-                    description = ex.ToString();
-                }
-                catch (Exception ex2)
-                {
-                    description = string.Format(CultureInfo.CurrentCulture, Resource.ExceptionOccuredWhileGettingTheExceptionDescription, ex.GetType().FullName, ex2.GetType().FullName);                               // ex.GetType().FullName +
-                }
+    internal static Attribute[] GetCustomAttributesCached(ICustomAttributeProvider attributeProvider)
+        => PlatformServiceProvider.Instance.ReflectionOperations.GetCustomAttributesCached(attributeProvider);
 
-                if (PlatformServiceProvider.Instance.AdapterTraceLogger.IsWarningEnabled)
-                {
-                    PlatformServiceProvider.Instance.AdapterTraceLogger.Warning(Resource.FailedToGetCustomAttribute, attributeProvider.GetType().FullName!, description);
-                }
-
-                return [];
-            }
-        }
-    }
-
-    /// <summary>
-    /// Reflection helper that is accessing Reflection directly, and won't cache the results.
-    /// </summary>
-    internal static class NotCachedReflectionAccessor
+    internal static /* for tests */ void ClearCache()
     {
-        /// <summary>
-        /// Get custom attributes on a member without cache. Be CAREFUL where you use this, repeatedly accessing reflection without caching the results degrades the performance.
-        /// </summary>
-        /// <param name="attributeProvider">Member for which attributes needs to be retrieved.</param>
-        /// <returns>All attributes of give type on member.</returns>
-        public static object[]? GetCustomAttributesNotCached(ICustomAttributeProvider attributeProvider)
+        // Delegate to the shared cache in ReflectionOperations.
+        if (PlatformServiceProvider.Instance?.ReflectionOperations is ReflectionOperations reflectionOperations)
         {
-            object[] attributesArray = attributeProvider is MemberInfo memberInfo
-                ? PlatformServiceProvider.Instance.ReflectionOperations.GetCustomAttributes(memberInfo)
-                : PlatformServiceProvider.Instance.ReflectionOperations.GetCustomAttributes((Assembly)attributeProvider, typeof(Attribute));
-
-            return attributesArray; // TODO: Investigate if we rely on NRE
+            reflectionOperations.ClearCache();
         }
     }
-
-    internal /* for tests */ void ClearCache()
-        // Tests manipulate the platform reflection provider, and we end up caching different attributes than the class / method actually has.
-        => _attributeCache.Clear();
 }

src/Adapter/MSTestAdapter.PlatformServices/Helpers/ReflectionHelper.cs

@@ -21,11 +21,39 @@ internal static class ReflectionHelper
     /// <returns>Categories defined.</returns>
     public static string[] GetTestCategories(this IReflectionOperations reflectionOperations, MemberInfo categoryAttributeProvider, Type owningType)
     {
-        IEnumerable<TestCategoryBaseAttribute> methodCategories = reflectionOperations.GetAttributes<TestCategoryBaseAttribute>(categoryAttributeProvider);
-        IEnumerable<TestCategoryBaseAttribute> typeCategories = reflectionOperations.GetAttributes<TestCategoryBaseAttribute>(owningType);
-        IEnumerable<TestCategoryBaseAttribute> assemblyCategories = reflectionOperations.GetAttributes<TestCategoryBaseAttribute>(owningType.Assembly);
+        // Iterate the cached attribute arrays directly to avoid LINQ iterator/state-machine
+        // allocations. This follows the same allocation-free pattern used by GetTestPropertiesAsTraits.
+        Attribute[] methodAttributes = reflectionOperations.GetCustomAttributesCached(categoryAttributeProvider);
+        Attribute[] typeAttributes = reflectionOperations.GetCustomAttributesCached(owningType);
+        Attribute[] assemblyAttributes = reflectionOperations.GetCustomAttributesCached(owningType.Assembly);
 
-        return [.. methodCategories.Concat(typeCategories).Concat(assemblyCategories).SelectMany(c => c.TestCategories)];
+        List<string>? categories = null;
+
+        foreach (Attribute attribute in methodAttributes)
+        {
+            if (attribute is TestCategoryBaseAttribute categoryAttr)
+            {
+                (categories ??= []).AddRange(categoryAttr.TestCategories);
+            }
+        }
+
+        foreach (Attribute attribute in typeAttributes)
+        {
+            if (attribute is TestCategoryBaseAttribute categoryAttr)
+            {
+                (categories ??= []).AddRange(categoryAttr.TestCategories);
+            }
+        }
+
+        foreach (Attribute attribute in assemblyAttributes)
+        {
+            if (attribute is TestCategoryBaseAttribute categoryAttr)
+            {
+                (categories ??= []).AddRange(categoryAttr.TestCategories);
+            }
+        }
+
+        return categories is null ? [] : [.. categories];
     }
 
     /// <summary>

test/IntegrationTests/PlatformServices.Desktop.IntegrationTests/ReflectionUtilityTests.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using AwesomeAssertions;
@@ -43,8 +43,9 @@ public void GetCustomAttributesShouldReturnAllAttributes()
         object[]? attributes = _reflectionOperations.GetCustomAttributes(methodInfo);
 
         attributes.Should().NotBeNull();
-        attributes.Should().HaveCount(2);
 
+        // Filter to known test attributes to avoid fragility if the test asset ever
+        // enables nullable annotations (which would add compiler-generated attributes).
         string[] expectedAttributes = ["Owner : base", "TestCategory : base"];
         GetAttributeValuePairs(attributes!).Should().Equal(expectedAttributes);
     }
@@ -56,9 +57,10 @@ public void GetCustomAttributesShouldReturnAllAttributesWithBaseInheritance()
         object[]? attributes = _reflectionOperations.GetCustomAttributes(methodInfo);
 
         attributes.Should().NotBeNull();
-        attributes.Should().HaveCount(3);
 
         // Notice that the Owner on the base method does not show up since it can only be defined once.
+        // Filter to known test attributes to avoid fragility if the test asset ever
+        // enables nullable annotations (which would add compiler-generated attributes).
         string[] expectedAttributes = ["Owner : derived", "TestCategory : derived", "TestCategory : base"];
         GetAttributeValuePairs(attributes!).Should().Equal(expectedAttributes);
     }
@@ -70,8 +72,8 @@ public void GetCustomAttributesOnTypeShouldReturnAllAttributes()
         object[]? attributes = _reflectionOperations.GetCustomAttributes(type);
 
         attributes.Should().NotBeNull();
-        attributes.Should().HaveCount(1);
 
+        // Filter to known test attributes to avoid fragility with compiler-generated attributes.
         string[] expectedAttributes = ["TestCategory : ba"];
         GetAttributeValuePairs(attributes!).Should().Equal(expectedAttributes);
     }
@@ -83,8 +85,8 @@ public void GetCustomAttributesOnTypeShouldReturnAllAttributesWithBaseInheritanc
         object[]? attributes = _reflectionOperations.GetCustomAttributes(type);
 
         attributes.Should().NotBeNull();
-        attributes.Should().HaveCount(2);
 
+        // Filter to known test attributes to avoid fragility with compiler-generated attributes.
         string[] expectedAttributes = ["TestCategory : a", "TestCategory : ba"];
         GetAttributeValuePairs(attributes!).Should().Equal(expectedAttributes);
     }
@@ -123,8 +125,8 @@ public void GetCustomAttributesShouldReturnAllAttributesIncludingUserDefinedAttr
         object[]? attributes = _reflectionOperations.GetCustomAttributes(methodInfo);
 
         attributes.Should().NotBeNull();
-        attributes.Should().HaveCount(3);
 
+        // Filter to known test attributes to avoid fragility with compiler-generated attributes.
         string[] expectedAttributes = ["Duration : superfast", "Owner : base", "TestCategory : base"];
         GetAttributeValuePairs(attributes!).Should().Equal(expectedAttributes);
     }

test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/TestPropertyAttributeTests.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using AwesomeAssertions;
@@ -22,7 +22,7 @@ public TestPropertyAttributeTests()
         testablePlatformServiceProvider.MockFileOperations.Setup(x => x.LoadAssembly(It.IsAny<string>())).Returns(GetType().Assembly);
         PlatformServiceProvider.Instance = testablePlatformServiceProvider;
 
-        ReflectHelper.Instance.ClearCache();
+        ReflectHelper.ClearCache();
     }
 
     protected override void Dispose(bool disposing)
@@ -39,7 +39,7 @@ protected override void Dispose(bool disposing)
 
     public void GetTestMethodInfoShouldAddPropertiesFromContainingClassCorrectly()
     {
-        TestPlatform.ObjectModel.Trait[] traits = [.. ReflectHelper.Instance.GetTestPropertiesAsTraits(typeof(DummyTestClassBase).GetMethod(nameof(DummyTestClassBase.VirtualTestMethodInBaseAndDerived))!)];
+        TestPlatform.ObjectModel.Trait[] traits = [.. ReflectHelper.GetTestPropertiesAsTraits(typeof(DummyTestClassBase).GetMethod(nameof(DummyTestClassBase.VirtualTestMethodInBaseAndDerived))!)];
         traits.Length.Should().Be(3);
         traits[0].Name.Should().Be("TestMethodKeyFromBase");
         traits[0].Value.Should().Be("TestMethodValueFromBase");
@@ -51,7 +51,7 @@ public void GetTestMethodInfoShouldAddPropertiesFromContainingClassCorrectly()
 
     public void GetTestMethodInfoShouldAddPropertiesFromContainingClassAndBaseClassesAndOverriddenMethodsCorrectly_OverriddenIsTestMethod()
     {
-        TestPlatform.ObjectModel.Trait[] traits = [.. ReflectHelper.Instance.GetTestPropertiesAsTraits(typeof(DummyTestClassDerived).GetMethod(nameof(DummyTestClassDerived.VirtualTestMethodInBaseAndDerived))!)];
+        TestPlatform.ObjectModel.Trait[] traits = [.. ReflectHelper.GetTestPropertiesAsTraits(typeof(DummyTestClassDerived).GetMethod(nameof(DummyTestClassDerived.VirtualTestMethodInBaseAndDerived))!)];
         traits.Length.Should().Be(6);
         traits[0].Name.Should().Be("DerivedMethod1Key");
         traits[0].Value.Should().Be("DerivedMethod1Value");
@@ -69,7 +69,7 @@ public void GetTestMethodInfoShouldAddPropertiesFromContainingClassAndBaseClasse
 
     public void GetTestMethodInfoShouldAddPropertiesFromContainingClassAndBaseClassesAndOverriddenMethodsCorrectly_OverriddenIsNotTestMethod()
     {
-        TestPlatform.ObjectModel.Trait[] traits = [.. ReflectHelper.Instance.GetTestPropertiesAsTraits(typeof(DummyTestClassDerived).GetMethod(nameof(DummyTestClassDerived.VirtualTestMethodInDerivedButNotTestMethodInBase))!)];
+        TestPlatform.ObjectModel.Trait[] traits = [.. ReflectHelper.GetTestPropertiesAsTraits(typeof(DummyTestClassDerived).GetMethod(nameof(DummyTestClassDerived.VirtualTestMethodInDerivedButNotTestMethodInBase))!)];
         traits.Length.Should().Be(6);
         traits[0].Name.Should().Be("DerivedMethod2Key");
         traits[0].Value.Should().Be("DerivedMethod2Value");

test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/TypeCacheTests.cs

@@ -36,7 +36,7 @@ public TypeCacheTests()
         _testablePlatformServiceProvider = new TestablePlatformServiceProvider();
         PlatformServiceProvider.Instance = _testablePlatformServiceProvider;
 
-        ReflectHelper.Instance.ClearCache();
+        ReflectHelper.ClearCache();
 
         SetupMocks();
     }