From 2b991212c50ddd9521c489499730d4b160b7f272 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Jare=C5=A1?= <me@jakubjares.com>
Date: Thu, 5 Sep 2024 13:52:00 +0200
Subject: [PATCH] Fix running cleanup after first test method when the same
 method is requested to run multiple times

---
 .../Execution/ClassCleanupManager.cs          |  6 +-
 .../Execution/ClassCleanupManagerTests.cs     | 63 +++++++++++++++++++
 2 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 test/UnitTests/MSTestAdapter.UnitTests/Execution/ClassCleanupManagerTests.cs

diff --git a/src/Adapter/MSTest.TestAdapter/Execution/ClassCleanupManager.cs b/src/Adapter/MSTest.TestAdapter/Execution/ClassCleanupManager.cs
index c64769ef5d..8f6d5d75f2 100644
--- a/src/Adapter/MSTest.TestAdapter/Execution/ClassCleanupManager.cs
+++ b/src/Adapter/MSTest.TestAdapter/Execution/ClassCleanupManager.cs
@@ -14,7 +14,7 @@ internal sealed class ClassCleanupManager
     private readonly ClassCleanupBehavior? _lifecycleFromMsTest;
     private readonly ClassCleanupBehavior _lifecycleFromAssembly;
     private readonly ReflectHelper _reflectHelper;
-    private readonly ConcurrentDictionary<string, HashSet<string>> _remainingTestsByClass;
+    private readonly ConcurrentDictionary<string, List<string>> _remainingTestsByClass;
 
     public ClassCleanupManager(
         IEnumerable<UnitTestElement> testsToRun,
@@ -27,7 +27,7 @@ public ClassCleanupManager(
             new(runnableTests.GroupBy(t => t.TestMethod.FullClassName)
                 .ToDictionary(
                     g => g.Key,
-                    g => new HashSet<string>(g.Select(t => t.TestMethod.UniqueName))));
+                    g => new List<string>(g.Select(t => t.TestMethod.UniqueName))));
         _lifecycleFromMsTest = lifecycleFromMsTest;
         _lifecycleFromAssembly = lifecycleFromAssembly;
         _reflectHelper = reflectHelper;
@@ -38,7 +38,7 @@ public ClassCleanupManager(
     public void MarkTestComplete(TestMethodInfo testMethodInfo, TestMethod testMethod, out bool shouldRunEndOfClassCleanup)
     {
         shouldRunEndOfClassCleanup = false;
-        if (!_remainingTestsByClass.TryGetValue(testMethodInfo.TestClassName, out HashSet<string>? testsByClass))
+        if (!_remainingTestsByClass.TryGetValue(testMethodInfo.TestClassName, out List<string>? testsByClass))
         {
             return;
         }
diff --git a/test/UnitTests/MSTestAdapter.UnitTests/Execution/ClassCleanupManagerTests.cs b/test/UnitTests/MSTestAdapter.UnitTests/Execution/ClassCleanupManagerTests.cs
new file mode 100644
index 0000000000..223668be9e
--- /dev/null
+++ b/test/UnitTests/MSTestAdapter.UnitTests/Execution/ClassCleanupManagerTests.cs
@@ -0,0 +1,63 @@
+// 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.Reflection;
+
+using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Execution;
+using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Helpers;
+using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+using Moq;
+
+using TestFramework.ForTestingMSTest;
+
+namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.UnitTests.Execution;
+
+public class ClassCleanupManagerTests : TestContainer
+{
+    public void AssemblyCleanupRunsAfterAllTestsFinishEvenIfWeScheduleTheSameTestMultipleTime()
+    {
+        ReflectHelper reflectHelper = Mock.Of<ReflectHelper>();
+        MethodInfo methodInfo = typeof(ClassCleanupManagerTests).GetMethod(nameof(FakeTestMethod), BindingFlags.Instance | BindingFlags.NonPublic);
+        MethodInfo classCleanupMethodInfo = typeof(ClassCleanupManagerTests).GetMethod(nameof(FakeClassCleanupMethod), BindingFlags.Instance | BindingFlags.NonPublic);
+        // Full class name must agree between unitTestElement.TestMethod.FullClassName and testMethod.FullClassName;
+        string fullClassName = methodInfo.DeclaringType.FullName;
+        TestMethod testMethod = new(nameof(FakeTestMethod), fullClassName, typeof(ClassCleanupManagerTests).Assembly.FullName, isAsync: false);
+
+        // Setting 2 of the same test to run, we should run assembly cleanup after both these tests
+        // finish, not after the first one finishes.
+        List<UnitTestElement> testsToRun = new()
+        {
+            new(testMethod),
+            new(testMethod),
+        };
+
+        var classCleanupManager = new ClassCleanupManager(testsToRun, ClassCleanupBehavior.EndOfClass, ClassCleanupBehavior.EndOfClass, reflectHelper);
+
+        TestClassInfo testClassInfo = new(typeof(ClassCleanupManagerTests), null, true, null, null, null)
+        {
+            // This needs to be set, to allow running class cleanup.
+            ClassCleanupMethod = classCleanupMethodInfo,
+        };
+        TestMethodInfo testMethodInfo = new(methodInfo, testClassInfo, null!);
+        classCleanupManager.MarkTestComplete(testMethodInfo, testMethod, out bool shouldRunEndOfClassCleanup);
+
+        // The cleanup should not run here yet, we have 1 remaining test to run.
+        Assert.IsFalse(shouldRunEndOfClassCleanup);
+        Assert.IsFalse(classCleanupManager.ShouldRunEndOfAssemblyCleanup);
+
+        classCleanupManager.MarkTestComplete(testMethodInfo, testMethod, out shouldRunEndOfClassCleanup);
+        // The cleanup should run here.
+        Assert.IsTrue(shouldRunEndOfClassCleanup);
+        Assert.IsTrue(classCleanupManager.ShouldRunEndOfAssemblyCleanup);
+    }
+
+    private void FakeTestMethod()
+    {
+    }
+
+    private void FakeClassCleanupMethod()
+    {
+    }
+}