diff --git a/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs b/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs index fbf08aae43..d8fc5529d3 100644 --- a/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs +++ b/src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs @@ -237,11 +237,13 @@ internal async Task 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 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 diff --git a/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs b/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs index 879021ff7e..bfa71fc029 100644 --- a/test/UnitTests/MSTestAdapter.PlatformServices.UnitTests/Execution/UnitTestRunnerTests.cs +++ b/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;