Fix AssemblyCleanup running when not all class cleanups might have finished

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

@@ -24,26 +24,31 @@ public ClassCleanupManager(IEnumerable<UnitTestElement> testsToRun)
 
     public bool ShouldRunEndOfAssemblyCleanup => _remainingTestCountsByClass.IsEmpty;
 
-    public void MarkTestComplete(TestMethodInfo testMethodInfo, out bool shouldRunEndOfClassCleanup)
+    public void MarkTestComplete(TestMethod testMethod, out bool isLastTestInClass)
     {
-        shouldRunEndOfClassCleanup = false;
-
         lock (_remainingTestCountsByClass)
         {
-            if (!_remainingTestCountsByClass.TryGetValue(testMethodInfo.TestClassName, out int remainingCount))
+            if (!_remainingTestCountsByClass.TryGetValue(testMethod.FullClassName, out int remainingCount))
             {
-                return;
+                throw ApplicationStateGuard.Unreachable();
             }
 
             remainingCount--;
-            _remainingTestCountsByClass[testMethodInfo.TestClassName] = remainingCount;
-            if (remainingCount == 0)
+            _remainingTestCountsByClass[testMethod.FullClassName] = remainingCount;
+            isLastTestInClass = remainingCount == 0;
+        }
+    }
+
+    public void MarkClassComplete(string fullClassName)
+    {
+        lock (_remainingTestCountsByClass)
+        {
+            if (!_remainingTestCountsByClass.TryRemove(fullClassName, out int remainingTests) ||
+                remainingTests != 0)
             {
-                _remainingTestCountsByClass.TryRemove(testMethodInfo.TestClassName, out _);
-                if (testMethodInfo.Parent.HasExecutableCleanupMethod)
-                {
-                    shouldRunEndOfClassCleanup = true;
-                }
+                // We failed to remove the class, or we are incorrectly marking the class as complete while there are remaining tests.
+                // This should never happen.
+                throw ApplicationStateGuard.Unreachable();
             }
         }
     }

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

@@ -620,20 +620,9 @@ async Task<TestResult> DoRunAsync()
         return testFailedException;
     }
 
-    internal async Task RunClassCleanupAsync(ITestContext testContext, ClassCleanupManager classCleanupManager, TestMethodInfo testMethodInfo, TestResult[] results)
+    internal async Task RunClassCleanupAsync(ITestContext testContext, TestResult[] results)
     {
-        DebugEx.Assert(testMethodInfo.Parent == this, "Parent of testMethodInfo should be this TestClassInfo.");
-
-        classCleanupManager.MarkTestComplete(testMethodInfo, out bool shouldRunEndOfClassCleanup);
-        if (!shouldRunEndOfClassCleanup)
-        {
-            return;
-        }
-
-        // TODO: Looks like 'ClassCleanupMethod is null && BaseClassCleanupMethods.Count == 0' is always false?
-        // shouldRunEndOfClassCleanup should be false if there are no class cleanup methods at all.
-        if ((ClassCleanupMethod is null && BaseClassCleanupMethods.Count == 0)
-                || IsClassCleanupExecuted)
+        if (!HasExecutableCleanupMethod || IsClassCleanupExecuted)
         {
             // DoRun will already do nothing for this condition. So, we gain a bit of performance.
             return;

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

@@ -204,15 +204,23 @@ internal async Task<TestResult[]> RunSingleTestAsync(TestMethod testMethod, IDic
             }
 
             testContextForClassCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, testMethod.FullClassName, testContextProperties, messageLogger, testContextForTestExecution.Context.CurrentTestOutcome);
-            if (testMethodInfo is not null)
+
+            _classCleanupManager.MarkTestComplete(testMethod, out bool isLastTestInClass);
+            if (isLastTestInClass && testMethodInfo is not null)
             {
-                await testMethodInfo.Parent.RunClassCleanupAsync(testContextForClassCleanup, _classCleanupManager, testMethodInfo, result).ConfigureAwait(false);
+                await testMethodInfo.Parent.RunClassCleanupAsync(testContextForClassCleanup, result).ConfigureAwait(false);
+
+                // Mark the class as complete when all class cleanups are complete. When all classes are complete we progress to running assembly cleanup.
+                // Class is not complete until after all class cleanups are done, to prevent running assembly cleanup too early.
+                // Do not mark the class as complete when the last test method in the class completed. That is too early, we need to run class cleanups before marking class as complete.
+                _classCleanupManager.MarkClassComplete(testMethod.FullClassName);
             }
 
-            if (testMethodInfo?.Parent.Parent.IsAssemblyInitializeExecuted == true)
+            if (testMethodInfo?.Parent.Parent.IsAssemblyInitializeExecuted == true &&
+                _classCleanupManager.ShouldRunEndOfAssemblyCleanup)
             {
                 testContextForAssemblyCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, null, testContextProperties, messageLogger, testContextForClassCleanup.Context.CurrentTestOutcome);
-                await RunAssemblyCleanupIfNeededAsync(testContextForAssemblyCleanup, _classCleanupManager, _typeCache, result).ConfigureAwait(false);
+                await RunAssemblyCleanupAsync(testContextForAssemblyCleanup, _typeCache, result).ConfigureAwait(false);
             }
 
             return result;
