Fix [ClassCleanup] leak when ITestFilter drops the last test of an initialized class

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

@@ -584,9 +584,11 @@ private static TestFilterContext CreateFilterContext(UnitTestElement element)
     }
 
     /// <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.
+    /// Handles the bookkeeping (class-cleanup countdown, class cleanup, end-of-assembly cleanup) for a
+    /// test that was filtered out by a <see cref="ITestFilter"/>. Mirrors the tail of
+    /// <see cref="RunSingleTestAsync"/>. The filtered-out test never loaded its own type, but if a
+    /// sibling test of the same class already ran in this worker the class was initialized and still
+    /// owes its <c>[ClassCleanup]</c>, so it is executed here when this is the last test of the class.
     /// </summary>
     private async Task<TestResult[]> FinishFilteredOutTestAsync(
         TestMethod testMethod,
@@ -598,8 +600,53 @@ private async Task<TestResult[]> FinishFilteredOutTestAsync(
         _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.
+            // The class-cleanup countdown spans the full (pre-filter) set of tests, so the "last test
+            // in class" can land on a filtered-out test. The filtered-out test itself never loaded the
+            // type, but a SIBLING test of the same class may have run earlier in this worker — which
+            // means [ClassInitialize] already executed and [ClassCleanup] is still owed. We must run it
+            // here; otherwise the class leaks its cleanup whenever its last-in-order test is dropped.
+            //
+            // _lastRunnableTestByClass is populated only for classes that both have an executable
+            // cleanup method AND ran at least one non-filtered test in this worker, so its presence is
+            // exactly the signal that the type is already loaded and cleanup is pending. Resolving the
+            // test method info therefore hits the TypeCache and never loads a new type.
+            if (_lastRunnableTestByClass.TryGetValue(testMethod.FullClassName, out UnitTestElement? lastRunnableTest))
+            {
+                TestMethodInfo? testMethodInfo = _typeCache.GetTestMethodInfo(lastRunnableTest.TestMethod);
+                if (testMethodInfo is not null)
+                {
+                    ITestContext testContextForClassCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, testMethod.FullClassName, testContextProperties, messageLogger, testContextForTestExecution.Context.CurrentTestOutcome);
+                    try
+                    {
+                        // Flow properties set during AssemblyInitialize and ClassInitialize so the
+                        // ClassCleanup method observes them, mirroring the run path in RunSingleTestAsync.
+                        var classCleanupImpl = (TestContextImplementation)testContextForClassCleanup.Context;
+                        classCleanupImpl.MergeProperties(testMethodInfo.Parent.Parent.PostAssemblyInitProperties);
+                        classCleanupImpl.MergeProperties(testMethodInfo.Parent.PostClassInitProperties);
+
+                        // Note: filterResult is empty for a dropped test, so any TestContext output
+                        // written by a *successful* [ClassCleanup] is not attached to a result here
+                        // (RunClassCleanupAsync only flushes TestContext output onto an existing result).
+                        // Console output is unaffected — it goes to process stdout, which the test host
+                        // still surfaces. A *failing* cleanup produces its own result, handled below.
+                        TestResult? cleanupResult = await testMethodInfo.Parent.RunClassCleanupAsync(testContextForClassCleanup, filterResult).ConfigureAwait(false);
+                        if (cleanupResult is not null)
+                        {
+                            // The current test was filtered out (no result of its own), so a class
+                            // cleanup failure must be attached to the last real test that ran in the class.
+                            cleanupResult.AssociatedUnitTestElement = lastRunnableTest;
+                            filterResult = [.. filterResult, cleanupResult];
+                        }
+                    }
+                    finally
+                    {
+                        (testContextForClassCleanup as IDisposable)?.Dispose();
+                    }
+                }
+            }
+
+            // Mark the class as complete so end-of-assembly cleanup is gated correctly. Done after the
+            // class cleanup above so assembly cleanup never runs before this class is fully torn down.
             _classCleanupManager.MarkClassComplete(testMethod.FullClassName);
         }
 

test/IntegrationTests/MSTest.Acceptance.IntegrationTests/TestFilterClassCleanupTests.cs

