[perf-improver] perf: defer class-cleanup TestContextImplementation allocation to last test in class

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

@@ -237,11 +237,13 @@ internal async Task<TestResult[]> RunSingleTestAsync(UnitTestElement unitTestEle
                 }
             }
 
-            testContextForClassCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, testMethod.FullClassName, testContextProperties, messageLogger, testContextForTestExecution.Context.CurrentTestOutcome);
-
             _classCleanupManager.MarkTestComplete(testMethod, out bool isLastTestInClass);
             if (isLastTestInClass)
             {
+                // Defer TestContextImplementation allocation to only the last test in each class,
+                // saving one dict-copy + CancellationTokenRegistration per non-last test.
+                testContextForClassCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, testMethod.FullClassName, testContextProperties, messageLogger, testContextForTestExecution.Context.CurrentTestOutcome);
+
                 if (testMethodInfo is not null)
                 {
                     // Flow properties set during AssemblyInitialize and ClassInitialize so the
@@ -276,6 +278,10 @@ internal async Task<TestResult[]> RunSingleTestAsync(UnitTestElement unitTestEle
             if (testMethodInfo?.Parent.Parent.IsAssemblyInitializeExecuted == true &&
                 _classCleanupManager.ShouldRunEndOfAssemblyCleanup)
             {
+                // testContextForClassCleanup is guaranteed non-null here: ShouldRunEndOfAssemblyCleanup
+                // becomes true only after MarkClassComplete, which is called exclusively inside the
+                // isLastTestInClass block above — where testContextForClassCleanup is allocated.
+                DebugEx.Assert(testContextForClassCleanup is not null, "testContextForClassCleanup should not be null when running assembly cleanup.");
                 testContextForAssemblyCleanup = PlatformServiceProvider.Instance.GetTestContext(testMethod: null, null, testContextProperties, messageLogger, testContextForClassCleanup.Context.CurrentTestOutcome);
 
                 // Flow properties set during AssemblyInitialize so the AssemblyCleanup method

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

@@ -385,6 +385,48 @@ public async Task RunSingleTestWhenClassInitializeAlreadyFailedShouldTakeFastPat
         contextsAllocatedForSecondRun.Should().BeLessThan(contextsAllocatedForFirstRun);
     }
 
+    public async Task RunSingleTestShouldDeferClassCleanupContextAllocationToLastTestInClass()
+    {
+        Type type = typeof(DummyTestClassWithCleanupMethods);
+        TestMethod testMethod1 = CreateTestMethod("TestMethod", type.FullName!, "A", displayName: null);
+        TestMethod testMethod2 = CreateTestMethod("TestMethod2", type.FullName!, "A", displayName: null);
+        TestMethod testMethod3 = CreateTestMethod("TestMethod3", type.FullName!, "A", displayName: null);
+        var unitTestElement1 = new UnitTestElement(testMethod1);
+        var unitTestElement2 = new UnitTestElement(testMethod2);
+        var unitTestElement3 = new UnitTestElement(testMethod3);
+        UnitTestRunner unitTestRunner = CreateUnitTestRunner([unitTestElement1, unitTestElement2, unitTestElement3]);
+
+        _testablePlatformServiceProvider.MockFileOperations.Setup(fo => fo.LoadAssembly("A"))
+            .Returns(Assembly.GetExecutingAssembly());
+
+        DummyTestClassWithCleanupMethods.ClassCleanupMethodBody = () => { };
+        DummyTestClassWithCleanupMethods.AssemblyCleanupMethodBody = () => { };
+
+        // Run the first test to warm up assembly/class initialize (slow path, sets cached results).
+        TestResult[] firstResults = await unitTestRunner.RunSingleTestAsync(unitTestElement1, _testRunParameters, null!);
+        firstResults[0].Outcome.Should().Be(UnitTestOutcome.Passed);
+
+        // Second test (non-last): both assembly-init and class-init take the fast path (no contexts),
+        // and the class-cleanup context is NOT allocated (deferred to the last test).
+        int countBeforeSecond = _testablePlatformServiceProvider.GetTestContextCallCount;
+        TestResult[] secondResults = await unitTestRunner.RunSingleTestAsync(unitTestElement2, _testRunParameters, null!);
+        secondResults[0].Outcome.Should().Be(UnitTestOutcome.Passed);
+        int allocationsForSecond = _testablePlatformServiceProvider.GetTestContextCallCount - countBeforeSecond;
+
+        // Third test (last in class): takes fast paths for init, but allocates a class-cleanup context
+        // and an assembly-cleanup context. Allocation count must exceed the middle test's count.
+        int countBeforeThird = _testablePlatformServiceProvider.GetTestContextCallCount;
+        TestResult[] thirdResults = await unitTestRunner.RunSingleTestAsync(unitTestElement3, _testRunParameters, null!);
+        thirdResults[0].Outcome.Should().Be(UnitTestOutcome.Passed);
+        int allocationsForThird = _testablePlatformServiceProvider.GetTestContextCallCount - countBeforeThird;
+
+        // The last-test run allocates exactly two more contexts than a non-last test: the class-cleanup
+        // context and the assembly-cleanup context, both of which are deferred from all non-last tests.
+        // Asserting the exact delta (rather than a loose "greater than") catches the regression this test
+        // guards against, where a non-last test mistakenly allocates the deferred class-cleanup context.
+        allocationsForThird.Should().Be(allocationsForSecond + 2);
+    }
+
     #endregion
 
     #region private helpers
@@ -498,6 +540,12 @@ private class DummyTestClassWithCleanupMethods
 
         [TestMethod]
         public void TestMethod() => TestMethodBody?.Invoke(TestContext);
+
+        [TestMethod]
+        public void TestMethod2() => TestMethodBody?.Invoke(TestContext);
+
+        [TestMethod]
+        public void TestMethod3() => TestMethodBody?.Invoke(TestContext);
     }
 
     private class DummyTestClassAttribute : TestClassAttribute;