Add class/assembly cleanup/init messages to first/last test

src/Adapter/MSTest.TestAdapter/Execution/TestClassInfo.cs

@@ -286,7 +286,6 @@ public void RunClassInitialize(TestContext testContext)
                     finally
                     {
                         IsClassInitializeExecuted = true;
-                        (testContext as ITestContext)?.ClearDiagnosticMessages();
                     }
                 }
             }

src/Adapter/MSTest.TestAdapter/Execution/TestExecutionManager.cs

@@ -347,8 +347,6 @@ private void ExecuteTestsInSource(IEnumerable<TestCase> tests, IRunContext? runC
             ExecuteTestsWithTestRunner(testsToRun, frameworkHandle, source, sourceLevelParameters, testRunner);
         }
 
-        RunCleanup(frameworkHandle, testRunner);
-
         PlatformServiceProvider.Instance.AdapterTraceLogger.LogInfo("Executed tests belonging to source {0}", source);
     }
 
@@ -429,24 +427,6 @@ private void ExecuteTestsWithTestRunner(
         return testContextProperties;
     }
 
-    private static void RunCleanup(
-        ITestExecutionRecorder testExecutionRecorder,
-        UnitTestRunner testRunner)
-    {
-        // All cleanups (Class and Assembly) run at the end of test execution. Failures in these cleanup
-        // methods will be reported as Warnings.
-        PlatformServiceProvider.Instance.AdapterTraceLogger.LogInfo("Executing cleanup methods.");
-        var cleanupResult = testRunner.RunCleanup();
-        PlatformServiceProvider.Instance.AdapterTraceLogger.LogInfo("Executed cleanup methods.");
-        if (cleanupResult != null)
-        {
-            // Do not attach the standard output and error messages to any test result. It is not
-            // guaranteed that a test method of same class would have run last. We will end up
-            // adding stdout to test method of another class.
-            LogCleanupResult(testExecutionRecorder, cleanupResult);
-        }
-    }
-
     [SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "Requirement is to handle errors in user specified run parameters")]
     private void CacheSessionParameters(IRunContext? runContext, ITestExecutionRecorder testExecutionRecorder)
     {
@@ -474,44 +454,4 @@ private void CacheSessionParameters(IRunContext? runContext, ITestExecutionRecor
             testExecutionRecorder.SendMessage(TestMessageLevel.Error, ex.Message);
         }
     }
-
-    /// <summary>
-    /// Log the parameter warnings on the parameter logger.
-    /// </summary>
-    /// <param name="testExecutionRecorder">Handle to record test start/end/results/messages.</param>
-    /// <param name="result">Result of the run operation.</param>
-    private static void LogCleanupResult(ITestExecutionRecorder testExecutionRecorder, RunCleanupResult result)
-    {
-        DebugEx.Assert(testExecutionRecorder != null, "Logger should not be null");
-
-        if (!StringEx.IsNullOrWhiteSpace(result.StandardOut))
-        {
-            testExecutionRecorder.SendMessage(TestMessageLevel.Informational, result.StandardOut);
-        }
-
-        if (!StringEx.IsNullOrWhiteSpace(result.DebugTrace))
-        {
-            testExecutionRecorder.SendMessage(TestMessageLevel.Informational, result.DebugTrace);
-        }
-
-        if (!StringEx.IsNullOrWhiteSpace(result.StandardError))
-        {
-            testExecutionRecorder.SendMessage(
-                MSTestSettings.CurrentSettings.TreatClassAndAssemblyCleanupWarningsAsErrors ? TestMessageLevel.Error : TestMessageLevel.Warning,
-                result.StandardError);
-        }
-
-        if (result.Warnings != null)
-        {
-            foreach (string warning in result.Warnings)
-            {
-                if (!StringEx.IsNullOrWhiteSpace(warning))
-                {
-                    testExecutionRecorder.SendMessage(
-                        MSTestSettings.CurrentSettings.TreatClassAndAssemblyCleanupWarningsAsErrors ? TestMessageLevel.Error : TestMessageLevel.Warning,
-                        warning);
-                }
-            }
-        }
-    }
 }

src/Adapter/MSTest.TestAdapter/Execution/TypeCache.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Collections.Immutable;
 using System.Globalization;
 using System.Linq;
 using System.Reflection;
