Add ITestFilter / [TestFilterProvider] for programmatic test filtering (#8894)

src/Adapter/MSTestAdapter.PlatformServices/Execution/TypeCache.FilterDiscovery.cs

@@ -0,0 +1,161 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+namespace Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Execution;
+
+internal sealed partial class TypeCache
+{
+    // Single filter instance cached per test assembly source path. Computed lazily on the first
+    // request for that source so the cost is paid at most once per run, even when many tests
+    // target the same assembly. Stored as a TestFilterBox so the dictionary can cache the
+    // "no filter" answer alongside real filter instances.
+    private readonly ConcurrentDictionary<string, TestFilterBox> _testFilterBySource =
+        new(StringComparer.Ordinal);
+
+    /// <summary>
+    /// Returns the cached <see cref="ITestFilter"/> instance registered via
+    /// <see cref="TestFilterProviderAttribute"/> on the given test assembly, or
+    /// <see langword="null"/> if the assembly does not register one.
+    /// </summary>
+    /// <param name="assemblySource">The test assembly source path (typically <c>TestMethod.AssemblyName</c>).</param>
+    /// <remarks>
+    /// Discovery is metadata-only for the probe step and never forces the test types of the
+    /// assembly to load. The filter <em>type</em> is loaded the first time the filter for a
+    /// given source is requested. Only the test assembly itself is inspected — registering a
+    /// <see cref="TestFilterProviderAttribute"/> in a referenced library has no effect.
+    /// </remarks>
+    internal ITestFilter? GetOrLoadTestFilter(string assemblySource)
+        => _testFilterBySource
+            .GetOrAdd(assemblySource, static src => new TestFilterBox(LoadTestFilterForSource(src)))
+            .Filter;
+
+    private static ITestFilter? LoadTestFilterForSource(string assemblySource)
+    {
+        Assembly assembly;
+        try
+        {
+            assembly = PlatformServiceProvider.Instance.FileOperations.LoadAssembly(assemblySource);
+        }
+        catch (Exception ex)
+        {
+            if (PlatformServiceProvider.Instance.AdapterTraceLogger.IsWarningEnabled)
+            {
+                PlatformServiceProvider.Instance.AdapterTraceLogger.Warning(
+                    "TypeCache: Could not load test assembly {0} for TestFilterProvider discovery. {1}",
+                    assemblySource,
+                    ex);
+            }
+
+            return null;
+        }
+
+        return DiscoverTestFilterFromProvider(assembly);
+    }
+
+    private static ITestFilter? DiscoverTestFilterFromProvider(Assembly testAssembly)
+    {
+        // Cheap metadata-only probe first: avoid loading the filter's Type unless the attribute is
+        // actually present. Mirrors the AssemblyFixtureProvider probe pattern.
+        if (!HasTestFilterProviderMarker(testAssembly))
+        {
+            return null;
+        }
+
+        object[] markers;
+        try
+        {
+            markers = PlatformServiceProvider.Instance.ReflectionOperations.GetCustomAttributes(testAssembly, typeof(TestFilterProviderAttribute));
+        }
+        catch (Exception ex)
+        {
+            // Marker is present (CustomAttributeData saw it) but the attribute cannot be
+            // instantiated. This typically means the type referenced by typeof(...) cannot be
+            // loaded. [TestFilterProvider] is explicit opt-in: silently dropping the marker
+            // would let the user's filter logic disappear at runtime, which is a more
+            // dangerous failure mode than a clear diagnostic.
+            string message = string.Format(
+                CultureInfo.CurrentCulture,
+                Resource.UTA_TestFilterProviderLoadFailed,
+                SafeGetAssemblyName(testAssembly) ?? "<unknown>",
+                ex.Message);
+            throw new TypeInspectionException(message, ex);
+        }
+
+        if (markers is null || markers.Length == 0)
+        {
+            return null;
+        }
+
+        if (markers.Length > 1)
+        {
+            string message = string.Format(
+                CultureInfo.CurrentCulture,
+                Resource.UTA_TestFilterProviderMultipleDeclared,
+                SafeGetAssemblyName(testAssembly) ?? "<unknown>");
+            throw new TypeInspectionException(message);
+        }
+
+        return markers[0] is TestFilterProviderAttribute { FilterType: { } filterType }
+            ? InstantiateTestFilter(filterType)
+            : null;
+    }
+
+    internal static ITestFilter InstantiateTestFilter(Type filterType)
+    {
+        if (filterType.IsGenericType)
+        {
+            string message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_TestFilterProviderTypeIsGeneric, filterType.FullName);
+            throw new TypeInspectionException(message);
+        }
+
+        if (filterType.IsAbstract || filterType.IsInterface)
+        {
+            string message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_TestFilterProviderTypeIsNotInstantiable, filterType.FullName);
+            throw new TypeInspectionException(message);
+        }
+
+        if (!typeof(ITestFilter).IsAssignableFrom(filterType))
+        {
+            string message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_TestFilterProviderTypeDoesNotImplementInterface, filterType.FullName, typeof(ITestFilter).FullName);
+            throw new TypeInspectionException(message);
+        }
+
+        try
+        {
+            return (ITestFilter)Activator.CreateInstance(filterType)!;
+        }
+        catch (Exception ex)
+        {
+            string message = string.Format(CultureInfo.CurrentCulture, Resource.UTA_TestFilterProviderInstantiationFailed, filterType.FullName, ex.Message);
+            throw new TypeInspectionException(message, ex);
+        }
+    }
+
+    private static bool HasTestFilterProviderMarker(Assembly assembly)
+    {
+        // Compare on the attribute type's FullName so we don't trigger attribute construction,
+        // mirroring the AssemblyFixtureProvider probe.
+        string markerFullName = typeof(TestFilterProviderAttribute).FullName!;
+        foreach (CustomAttributeData data in assembly.GetCustomAttributesData())
+        {
+            if (string.Equals(data.AttributeType.FullName, markerFullName, StringComparison.Ordinal))
+            {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    // Tiny holder so the cache can distinguish "not computed yet" (missing key) from
+    // "computed and result is no filter" (present key with Filter = null).
+    private sealed class TestFilterBox
+    {
+        public TestFilterBox(ITestFilter? filter) => Filter = filter;
+
+        public ITestFilter? Filter { get; }
+    }
+}

src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs

@@ -10,6 +10,7 @@
 using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel;
 using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices;
 using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Extensions;
+using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Helpers;
 using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Interface;
 using Microsoft.VisualStudio.TestPlatform.ObjectModel.Logging;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -36,6 +37,12 @@ internal sealed class UnitTestRunner
     // Used to attach assembly cleanup failures to the right test.
     private UnitTestElement? _lastRunnableTestInWholeAssembly;
 
+    // Tracks whether at least one test in this runner's lifetime triggered AssemblyInitialize.
+    // Needed so that end-of-assembly cleanup still runs when the very last test of the assembly
+    // was filtered out by a [TestFilterProvider] (in which case testMethodInfo is null for the
+    // cleanup decision in RunSingleTestAsync).
+    private bool _assemblyInitializeWasExecuted;
+
     /// <summary>
     /// Initializes a new instance of the <see cref="UnitTestRunner"/> class.
     /// </summary>
@@ -133,6 +140,21 @@ internal async Task<TestResult[]> RunSingleTestAsync(UnitTestElement unitTestEle
         {
             testContextForTestExecution = PlatformServiceProvider.Instance.GetTestContext(testMethod, null, testContextProperties, messageLogger, UnitTestOutcome.InProgress);
 
+            // Apply user-supplied [TestFilterProvider] filter BEFORE loading the test type, BEFORE
+            // running [AssemblyInitialize] and BEFORE [ClassInitialize]. This is the whole point of
+            // the feature: a Drop or Skip here pays none of those costs. See
+            // https://github.com/microsoft/testfx/issues/8894 for the design.
+            TestResult[]? filterResult = ApplyTestFilter(unitTestElement);
+            if (filterResult is not null)
+            {
+                return await FinishFilteredOutTestAsync(
+                    testMethod,
+                    testContextProperties,
+                    messageLogger,
+                    filterResult,
+                    testContextForTestExecution).ConfigureAwait(false);
+            }
+
             // Get the testMethod
             TestMethodInfo? testMethodInfo = _typeCache.GetTestMethodInfo(testMethod);
 
@@ -148,6 +170,7 @@ internal async Task<TestResult[]> RunSingleTestAsync(UnitTestElement unitTestEle
                 testContextForAssemblyInit = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, null, testContextProperties, messageLogger, testContextForTestExecution.Context.CurrentTestOutcome);
 
                 TestResult assemblyInitializeResult = await RunAssemblyInitializeIfNeededAsync(testMethodInfo, testContextForAssemblyInit).ConfigureAwait(false);
+                _assemblyInitializeWasExecuted |= testMethodInfo.Parent.Parent.IsAssemblyInitializeExecuted;
 
                 if (assemblyInitializeResult.Outcome != UnitTestOutcome.Passed)
                 {
@@ -390,4 +413,171 @@ private static bool IsTestMethodRunnable(
     }
 
     internal void ForceCleanup(IDictionary<string, object?> sourceLevelParameters, IMessageLogger logger) => ClassCleanupManager.ForceCleanup(_typeCache, sourceLevelParameters, logger);
+
+    /// <summary>
+    /// Invokes the user-supplied <see cref="ITestFilter"/> registered via
+    /// <see cref="TestFilterProviderAttribute"/> for the test assembly, if any. Returns
+    /// <see langword="null"/> if no filter is registered or the filter returned
+    /// <see cref="TestFilterResult.Run"/> (test should run normally), an empty array if the
+    /// filter returned <see cref="TestFilterResult.Drop"/>, or a single Skipped
+    /// <see cref="TestResult"/> if the filter returned <see cref="TestFilterResult.Skip(string)"/>.
+    /// </summary>
+    /// <remarks>
+    /// A filter exception is surfaced as an Error test result so the failure is visible to the
+    /// user instead of silently affecting test selection. <see cref="TestFilterProviderAttribute"/>
+    /// is single-per-assembly by design: callers that want to combine multiple strategies should
+    /// compose them explicitly inside their <see cref="ITestFilter"/> implementation.
+    /// </remarks>
+    private TestResult[]? ApplyTestFilter(UnitTestElement unitTestElement)
+    {
+        ITestFilter? filter = _typeCache.GetOrLoadTestFilter(unitTestElement.TestMethod.AssemblyName);
+        if (filter is null)
+        {
+            return null;
+        }
+
+        TestFilterContext context = CreateFilterContext(unitTestElement);
+
+        TestFilterResult result;
+        try
+        {
+            result = filter.Filter(context);
+        }
+        catch (Exception ex)
+        {
+            string message = string.Format(
+                CultureInfo.CurrentCulture,
+                Resource.UTA_TestFilterProviderThrew,
+                filter.GetType().FullName,
+                context.FullyQualifiedName,
+                ex.Message);
+            return
+            [
+                new TestResult
+                {
+                    Outcome = UnitTestOutcome.Error,
+                    TestFailureException = new TestFailedException(UnitTestOutcome.Error, message, ex.TryGetStackTraceInformation()),
+                }
+            ];
+        }
+
+        return result.Action switch
+        {
+            TestFilterAction.Drop => [],
+            TestFilterAction.Skip => [TestResult.CreateIgnoredResult(result.SkipReason)],
+            _ => null,
+        };
+    }
+
+    private static TestFilterContext CreateFilterContext(UnitTestElement element)
+    {
+        TestMethod testMethod = element.TestMethod;
+        string[] categories = element.TestCategory ?? [];
+
+        KeyValuePair<string, string?>[] traits;
+        if (element.Traits is { Length: > 0 } source)
+        {
+            traits = new KeyValuePair<string, string?>[source.Length];
+            for (int i = 0; i < source.Length; i++)
+            {
+                traits[i] = new KeyValuePair<string, string?>(source[i].Name, source[i].Value);
+            }
+        }
+        else
+        {
+            traits = [];
+        }
+
+        // Pull namespace + simple class name from the hierarchy when available — this is the
+        // same source the IDE / Test Explorer uses, so it correctly handles nested types and
+        // generic classes (where naïve FullClassName splitting would lie).
+        string? hierarchyNamespace = null;
+        string? hierarchyClassName = null;
+        if (testMethod.Hierarchy is IReadOnlyList<string?> hierarchy && hierarchy.Count > HierarchyConstants.Levels.ClassIndex)
+        {
+            hierarchyNamespace = hierarchy[HierarchyConstants.Levels.NamespaceIndex];
+            hierarchyClassName = hierarchy[HierarchyConstants.Levels.ClassIndex];
+        }
+
+        // ManagedMethodName is an ECMA-335 string like `MyMethod`1(System.Int32)` — parse it
+        // cheaply (no MethodInfo reflection) to surface arity and parameter type names.
+        int? methodArity = null;
+        IReadOnlyList<string>? parameterTypeFullNames = null;
+        if (testMethod.ManagedMethodName is { } managedMethod)
+        {
+            try
+            {
+                ManagedNameParser.ParseManagedMethodName(managedMethod, out _, out int arity, out string[]? parameterTypes);
+                methodArity = arity;
+                parameterTypeFullNames = parameterTypes ?? (IReadOnlyList<string>)[];
+            }
+            catch
+            {
+                // Defensive: if the managed name is malformed for any reason, surface what we
+                // can via the flat strings rather than failing the filter.
+            }
+        }
+
+        return new TestFilterContext
+        {
+            FullyQualifiedName = $"{testMethod.FullClassName}.{testMethod.Name}",
+            DisplayName = testMethod.DisplayName,
+            MethodName = testMethod.Name,
+            Source = testMethod.AssemblyName,
+            Namespace = hierarchyNamespace,
+            ClassName = hierarchyClassName,
+            ManagedTypeName = testMethod.ManagedTypeName,
+            ManagedMethodName = testMethod.ManagedMethodName,
+            MethodArity = methodArity,
+            ParameterTypeFullNames = parameterTypeFullNames,
+            Categories = categories,
+            Traits = traits,
+            Priority = element.Priority,
+        };
+    }
+
+    /// <summary>
+    /// Handles the bookkeeping (class-cleanup countdown, end-of-assembly cleanup) for a test that
+    /// was filtered out by a <see cref="ITestFilter"/>. Mirrors the tail of <see cref="RunSingleTestAsync"/>
+    /// without ever requiring <c>testMethodInfo</c>, since the filter ran before any type was loaded.
+    /// </summary>
+    private async Task<TestResult[]> FinishFilteredOutTestAsync(
+        TestMethod testMethod,
+        IDictionary<string, object?> testContextProperties,
+        IMessageLogger messageLogger,
+        TestResult[] filterResult,
+        ITestContext testContextForTestExecution)
+    {
+        _classCleanupManager.MarkTestComplete(testMethod, out bool isLastTestInClass);
+        if (isLastTestInClass)
+        {
+            // No class cleanup possible: we never loaded the test type, so there's nothing to
+            // execute. Still mark the class as complete so end-of-assembly cleanup is gated correctly.
+            _classCleanupManager.MarkClassComplete(testMethod.FullClassName);
+        }
+
+        if (_assemblyInitializeWasExecuted && _classCleanupManager.ShouldRunEndOfAssemblyCleanup)
+        {
+            ITestContext? testContextForAssemblyCleanup = null;
+            try
+            {
+                testContextForAssemblyCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, null, testContextProperties, messageLogger, testContextForTestExecution.Context.CurrentTestOutcome);
+
+                TestResult? assemblyCleanupResult = await RunAssemblyCleanupAsync(testContextForAssemblyCleanup, _typeCache, filterResult).ConfigureAwait(false);
+                if (assemblyCleanupResult is not null)
+                {
+                    // Current test was filtered (no result), so an assembly cleanup failure needs to
+                    // be associated with the last real test that ran in the assembly.
+                    assemblyCleanupResult.AssociatedUnitTestElement = _lastRunnableTestInWholeAssembly;
+                    filterResult = [.. filterResult, assemblyCleanupResult];
+                }
+            }
+            finally
+            {
+                (testContextForAssemblyCleanup as IDisposable)?.Dispose();
+            }
+        }
+
+        return filterResult;
+    }
 }