@@ -268,13 +276,8 @@ private static async Task<TestResult> RunAssemblyInitializeIfNeededAsync(TestMet
         return result;
     }
 
-    private static async Task RunAssemblyCleanupIfNeededAsync(ITestContext testContext, ClassCleanupManager classCleanupManager, TypeCache typeCache, TestResult[] results)
+    private static async Task RunAssemblyCleanupAsync(ITestContext testContext, TypeCache typeCache, TestResult[] results)
     {
-        if (!classCleanupManager.ShouldRunEndOfAssemblyCleanup)
-        {
-            return;
-        }
-
         try
         {
             IEnumerable<TestAssemblyInfo> assemblyInfoCache = typeCache.AssemblyInfoListWithExecutableCleanupMethods;

test/IntegrationTests/MSTest.Acceptance.IntegrationTests/AssemblyCleanupTests.cs

@@ -0,0 +1,127 @@
+// 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 AssemblyCleanupTests : AcceptanceTestBase<AssemblyCleanupTests.TestAssetFixture>
+{
+    [TestMethod]
+    public async Task AssemblyCleanupShouldRunAfterAllClassCleanupsHaveCompleted()
+    {
+        var testHost = TestHost.LocateFrom(AssetFixture.ProjectPath, TestAssetFixture.ProjectName, TargetFrameworks.NetCurrent);
+        TestHostResult testHostResult = await testHost.ExecuteAsync("--settings my.runsettings", cancellationToken: TestContext.CancellationToken);
+
+        testHostResult.AssertExitCodeIs(ExitCodes.Success);
+        testHostResult.AssertOutputContainsSummary(failed: 0, passed: 2, skipped: 0);
+        testHostResult.AssertOutputContains("""
+            TestClass1.Test1.
+            TestClass1.Cleanup1 started.
+            TestClass1.Cleanup1 finished.
+            In AsmCleanup
+            """);
+    }
+
+    public sealed class TestAssetFixture() : TestAssetFixtureBase(AcceptanceFixture.NuGetGlobalPackagesFolder)
+    {
+        public const string ProjectName = "AssemblyCleanupTests";
+
+        public string ProjectPath => GetAssetPath(ProjectName);
+
+        public override IEnumerable<(string ID, string Name, string Code)> GetAssetsToGenerate()
+        {
+            yield return (ProjectName, ProjectName,
+                SourceCode
+                .PatchTargetFrameworks(TargetFrameworks.All)
+                .PatchCodeWithReplace("$MSTestVersion$", MSTestVersion));
+        }
+
+        private const string SourceCode = """
+#file AssemblyCleanupTests.csproj
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <OutputType>Exe</OutputType>
+    <EnableMSTestRunner>true</EnableMSTestRunner>
+    <TargetFrameworks>$TargetFrameworks$</TargetFrameworks>
+    <LangVersion>preview</LangVersion>
+  </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 TestClass1.cs
+using System;
+using System.Collections.Generic;
+using System.Text;
+using System.Threading;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+[assembly: Parallelize(Scope = ExecutionScope.MethodLevel, Workers = 0)]
+
+[TestClass]
+public class TestClass1
+{
+    public static bool ClassCleanupFinished { get; private set; }
+
+    [TestMethod]
+    public void Test1()
+    {
+        Console.WriteLine("TestClass1.Test1.");
+    }
+
+    [ClassCleanup]
+    public static void Cleanup1()
+    {
+        Console.WriteLine("TestClass1.Cleanup1 started.");
+        Thread.Sleep(4000);
+        Console.WriteLine("TestClass1.Cleanup1 finished.");
+    }
+}
+
+[TestClass]
+public class TestClass2
+{
+    [TestMethod]
+    public void Test2()
+    {
+    }
+
+    [ClassCleanup]
+    public static void Cleanup2()
+        => Thread.Sleep(2000);
+}
+
+[TestClass]
+public static class Asm
+{
+    [AssemblyCleanup]
+    public static void AsmCleanup()
+        => Console.WriteLine("In AsmCleanup");
+}
+
+#file my.runsettings
+<RunSettings>
+  <MSTest>
+    <CaptureTraceOutput>false</CaptureTraceOutput>
+  </MSTest>
+</RunSettings>
+""";
+    }
+
+    public TestContext TestContext { get; set; }
+}

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

@@ -18,10 +18,9 @@ public class ClassCleanupManagerTests : TestContainer
     public void AssemblyCleanupRunsAfterAllTestsFinishEvenIfWeScheduleTheSameTestMultipleTime()
     {
         ReflectHelper reflectHelper = Mock.Of<ReflectHelper>();
-        MethodInfo methodInfo = typeof(FakeTestClass).GetMethod(nameof(FakeTestClass.FakeTestMethod), BindingFlags.Instance | BindingFlags.NonPublic)!;
         MethodInfo classCleanupMethodInfo = typeof(FakeTestClass).GetMethod(nameof(FakeTestClass.FakeClassCleanupMethod), BindingFlags.Instance | BindingFlags.NonPublic)!;
         // Full class name must agree between unitTestElement.TestMethod.FullClassName and testMethod.FullClassName;
-        string fullClassName = methodInfo.DeclaringType!.FullName!;
+        string fullClassName = typeof(FakeTestClass).FullName!;
         TestMethod testMethod = new(nameof(FakeTestClass.FakeTestMethod), fullClassName, typeof(FakeTestClass).Assembly.FullName!, displayName: null);
 
         // Setting 2 of the same test to run, we should run assembly cleanup after both these tests
@@ -39,14 +38,14 @@ public void AssemblyCleanupRunsAfterAllTestsFinishEvenIfWeScheduleTheSameTestMul
             // This needs to be set, to allow running class cleanup.
             ClassCleanupMethod = classCleanupMethodInfo,
         };
-        TestMethodInfo testMethodInfo = new(methodInfo, testClassInfo, null!);
-        classCleanupManager.MarkTestComplete(testMethodInfo, out bool shouldRunEndOfClassCleanup);
+        classCleanupManager.MarkTestComplete(testMethod, out bool shouldRunEndOfClassCleanup);
 
         // The cleanup should not run here yet, we have 1 remaining test to run.
         shouldRunEndOfClassCleanup.Should().BeFalse();
         classCleanupManager.ShouldRunEndOfAssemblyCleanup.Should().BeFalse();
 
-        classCleanupManager.MarkTestComplete(testMethodInfo, out shouldRunEndOfClassCleanup);
+        classCleanupManager.MarkTestComplete(testMethod, out shouldRunEndOfClassCleanup);
+        classCleanupManager.MarkClassComplete(fullClassName);
         // The cleanup should run here.
         shouldRunEndOfClassCleanup.Should().BeTrue();
         classCleanupManager.ShouldRunEndOfAssemblyCleanup.Should().BeTrue();