@@ -0,0 +1,211 @@
+// 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.Acceptance.IntegrationTests;
+using Microsoft.Testing.Platform.Acceptance.IntegrationTests.Helpers;
+using Microsoft.Testing.Platform.Helpers;
+
+namespace MSTest.Acceptance.IntegrationTests;
+
+[TestClass]
+public sealed class TestFilterClassCleanupTests : AcceptanceTestBase<TestFilterClassCleanupTests.TestAssetFixture>
+{
+    public TestContext TestContext { get; set; } = null!;
+
+    // Regression test for the [ClassCleanup] leak when an ITestFilter drops the *last-in-order* test of
+    // a class that was already initialized by an earlier (non-dropped) test in the same worker. This is
+    // the exact scenario reported for the per-worker sharding pattern (each worker enumerates the full
+    // set per class, runs its subset, and Drops the rest). Before the fix, ClassCleanupManager counted
+    // the dropped test as the last test of the class, so the drop path skipped class cleanup even though
+    // the type had been loaded and [ClassInitialize] had run.
+    [TestMethod]
+    public async Task ClassCleanup_RunsWhenLastTestOfInitializedClassIsDropped()
+    {
+        var testHost = TestHost.LocateFrom(AssetFixture.ProjectPath, TestAssetFixture.ProjectName, TargetFrameworks.NetCurrent);
+        TestHostResult testHostResult = await testHost.ExecuteAsync(
+            "--settings my.runsettings",
+            environmentVariables: new Dictionary<string, string?>
+            {
+                // Drop the last-in-order test of the partially-run class and the only test of the
+                // fully-dropped class.
+                ["DROP_METHODS"] = "Test_Z_Dropped,OnlyTest",
+            },
+            cancellationToken: TestContext.CancellationToken);
+
+        // Only the single non-dropped test of PartiallyDroppedClass runs.
+        testHostResult.AssertExitCodeIs(ExitCode.Success);
+        testHostResult.AssertOutputContainsSummary(failed: 0, passed: 1, skipped: 0);
+
+        // The partially-dropped class was initialized (a real test ran) so its class cleanup MUST run,
+        // even though its last-in-order test was filtered out.
+        testHostResult.AssertOutputContains("PartiallyDroppedClass.ClassInitialize");
+        testHostResult.AssertOutputContains("PartiallyDroppedClass.Test_A_Run");
+        testHostResult.AssertOutputContains("PartiallyDroppedClass.ClassCleanup");
+        testHostResult.AssertOutputDoesNotContain("PartiallyDroppedClass.Test_Z_Dropped");
+
+        // The fully-dropped class never loaded its type, so neither [ClassInitialize] nor [ClassCleanup]
+        // should run. This guards the fix from over-correcting (running cleanup for a class that was
+        // never initialized).
+        testHostResult.AssertOutputDoesNotContain("FullyDroppedClass.ClassInitialize");
+        testHostResult.AssertOutputDoesNotContain("FullyDroppedClass.ClassCleanup");
+        testHostResult.AssertOutputDoesNotContain("FullyDroppedClass.OnlyTest");
+    }
+
+    // Covers the failure branch of the drop path: when the dropped test is the last in its class and
+    // the class's [ClassCleanup] throws, the failure must surface and be attributed to the last real
+    // test that ran in the class (Test_A_Run), not to the dropped test (which produced no result of
+    // its own). This exercises the `cleanupResult is not null` branch and the AssociatedUnitTestElement
+    // attribution in FinishFilteredOutTestAsync.
+    [TestMethod]
+    public async Task ClassCleanup_FailureIsAttributedToLastRealTest_WhenLastTestIsDropped()
+    {
+        var testHost = TestHost.LocateFrom(AssetFixture.ProjectPath, TestAssetFixture.ProjectName, TargetFrameworks.NetCurrent);
+        TestHostResult testHostResult = await testHost.ExecuteAsync(
+            "--settings my.runsettings",
+            environmentVariables: new Dictionary<string, string?>
+            {
+                ["DROP_METHODS"] = "Test_Z_Dropped,OnlyTest",
+                // Make PartiallyDroppedClass.ClassCleanup throw so we exercise the failure branch.
+                ["THROW_CLEANUP"] = "1",
+            },
+            cancellationToken: TestContext.CancellationToken);
+
+        // The class was initialized and its cleanup ran (in the drop path) and threw, so the run fails.
+        testHostResult.AssertExitCodeIs(ExitCode.AtLeastOneTestFailed);
+
+        // The cleanup actually ran in the drop path and its failure surfaced (proves the
+        // `cleanupResult is not null` branch and the filterResult spread are exercised).
+        testHostResult.AssertOutputContains("PartiallyDroppedClass.ClassInitialize");
+        testHostResult.AssertOutputContains("PartiallyDroppedClass.Test_A_Run");
+        testHostResult.AssertOutputContains("ClassCleanup failed on purpose");
+
+        // The failure is attributed to the last real test that ran (Test_A_Run), not the dropped test.
+        testHostResult.AssertOutputContains("PartiallyDroppedClass.Test_A_Run");
+        testHostResult.AssertOutputDoesNotContain("PartiallyDroppedClass.Test_Z_Dropped");
+    }
+
+    public sealed class TestAssetFixture() : TestAssetFixtureBase()
+    {
+        public const string ProjectName = "TestFilterClassCleanup";
+
+        public string ProjectPath => GetAssetPath(ProjectName);
+
+        public override (string ID, string Name, string Code) GetAssetsToGenerate() => (ProjectName, ProjectName,
+                SourceCode
+                .PatchTargetFrameworks(TargetFrameworks.NetCurrent)
+                .PatchCodeWithReplace("$MSTestVersion$", MSTestVersion));
+
+        private const string SourceCode = """
+#file TestFilterClassCleanup.csproj
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <EnableMSTestRunner>true</EnableMSTestRunner>
+    <TargetFrameworks>$TargetFrameworks$</TargetFrameworks>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="MSTest.TestAdapter" Version="$MSTestVersion$" />
+    <PackageReference Include="MSTest.TestFramework" Version="$MSTestVersion$" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <None Update="*.runsettings">
+      <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+    </None>
+  </ItemGroup>
+
+</Project>
+
+#file my.runsettings
+<RunSettings>
+  <MSTest>
+    <CaptureTraceOutput>false</CaptureTraceOutput>
+    <!-- Force alphabetical ordering so Test_A_Run always executes before Test_Z_Dropped, making the
+         "last-in-order test is dropped" scenario deterministic instead of relying on discovery order. -->
+    <OrderTestsByNameInClass>true</OrderTestsByNameInClass>
+  </MSTest>
+</RunSettings>
+
+#file ShardFilter.cs
+using System;
+using System.Collections.Generic;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+[assembly: TestFilterProvider(typeof(ShardFilter))]
+
+// Drops every test method whose name is listed in the DROP_METHODS environment variable. This mimics a
+// worker that claims a subset of tests and drops the rest before the test type is loaded.
+public sealed class ShardFilter : ITestFilter
+{
+    private static readonly HashSet<string> DropMethods = new(
+        (Environment.GetEnvironmentVariable("DROP_METHODS") ?? string.Empty)
+            .Split(new[] { ',' }, StringSplitOptions.RemoveEmptyEntries),
+        StringComparer.Ordinal);
+
+    public TestFilterResult Filter(TestFilterContext context)
+        => DropMethods.Contains(context.MethodName) ? TestFilterResult.Drop : TestFilterResult.Run;
+}
+
+#file PartiallyDroppedClass.cs
+using System;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+// This class runs at least one real test (Test_A_Run) so [ClassInitialize] executes, but its
+// last-in-order test (Test_Z_Dropped) is filtered out. Its [ClassCleanup] must still run.
+[TestClass]
+public class PartiallyDroppedClass
+{
+    [ClassInitialize]
+    public static void ClassInitialize(TestContext context)
+        => Console.WriteLine("PartiallyDroppedClass.ClassInitialize");
+
+    [ClassCleanup]
+    public static void ClassCleanup()
+    {
+        if (Environment.GetEnvironmentVariable("THROW_CLEANUP") == "1")
+        {
+            throw new InvalidOperationException("PartiallyDroppedClass.ClassCleanup failed on purpose");
+        }
+
+        Console.WriteLine("PartiallyDroppedClass.ClassCleanup");
+    }
+
+    // Method names are intentional: with <OrderTestsByNameInClass> enabled, alphabetical ordering
+    // guarantees Test_A_Run executes first (triggering [ClassInitialize]) and Test_Z_Dropped is the
+    // last-in-order test that the filter drops — exactly the scenario the fix addresses. Renaming or
+    // reordering these would silently stop exercising the bug.
+    [TestMethod]
+    public void Test_A_Run()
+        => Console.WriteLine("PartiallyDroppedClass.Test_A_Run");
+
+    [TestMethod]
+    public void Test_Z_Dropped()
+        => Console.WriteLine("PartiallyDroppedClass.Test_Z_Dropped");
+}
+
+#file FullyDroppedClass.cs
+using System;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+// Every test of this class is dropped, so the type is never loaded: neither [ClassInitialize] nor
+// [ClassCleanup] should run.
+[TestClass]
+public class FullyDroppedClass
+{
+    [ClassInitialize]
+    public static void ClassInitialize(TestContext context)
+        => Console.WriteLine("FullyDroppedClass.ClassInitialize");
+
+    [ClassCleanup]
+    public static void ClassCleanup()
+        => Console.WriteLine("FullyDroppedClass.ClassCleanup");
+
+    [TestMethod]
+    public void OnlyTest()
+        => Console.WriteLine("FullyDroppedClass.OnlyTest");
+}
+""";
+    }
+}