Support '--maximum-failed-tests' to abort test run when failure threshold is reached (#4238)

Author: Youssef1313

Diff for: src/Adapter/MSTest.TestAdapter/Execution/ClassCleanupManager.cs

@@ -62,4 +62,19 @@ public void MarkTestComplete(TestMethodInfo testMethodInfo, TestMethod testMetho
             ShouldRunEndOfAssemblyCleanup = _remainingTestsByClass.IsEmpty;
         }
     }
+
+    internal static void ForceCleanup(TypeCache typeCache)
+    {
+        IEnumerable<TestClassInfo> classInfoCache = typeCache.ClassInfoListWithExecutableCleanupMethods;
+        foreach (TestClassInfo classInfo in classInfoCache)
+        {
+            classInfo.ExecuteClassCleanup();
+        }
+
+        IEnumerable<TestAssemblyInfo> assemblyInfoCache = typeCache.AssemblyInfoListWithExecutableCleanupMethods;
+        foreach (TestAssemblyInfo assemblyInfo in assemblyInfoCache)
+        {
+            assemblyInfo.ExecuteAssemblyCleanup();
+        }
+    }
 }

Diff for: src/Adapter/MSTest.TestAdapter/Execution/TestClassInfo.cs

@@ -252,6 +252,7 @@ public void RunClassInitialize(TestContext testContext)
         // If no class initialize and no base class initialize, return
         if (ClassInitializeMethod is null && BaseClassInitMethods.Count == 0)
         {
+            IsClassInitializeExecuted = true;
             return;
         }
 
@@ -558,8 +559,10 @@ internal void ExecuteClassCleanup()
         lock (_testClassExecuteSyncObject)
         {
             if (IsClassCleanupExecuted
-                // If there is a ClassInitialize method and it has not been executed, then we should not execute ClassCleanup
-                || (!IsClassInitializeExecuted && ClassInitializeMethod is not null))
+                // If ClassInitialize method has not been executed, then we should not execute ClassCleanup
+                // Note that if there is no ClassInitialze method at all, we will still set
+                // IsClassInitializeExecuted to true in RunClassInitialize
+                || !IsClassInitializeExecuted)
             {
                 return;
             }

Diff for: src/Adapter/MSTest.TestAdapter/Execution/TestExecutionManager.cs

@@ -399,6 +399,11 @@ private void ExecuteTestsInSource(IEnumerable<TestCase> tests, IRunContext? runC
             ExecuteTestsWithTestRunner(testsToRun, frameworkHandle, source, sourceLevelParameters, testRunner);
         }
 
+        if (MSTestGracefulStopTestExecutionCapability.Instance.IsStopRequested)
+        {
+            testRunner.ForceCleanup();
+        }
+
         PlatformServiceProvider.Instance.AdapterTraceLogger.LogInfo("Executed tests belonging to source {0}", source);
     }
 