@@ -71,24 +72,24 @@ internal TypeCache(ReflectHelper reflectionHelper)
     /// <summary>
     /// Gets Class Info cache which has cleanup methods to execute.
     /// </summary>
-    public IEnumerable<TestClassInfo> ClassInfoListWithExecutableCleanupMethods =>
-        _classInfoCache.Values.Where(classInfo => classInfo.HasExecutableCleanupMethod).ToList();
+    public ImmutableArray<TestClassInfo> ClassInfoListWithExecutableCleanupMethods =>
+        _classInfoCache.Values.Where(classInfo => classInfo.HasExecutableCleanupMethod).ToImmutableArray();
 
     /// <summary>
     /// Gets Assembly Info cache which has cleanup methods to execute.
     /// </summary>
-    public IEnumerable<TestAssemblyInfo> AssemblyInfoListWithExecutableCleanupMethods =>
-        _testAssemblyInfoCache.Values.Where(assemblyInfo => assemblyInfo.HasExecutableCleanupMethod).ToList();
+    public ImmutableArray<TestAssemblyInfo> AssemblyInfoListWithExecutableCleanupMethods =>
+        _testAssemblyInfoCache.Values.Where(assemblyInfo => assemblyInfo.HasExecutableCleanupMethod).ToImmutableArray();
 
     /// <summary>
     /// Gets the set of cached assembly info values.
     /// </summary>
-    public IEnumerable<TestAssemblyInfo> AssemblyInfoCache => _testAssemblyInfoCache.Values.ToList();
+    public ImmutableArray<TestAssemblyInfo> AssemblyInfoCache => _testAssemblyInfoCache.Values.ToImmutableArray();
 
     /// <summary>
     /// Gets the set of cached class info values.
     /// </summary>
-    public IEnumerable<TestClassInfo> ClassInfoCache => _classInfoCache.Values.ToList();
+    public ImmutableArray<TestClassInfo> ClassInfoCache => _classInfoCache.Values.ToImmutableArray();
 
     /// <summary>
     /// Get the test method info corresponding to the parameter test Element.

src/Adapter/MSTest.TestAdapter/Execution/UnitTestRunner.cs

@@ -3,16 +3,17 @@
 
 using System;
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
 using System.Linq;
 using System.Reflection;
 using System.Security;
 
+using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Extensions;
 using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Helpers;
 using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel;
 using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices;
+using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Interface;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 
 using UTF = Microsoft.VisualStudio.TestTools.UnitTesting;
@@ -140,20 +141,15 @@ internal UnitTestResult[] RunSingleTest(TestMethod testMethod, IDictionary<strin
                 }
 
                 DebugEx.Assert(testMethodInfo is not null, "testMethodInfo should not be null.");
-                _classCleanupManager.MarkTestComplete(testMethodInfo, testMethod, out bool shouldRunClassCleanup);
-                if (shouldRunClassCleanup)
-                {
-                    testMethodInfo.Parent.RunClassCleanup(ClassCleanupBehavior.EndOfClass);
-                }
+                RunCleanupsIfNeeded(testContext, testMethodInfo, testMethod, notRunnableResult);
 
                 return notRunnableResult;
             }
 
             DebugEx.Assert(testMethodInfo is not null, "testMethodInfo should not be null.");
             var testMethodRunner = new TestMethodRunner(testMethodInfo, testMethod, testContext, MSTestSettings.CurrentSettings.CaptureDebugTraces);
             var result = testMethodRunner.Execute();
-
-            RunClassCleanupIfEndOfClass(testMethodInfo, testMethod, result);
+            RunCleanupsIfNeeded(testContext, testMethodInfo, testMethod, result);
             return result;
         }
         catch (TypeInspectionException ex)
