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);
+
+                // We need to mark the class as complete to allow assembly cleanup to run.
+                // Note that we don't mark the class as complete when we are marking the last method.
+                // We need to only mark the class as complete after we are sure we have executed the class cleanup, to prevent concurrent runs of class cleanup and assembly cleanup.
+                _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,107 @@
+// 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(cancellationToken: TestContext.CancellationToken);
+
+        testHostResult.AssertExitCodeIs(ExitCodes.Success);
+        testHostResult.AssertOutputContainsSummary(failed: 0, passed: 2, skipped: 0);
+    }
+
+    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>
+
+</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()
+    {
+    }
+
+    [ClassCleanup]
+    public static void Cleanup1()
+    {
+        Thread.Sleep(4000);
+        ClassCleanupFinished = true;
+    }
+}
+
+[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()
+        => Assert.IsTrue(TestClass1.ClassCleanupFinished);
+}
+
+""";
+    }
+
+    public TestContext TestContext { get; set; }
+}