@@ -419,6 +424,10 @@ private void ExecuteTestsWithTestRunner(
         foreach (TestCase currentTest in orderedTests)
         {
             _testRunCancellationToken?.ThrowIfCancellationRequested();
+            if (MSTestGracefulStopTestExecutionCapability.Instance.IsStopRequested)
+            {
+                break;
+            }
 
             // If it is a fixture test, add it to the list of fixture tests and do not execute it.
             // It is executed by test itself.

Diff for: src/Adapter/MSTest.TestAdapter/Execution/UnitTestRunner.cs

@@ -169,6 +169,7 @@ internal UnitTestResult[] RunSingleTest(TestMethod testMethod, IDictionary<strin
                 else
                 {
                     UnitTestResult classInitializeResult = testMethodInfo.Parent.GetResultOrRunClassInitialize(testContext, assemblyInitializeResult.StandardOut!, assemblyInitializeResult.StandardError!, assemblyInitializeResult.DebugTrace!, assemblyInitializeResult.TestContextMessages!);
+                    DebugEx.Assert(testMethodInfo.Parent.IsClassInitializeExecuted, "IsClassInitializeExecuted should be true after attempting to run it.");
                     if (classInitializeResult.Outcome != UnitTestOutcome.Passed)
                     {
                         result = [classInitializeResult];
@@ -361,4 +362,6 @@ private bool IsTestMethodRunnable(
         notRunnableResult = null;
         return true;
     }
+
+    internal void ForceCleanup() => ClassCleanupManager.ForceCleanup(_typeCache);
 }

Diff for: src/Adapter/MSTest.TestAdapter/TestingPlatformAdapter/MSTestGracefulStopTestExecutionCapability.cs

@@ -0,0 +1,25 @@
+// 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.Testing.Platform.Capabilities.TestFramework;
+
+namespace Microsoft.VisualStudio.TestTools.UnitTesting;
+
+#pragma warning disable TPEXP // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+internal sealed class MSTestGracefulStopTestExecutionCapability : IGracefulStopTestExecutionCapability
+#pragma warning restore TPEXP // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+{
+    private MSTestGracefulStopTestExecutionCapability()
+    {
+    }
+
+    public static MSTestGracefulStopTestExecutionCapability Instance { get; } = new();
+
+    public bool IsStopRequested { get; private set; }
+
+    public Task StopTestExecutionAsync(CancellationToken cancellationToken)
+    {
+        IsStopRequested = true;
+        return Task.CompletedTask;
+    }
+}

Diff for: src/Adapter/MSTest.TestAdapter/TestingPlatformAdapter/TestApplicationBuilderExtensions.cs

@@ -8,6 +8,7 @@
 using Microsoft.Testing.Extensions.VSTestBridge.Helpers;
 using Microsoft.Testing.Platform.Builder;
 using Microsoft.Testing.Platform.Capabilities.TestFramework;
+using Microsoft.Testing.Platform.Helpers;
 using Microsoft.Testing.Platform.Services;
 
 namespace Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -20,12 +21,16 @@ public static void AddMSTest(this ITestApplicationBuilder testApplicationBuilder
         testApplicationBuilder.AddRunSettingsService(extension);
         testApplicationBuilder.AddTestCaseFilterService(extension);
         testApplicationBuilder.AddTestRunParametersService(extension);
+#pragma warning disable TPEXP // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+        testApplicationBuilder.AddMaximumFailedTestsService(extension);
+#pragma warning restore TPEXP // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
         testApplicationBuilder.AddRunSettingsEnvironmentVariableProvider(extension);
         testApplicationBuilder.RegisterTestFramework(
             serviceProvider => new TestFrameworkCapabilities(
                 new VSTestBridgeExtensionBaseCapabilities(),
 #pragma warning disable TPEXP // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
-                new MSTestBannerCapability(serviceProvider.GetRequiredService<IPlatformInformation>())),
+                new MSTestBannerCapability(serviceProvider.GetRequiredService<IPlatformInformation>()),
+                MSTestGracefulStopTestExecutionCapability.Instance),
 #pragma warning restore TPEXP // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
             (capabilities, serviceProvider) => new MSTestBridgedTestFramework(extension, getTestAssemblies, serviceProvider, capabilities));
     }

Diff for: src/Platform/Microsoft.Testing.Platform/Capabilities/TestFramework/IGracefulStopTestExecutionCapability.cs

@@ -0,0 +1,19 @@
+// 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.Diagnostics.CodeAnalysis;
+
+namespace Microsoft.Testing.Platform.Capabilities.TestFramework;
+
+/// <summary>
+/// A capability to support stopping test execution gracefully, without cancelling/aborting everything.
+/// This is used to support '--maximum-failed-tests'.
+/// </summary>
+/// <remarks>
+/// Test frameworks can choose to run any needed cleanup when cancellation is requested.
+/// </remarks>
+[Experimental("TPEXP", UrlFormat = "https://aka.ms/testingplatform/diagnostics#{0}")]
+public interface IGracefulStopTestExecutionCapability : ITestFrameworkCapability
+{
+    Task StopTestExecutionAsync(CancellationToken cancellationToken);
+}

Diff for: src/Platform/Microsoft.Testing.Platform/CommandLine/MaxFailedTestsCommandLineOptionsProvider.cs

@@ -0,0 +1,53 @@
+// 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.Globalization;
+
+using Microsoft.Testing.Platform.Extensions;
+using Microsoft.Testing.Platform.Extensions.CommandLine;
+using Microsoft.Testing.Platform.Helpers;
+using Microsoft.Testing.Platform.Resources;
+
+namespace Microsoft.Testing.Platform.CommandLine;
+
+internal sealed class MaxFailedTestsCommandLineOptionsProvider(IExtension extension) : ICommandLineOptionsProvider
+{
+    internal const string MaxFailedTestsOptionKey = "maximum-failed-tests";
+
+    private static readonly IReadOnlyCollection<CommandLineOption> OptionsCache =
+    [
+        new(MaxFailedTestsOptionKey, PlatformResources.PlatformCommandLineMaxFailedTestsOptionDescription, ArgumentArity.ExactlyOne, isHidden: false),
+    ];
+
+    public string Uid => extension.Uid;
+
+    public string Version => extension.Version;
+
+    public string DisplayName => extension.DisplayName;
+
+    public string Description => extension.Description;
+
+    public IReadOnlyCollection<CommandLineOption> GetCommandLineOptions()
+        => OptionsCache;
+
+    public Task<bool> IsEnabledAsync() => Task.FromResult(true);
+
+    public Task<ValidationResult> ValidateCommandLineOptionsAsync(ICommandLineOptions commandLineOptions)
+        => ValidationResult.ValidTask;
+
+    public Task<ValidationResult> ValidateOptionArgumentsAsync(CommandLineOption commandOption, string[] arguments)
+    {
+        if (commandOption.Name == MaxFailedTestsOptionKey)
+        {
+            string arg = arguments[0];
+            // We consider --maximum-failed-tests 0 as invalid.
+            // The idea is that we stop the execution when we *reach* the max failed tests, not when *exceed*.
+            // So the value 1 means, stop execution on the first failure.
+            return int.TryParse(arg, out int maxFailedTestsResult) && maxFailedTestsResult > 0
+                ? ValidationResult.ValidTask
+                : ValidationResult.InvalidTask(string.Format(CultureInfo.InvariantCulture, PlatformResources.MaxFailedTestsMustBePositive, arg));
+        }
+
+        throw ApplicationStateGuard.Unreachable();
+    }
+}

Diff for: src/Platform/Microsoft.Testing.Platform/Extensions/AbortForMaxFailedTestsExtension.cs

@@ -0,0 +1,76 @@
+// 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.Testing.Platform.Capabilities.TestFramework;
+using Microsoft.Testing.Platform.CommandLine;
+using Microsoft.Testing.Platform.Extensions.Messages;
+using Microsoft.Testing.Platform.Extensions.TestHost;
+using Microsoft.Testing.Platform.Helpers;
+using Microsoft.Testing.Platform.Messages;
+using Microsoft.Testing.Platform.Resources;
+using Microsoft.Testing.Platform.Services;
+
+namespace Microsoft.Testing.Platform.Extensions;
+
+internal sealed class AbortForMaxFailedTestsExtension : IDataConsumer
+{
+    private readonly int? _maxFailedTests;
+    private readonly IGracefulStopTestExecutionCapability? _capability;
+    private readonly IStopPoliciesService _policiesService;
+    private readonly ITestApplicationCancellationTokenSource _testApplicationCancellationTokenSource;
+    private int _failCount;
+
+    public AbortForMaxFailedTestsExtension(
+        ICommandLineOptions commandLineOptions,
+        IGracefulStopTestExecutionCapability? capability,
+        IStopPoliciesService policiesService,
+        ITestApplicationCancellationTokenSource testApplicationCancellationTokenSource)
+    {
+        if (commandLineOptions.TryGetOptionArgumentList(MaxFailedTestsCommandLineOptionsProvider.MaxFailedTestsOptionKey, out string[]? args) &&
+            int.TryParse(args[0], out int maxFailedTests) &&
+            maxFailedTests > 0)
+        {
+            _maxFailedTests = maxFailedTests;
+        }
+
+        _capability = capability;
+        _policiesService = policiesService;
+        _testApplicationCancellationTokenSource = testApplicationCancellationTokenSource;
+    }
+
+    public Type[] DataTypesConsumed { get; } = [typeof(TestNodeUpdateMessage)];
+
+    /// <inheritdoc />
+    public string Uid { get; } = nameof(AbortForMaxFailedTestsExtension);
+
+    /// <inheritdoc />
+    public string Version { get; } = AppVersion.DefaultSemVer;
+
+    /// <inheritdoc />
+    public string DisplayName { get; } = nameof(AbortForMaxFailedTestsExtension);
+
+    /// <inheritdoc />
+    public string Description { get; } = PlatformResources.AbortForMaxFailedTestsDescription;
+
+    /// <inheritdoc />
+    public Task<bool> IsEnabledAsync() => Task.FromResult(_maxFailedTests.HasValue && _capability is not null);
+
+    public async Task ConsumeAsync(IDataProducer dataProducer, IData value, CancellationToken cancellationToken)
+    {
+        var node = (TestNodeUpdateMessage)value;
+
+        // If we are called, the extension is enabled, which means both _maxFailedTests and capability are not null.
+        RoslynDebug.Assert(_maxFailedTests is not null);
+        RoslynDebug.Assert(_capability is not null);
+
+        TestNodeStateProperty testNodeStateProperty = node.TestNode.Properties.Single<TestNodeStateProperty>();
+        if (TestNodePropertiesCategories.WellKnownTestNodeTestRunOutcomeFailedProperties.Any(t => t == testNodeStateProperty.GetType()) &&
+            ++_failCount >= _maxFailedTests.Value &&
+            // If already triggered, don't do it again.
+            !_policiesService.IsMaxFailedTestsTriggered)
+        {
+            await _capability.StopTestExecutionAsync(_testApplicationCancellationTokenSource.CancellationToken);
+            await _policiesService.ExecuteMaxFailedTestsCallbacksAsync(_maxFailedTests.Value, _testApplicationCancellationTokenSource.CancellationToken);
+        }
+    }
+}

Diff for: src/Platform/Microsoft.Testing.Platform/Helpers/ExitCodes.cs

@@ -22,4 +22,5 @@ internal static class ExitCodes
     public const int TestAdapterTestSessionFailure = 10;
     public const int DependentProcessExited = 11;
     public const int IncompatibleProtocolVersion = 12;
+    public const int TestExecutionStoppedForMaxFailedTests = 13;
 }

Diff for: src/Platform/Microsoft.Testing.Platform/Helpers/TestApplicationBuilderExtensions.cs

@@ -15,4 +15,12 @@ public static class TestApplicationBuilderExtensions
 {
     public static void AddTreeNodeFilterService(this ITestApplicationBuilder testApplicationBuilder, IExtension extension)
         => testApplicationBuilder.CommandLine.AddProvider(() => new TreeNodeFilterCommandLineOptionsProvider(extension));
+
+    /// <summary>
+    /// Registers the command-line options provider for '--maximum-failed-tests'.
+    /// </summary>
+    /// <param name="builder">The test application builder.</param>
+    [Experimental("TPEXP", UrlFormat = "https://aka.ms/testingplatform/diagnostics#{0}")]
+    public static void AddMaximumFailedTestsService(this ITestApplicationBuilder builder, IExtension extension)
+        => builder.CommandLine.AddProvider(() => new MaxFailedTestsCommandLineOptionsProvider(extension));
 }

Diff for: src/Platform/Microsoft.Testing.Platform/Hosts/TestHostBuilder.cs

@@ -206,6 +206,10 @@ public async Task<ITestHost> BuildAsync(
         // Set the concrete command line options to the proxy.
         commandLineOptionsProxy.SetCommandLineOptions(commandLineHandler);
 
+        // This is needed by output device.
+        var policiesService = new StopPoliciesService(testApplicationCancellationTokenSource);
+        serviceProvider.AddService(policiesService);
+
         bool hasServerFlag = commandLineHandler.TryGetOptionArgumentList(PlatformCommandLineProvider.ServerOptionKey, out string[]? protocolName);
         bool isJsonRpcProtocol = protocolName is null || protocolName.Length == 0 || protocolName[0].Equals(PlatformCommandLineProvider.JsonRpcProtocolName, StringComparison.OrdinalIgnoreCase);
 
@@ -313,9 +317,9 @@ public async Task<ITestHost> BuildAsync(
         // Register the ITestApplicationResult
         TestApplicationResult testApplicationResult = new(
             proxyOutputDevice,
-            serviceProvider.GetTestApplicationCancellationTokenSource(),
             serviceProvider.GetCommandLineOptions(),
-            serviceProvider.GetEnvironment());
+            serviceProvider.GetEnvironment(),
+            policiesService);
         serviceProvider.AddService(testApplicationResult);
 
         // ============= SETUP COMMON SERVICE USED IN ALL MODES END ===============//
@@ -376,6 +380,8 @@ await LogTestHostCreatedAsync(
         TestHostOrchestratorConfiguration testHostOrchestratorConfiguration = await TestHostOrchestratorManager.BuildAsync(serviceProvider);
         if (testHostOrchestratorConfiguration.TestHostOrchestrators.Length > 0 && !commandLineHandler.IsOptionSet(PlatformCommandLineProvider.DiscoverTestsOptionKey))
         {
+            policiesService.ProcessRole = TestProcessRole.TestHostOrchestrator;
+            await proxyOutputDevice.HandleProcessRoleAsync(TestProcessRole.TestHostOrchestrator);
             return new TestHostOrchestratorHost(testHostOrchestratorConfiguration, serviceProvider);
         }
 
@@ -411,6 +417,8 @@ await LogTestHostCreatedAsync(
             if (testHostControllers.RequireProcessRestart)
             {
                 testHostControllerInfo.IsCurrentProcessTestHostController = true;
+                policiesService.ProcessRole = TestProcessRole.TestHostController;
+                await proxyOutputDevice.HandleProcessRoleAsync(TestProcessRole.TestHostController);
                 TestHostControllersTestHost testHostControllersTestHost = new(testHostControllers, testHostControllersServiceProvider, passiveNode, systemEnvironment, loggerFactory, systemClock);
 
                 await LogTestHostCreatedAsync(
@@ -424,6 +432,8 @@ await LogTestHostCreatedAsync(
         }
 
         // ======= TEST HOST MODE ======== //
+        policiesService.ProcessRole = TestProcessRole.TestHost;
+        await proxyOutputDevice.HandleProcessRoleAsync(TestProcessRole.TestHost);
 
         // Setup the test host working folder.
         // Out of the test host controller extension the current working directory is the test host working directory.
@@ -724,6 +734,17 @@ private async Task<ITestFramework> BuildTestFrameworkAsync(TestFrameworkBuilderD
             dataConsumersBuilder.Add(pushOnlyProtocolDataConsumer);
         }
 
+        var abortForMaxFailedTestsExtension = new AbortForMaxFailedTestsExtension(
+            serviceProvider.GetCommandLineOptions(),
+            serviceProvider.GetTestFrameworkCapabilities().GetCapability<IGracefulStopTestExecutionCapability>(),
+            serviceProvider.GetRequiredService<IStopPoliciesService>(),
+            serviceProvider.GetTestApplicationCancellationTokenSource());
+
+        if (await abortForMaxFailedTestsExtension.IsEnabledAsync())
+        {
+            dataConsumersBuilder.Add(abortForMaxFailedTestsExtension);
+        }
+
         IDataConsumer[] dataConsumerServices = dataConsumersBuilder.ToArray();
 
         // Build the message bus

Diff for: src/Platform/Microsoft.Testing.Platform/OutputDevice/IPlatformOutputDevice.cs

@@ -15,4 +15,6 @@ internal interface IPlatformOutputDevice : IExtension
     Task DisplayAfterSessionEndRunAsync();
 
     Task DisplayAsync(IOutputDeviceDataProducer producer, IOutputDeviceData data);
+
+    Task HandleProcessRoleAsync(TestProcessRole processRole);
 }