@@ -168,12 +164,12 @@ internal UnitTestResult[] RunSingleTest(TestMethod testMethod, IDictionary<strin
     /// It returns any error information during the execution of the cleanup method.
     /// </summary>
     /// <returns> The <see cref="RunCleanupResult"/>. </returns>
-    internal RunCleanupResult? RunCleanup()
+    internal RunCleanupResult? RunClassAndAssemblyCleanup()
     {
         // No cleanup methods to execute, then return.
         var assemblyInfoCache = _typeCache.AssemblyInfoListWithExecutableCleanupMethods;
         var classInfoCache = _typeCache.ClassInfoListWithExecutableCleanupMethods;
-        if (!assemblyInfoCache.Any() && !classInfoCache.Any())
+        if (assemblyInfoCache.Length == 0 && classInfoCache.Length == 0)
         {
             return null;
         }
@@ -200,40 +196,72 @@ internal UnitTestResult[] RunSingleTest(TestMethod testMethod, IDictionary<strin
         return result;
     }
 
-    private void RunClassCleanupIfEndOfClass(TestMethodInfo testMethodInfo, TestMethod testMethod, UnitTestResult[] results)
+    private void RunCleanupsIfNeeded(ITestContext testContext, TestMethodInfo testMethodInfo, TestMethod testMethod, UnitTestResult[] results)
     {
         bool shouldRunClassCleanup = false;
-        _classCleanupManager?.MarkTestComplete(testMethodInfo, testMethod, out shouldRunClassCleanup);
-        if (!shouldRunClassCleanup)
-        {
-            return;
-        }
+        bool shouldRunClassAndAssemblyCleanup = false;
+        _classCleanupManager?.MarkTestComplete(testMethodInfo, testMethod, out shouldRunClassCleanup, out shouldRunClassAndAssemblyCleanup);
 
         try
         {
-            using LogMessageListener logListener =
-                new(MSTestSettings.CurrentSettings.CaptureDebugTraces);
-            try
+            using LogMessageListener logListener = new(MSTestSettings.CurrentSettings.CaptureDebugTraces);
+            if (shouldRunClassCleanup)
             {
-                // Class cleanup can throw exceptions in which case we need to ensure that we fail the test.
-                testMethodInfo.Parent.RunClassCleanup(ClassCleanupBehavior.EndOfClass);
+                try
+                {
+                    // Class cleanup can throw exceptions in which case we need to ensure that we fail the test.
+                    testMethodInfo.Parent.RunClassCleanup(ClassCleanupBehavior.EndOfClass);
+                }
+                finally
+                {
+                    string cleanupLogs = logListener.StandardOutput;
+                    string? cleanupTrace = logListener.DebugTrace;
+                    string cleanupErrorLogs = logListener.StandardError;
+                    var cleanupTestContextMessages = testContext.GetAndClearDiagnosticMessages();
+
+                    if (results.Length > 0)
+                    {
+                        var lastResult = results[results.Length - 1];
+                        lastResult.StandardOut += cleanupLogs;
+                        lastResult.StandardError += cleanupErrorLogs;
+                        lastResult.DebugTrace += cleanupTrace;
+                        lastResult.TestContextMessages += cleanupTestContextMessages;
+                    }
+                }
             }
-            finally
+
+            if (shouldRunClassAndAssemblyCleanup)
             {
-                string cleanupLogs = logListener.StandardOutput;
-                string? cleanupTrace = logListener.DebugTrace;
-                string cleanupErrorLogs = logListener.StandardError;
-                var lastResult = results[results.Length - 1];
-                lastResult.StandardOut += cleanupLogs;
-                lastResult.StandardError += cleanupErrorLogs;
-                lastResult.DebugTrace += cleanupTrace;
+                try
+                {
+                    RunClassAndAssemblyCleanup();
+                }
+                finally
+                {
+                    string cleanupLogs = logListener.StandardOutput;
+                    string? cleanupTrace = logListener.DebugTrace;
+                    string cleanupErrorLogs = logListener.StandardError;
+                    var cleanupTestContextMessages = testContext.GetAndClearDiagnosticMessages();
+
+                    if (results.Length > 0)
+                    {
+                        var lastResult = results[results.Length - 1];
+                        lastResult.StandardOut += cleanupLogs;
+                        lastResult.StandardError += cleanupErrorLogs;
+                        lastResult.DebugTrace += cleanupTrace;
+                        lastResult.TestContextMessages += cleanupTestContextMessages;
+                    }
+                }
             }
         }
         catch (Exception e)
         {
-            results[results.Length - 1].Outcome = ObjectModel.UnitTestOutcome.Failed;
-            results[results.Length - 1].ErrorMessage = e.Message;
-            results[results.Length - 1].ErrorStackTrace = e.StackTrace;
+            if (results.Length > 0)
+            {
+                results[results.Length - 1].Outcome = ObjectModel.UnitTestOutcome.Failed;
+                results[results.Length - 1].ErrorMessage = e.Message;
+                results[results.Length - 1].ErrorStackTrace = e.StackTrace;
+            }
         }
     }
 
@@ -363,9 +391,9 @@ public ClassCleanupManager(
             _reflectHelper = reflectHelper ?? new ReflectHelper();
         }
 
-        public void MarkTestComplete(TestMethodInfo testMethodInfo, TestMethod testMethod, out bool shouldCleanup)
+        public void MarkTestComplete(TestMethodInfo testMethodInfo, TestMethod testMethod, out bool shouldRunClassCleanup, out bool shouldRunAssemblyCleanup)
         {
-            shouldCleanup = false;
+            shouldRunClassCleanup = false;
             var testsByClass = _remainingTestsByClass[testMethodInfo.TestClassName];
             lock (testsByClass)
             {
@@ -376,8 +404,11 @@ public void MarkTestComplete(TestMethodInfo testMethodInfo, TestMethod testMetho
                         ?? _lifecycleFromMsTest
                         ?? _lifecycleFromAssembly;
 
-                    shouldCleanup = cleanupLifecycle == ClassCleanupBehavior.EndOfClass;
+                    shouldRunClassCleanup = cleanupLifecycle == ClassCleanupBehavior.EndOfClass;
+                    _remainingTestsByClass.Remove(testMethodInfo.TestClassName);
                 }
+
+                shouldRunAssemblyCleanup = _remainingTestsByClass.Count == 0;
             }
         }
     }

src/Adapter/MSTest.TestAdapter/ObjectModel/UnitTestResult.cs

@@ -183,14 +183,14 @@ internal TestResult ToTestResult(TestCase testCase, DateTimeOffset startTime, Da
 
         if (!StringEx.IsNullOrEmpty(DebugTrace))
         {
-            string debugTraceMessagesInStdOut = string.Format(CultureInfo.InvariantCulture, "\n\n{0}\n{1}", Resource.DebugTraceBanner, DebugTrace);
+            string debugTraceMessagesInStdOut = string.Format(CultureInfo.CurrentCulture, "{2}{2}{0}{2}{1}", Resource.DebugTraceBanner, DebugTrace, Environment.NewLine);
             TestResultMessage debugTraceMessage = new(TestResultMessage.StandardOutCategory, debugTraceMessagesInStdOut);
             testResult.Messages.Add(debugTraceMessage);
         }
 
         if (!StringEx.IsNullOrEmpty(TestContextMessages))
         {
-            string testContextMessagesInStdOut = string.Format(CultureInfo.InvariantCulture, "\n\n{0}\n{1}", Resource.TestContextMessageBanner, TestContextMessages);
+            string testContextMessagesInStdOut = string.Format(CultureInfo.InvariantCulture, "{2}{2}{0}{2}{1}", Resource.TestContextMessageBanner, TestContextMessages, Environment.NewLine);
             TestResultMessage testContextMessage = new(TestResultMessage.StandardOutCategory, testContextMessagesInStdOut);
             testResult.Messages.Add(testContextMessage);
         }

test/E2ETests/DiscoveryAndExecutionTests/OutputTests.cs

@@ -59,8 +59,8 @@ private void ValidateOutputForClass(string testAssembly, string className)
         ValidateInitializationsAndCleanups(testResults);
     }
 
-    private static readonly Func<TestResultMessage, bool> IsDebugMessage = m => m.Category == "StdOutMsgs" && m.Text.StartsWith("\n\nDebug Trace:\n");
-    private static readonly Func<TestResultMessage, bool> IsStandardOutputMessage = m => m.Category == "StdOutMsgs" && !m.Text.StartsWith("\n\nDebug Trace:\n");
+    private static readonly Func<TestResultMessage, bool> IsDebugMessage = m => m.Category == "StdOutMsgs" && m.Text.StartsWith("\r\n\r\nDebug Trace:\r\n");
+    private static readonly Func<TestResultMessage, bool> IsStandardOutputMessage = m => m.Category == "StdOutMsgs" && !m.Text.StartsWith("\r\n\r\nDebug Trace:\r\n");
     private static readonly Func<TestResultMessage, bool> IsStandardErrorMessage = m => m.Category == "StdErrMsgs";
 
     private static void ValidateOutputsAreNotMixed(ReadOnlyCollection<TestResult> testResults, string methodName, string[] shouldNotContain)