Allow AssemblyCleanup to have TestContext parameter

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

@@ -2,9 +2,12 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Collections.Concurrent;
+using System.Globalization;
 
 using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Helpers;
 using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel;
+using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices;
+using Microsoft.VisualStudio.TestPlatform.ObjectModel.Logging;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 
 namespace Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Execution;
@@ -63,8 +66,10 @@ public void MarkTestComplete(TestMethodInfo testMethodInfo, TestMethod testMetho
         }
     }
 
-    internal static void ForceCleanup(TypeCache typeCache)
+    internal static void ForceCleanup(TypeCache typeCache, IDictionary<string, object?> sourceLevelParameters, IMessageLogger logger)
     {
+        using var writer = new ThreadSafeStringWriter(CultureInfo.InvariantCulture, "context");
+        TestContext testContext = new TestContextImplementation(null, writer, sourceLevelParameters, logger);
         IEnumerable<TestClassInfo> classInfoCache = typeCache.ClassInfoListWithExecutableCleanupMethods;
         foreach (TestClassInfo classInfo in classInfoCache)
         {
@@ -74,7 +79,7 @@ internal static void ForceCleanup(TypeCache typeCache)
         IEnumerable<TestAssemblyInfo> assemblyInfoCache = typeCache.AssemblyInfoListWithExecutableCleanupMethods;
         foreach (TestAssemblyInfo assemblyInfo in assemblyInfoCache)
         {
-            assemblyInfo.ExecuteAssemblyCleanup();
+            assemblyInfo.ExecuteAssemblyCleanup(testContext);
         }
     }
 }

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

@@ -264,7 +264,7 @@ public void RunAssemblyInitialize(TestContext testContext)
     /// It is a replacement for RunAssemblyCleanup but as we are in a bug-fix version, we do not want to touch
     /// public API and so we introduced this method.
     /// </remarks>
-    internal void ExecuteAssemblyCleanup()
+    internal void ExecuteAssemblyCleanup(TestContext testContext)
     {
         if (AssemblyCleanupMethod == null)
         {
@@ -276,7 +276,17 @@ internal void ExecuteAssemblyCleanup()
             try
             {
                 AssemblyCleanupException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
-                     () => AssemblyCleanupMethod.InvokeAsSynchronousTask(null),
+                     () =>
+                     {
+                         if (AssemblyCleanupMethod.GetParameters().Length == 0)
+                         {
+                             AssemblyCleanupMethod.InvokeAsSynchronousTask(null);
+                         }
+                         else
+                         {
+                             AssemblyCleanupMethod.InvokeAsSynchronousTask(null, testContext);
+                         }
+                     },
                      new CancellationTokenSource(),
                      AssemblyCleanupMethodTimeoutMilliseconds,
                      AssemblyCleanupMethod,

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

@@ -414,7 +414,7 @@ private void ExecuteTestsInSource(IEnumerable<TestCase> tests, IRunContext? runC
 
         if (MSTestGracefulStopTestExecutionCapability.Instance.IsStopRequested)
         {
-            testRunner.ForceCleanup();
+            testRunner.ForceCleanup(sourceLevelParameters!, new RemotingMessageLogger(frameworkHandle));
         }
 
         PlatformServiceProvider.Instance.AdapterTraceLogger.LogInfo("Executed tests belonging to source {0}", source);
@@ -434,6 +434,8 @@ private void ExecuteTestsWithTestRunner(
             ? tests.OrderBy(t => t.GetManagedType()).ThenBy(t => t.GetManagedMethod())
             : tests;
 
+        var remotingMessageLogger = new RemotingMessageLogger(testExecutionRecorder);
+
         foreach (TestCase currentTest in orderedTests)
         {
             _testRunCancellationToken?.ThrowIfCancellationRequested();
@@ -465,7 +467,7 @@ private void ExecuteTestsWithTestRunner(
 
             // testRunner could be in a different AppDomain. We cannot pass the testExecutionRecorder directly.
             // Instead, we pass a proxy (remoting object) that is marshallable by ref.
-            UnitTestResult[] unitTestResult = testRunner.RunSingleTest(unitTestElement.TestMethod, testContextProperties, new RemotingMessageLogger(testExecutionRecorder));
+            UnitTestResult[] unitTestResult = testRunner.RunSingleTest(unitTestElement.TestMethod, testContextProperties, remotingMessageLogger);
 
             PlatformServiceProvider.Instance.AdapterTraceLogger.LogInfo("Executed test {0}", unitTestElement.TestMethod.Name);
 

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

@@ -266,7 +266,7 @@ private static void RunAssemblyCleanupIfNeeded(ITestContext testContext, ClassCl
                 IEnumerable<TestAssemblyInfo> assemblyInfoCache = typeCache.AssemblyInfoListWithExecutableCleanupMethods;
                 foreach (TestAssemblyInfo assemblyInfo in assemblyInfoCache)
                 {
-                    assemblyInfo.ExecuteAssemblyCleanup();
+                    assemblyInfo.ExecuteAssemblyCleanup(testContext.Context);
                 }
             }
             finally
@@ -364,5 +364,5 @@ private bool IsTestMethodRunnable(
         return true;
     }
 
-    internal void ForceCleanup() => ClassCleanupManager.ForceCleanup(_typeCache);
+    internal void ForceCleanup(IDictionary<string, object?> sourceLevelParameters, IMessageLogger logger) => ClassCleanupManager.ForceCleanup(_typeCache, sourceLevelParameters, logger);
 }

src/Adapter/MSTest.TestAdapter/Extensions/MethodInfoExtensions.cs

@@ -43,10 +43,16 @@ internal static bool HasCorrectClassOrAssemblyCleanupSignature(this MethodInfo m
 
         return
             method is { IsStatic: true, IsPublic: true } &&
-            (method.GetParameters().Length == 0) &&
+            HasCorrectClassOrAssemblyCleanupParameters(method) &&
             method.IsValidReturnType();
     }
 
+    private static bool HasCorrectClassOrAssemblyCleanupParameters(MethodInfo method)
+    {
+        ParameterInfo[] parameters = method.GetParameters();
+        return parameters.Length == 0 || (parameters.Length == 1 && parameters[0].ParameterType == typeof(TestContext));
+    }
+
     /// <summary>
     /// Verifies that the test Initialize/cleanup has the correct signature.
     /// </summary>

src/Adapter/MSTestAdapter.PlatformServices/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt

@@ -1 +1,2 @@
 #nullable enable
+Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.TestContextImplementation.TestContextImplementation(Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Interface.ObjectModel.ITestMethod? testMethod, System.IO.StringWriter! stringWriter, System.Collections.Generic.IDictionary<string!, object?>! properties) -> void

src/Adapter/MSTestAdapter.PlatformServices/Services/TestContextImplementation.cs

@@ -41,11 +41,6 @@ public class TestContextImplementation : TestContext, ITestContext
     private readonly StringWriter _stringWriter;
     private readonly ThreadSafeStringWriter? _threadSafeStringWriter;
 
-    /// <summary>
-    /// Test Method.
-    /// </summary>
-    private readonly ITestMethod _testMethod;
-
     /// <summary>
     /// Properties.
     /// </summary>
@@ -81,7 +76,7 @@ public class TestContextImplementation : TestContext, ITestContext
     /// <param name="stringWriter">The writer where diagnostic messages are written to.</param>
     /// <param name="properties">Properties/configuration passed in.</param>
     /// <param name="messageLogger">The message logger to use.</param>
-    internal TestContextImplementation(ITestMethod testMethod, StringWriter stringWriter, IDictionary<string, object?> properties, IMessageLogger messageLogger)
+    internal TestContextImplementation(ITestMethod? testMethod, StringWriter stringWriter, IDictionary<string, object?> properties, IMessageLogger messageLogger)
         : this(testMethod, stringWriter, properties)
         => _messageLogger = messageLogger;
 
@@ -91,27 +86,29 @@ internal TestContextImplementation(ITestMethod testMethod, StringWriter stringWr
     /// <param name="testMethod">The test method.</param>
     /// <param name="stringWriter">The writer where diagnostic messages are written to.</param>
     /// <param name="properties">Properties/configuration passed in.</param>
-    public TestContextImplementation(ITestMethod testMethod, StringWriter stringWriter, IDictionary<string, object?> properties)
+    public TestContextImplementation(ITestMethod? testMethod, StringWriter stringWriter, IDictionary<string, object?> properties)
     {
-        DebugEx.Assert(testMethod != null, "TestMethod is not null");
+        // testMethod can be null when running ForceCleanup (done when reaching --maximum-failed-tests.
         DebugEx.Assert(properties != null, "properties is not null");
 
 #if NETFRAMEWORK
         DebugEx.Assert(stringWriter != null, "StringWriter is not null");
 #endif
 
-        _testMethod = testMethod;
         _stringWriter = stringWriter;
 
         // Cannot get this type in constructor directly, because all signatures for all platforms need to be the same.
         _threadSafeStringWriter = stringWriter as ThreadSafeStringWriter;
-        _properties = new Dictionary<string, object?>(properties)
-        {
-            [FullyQualifiedTestClassNameLabel] = _testMethod.FullClassName,
-            [ManagedTypeLabel] = _testMethod.ManagedTypeName,
-            [ManagedMethodLabel] = _testMethod.ManagedMethodName,
-            [TestNameLabel] = _testMethod.Name,
-        };
+        _properties = testMethod is null
+            ? new Dictionary<string, object?>(properties)
+            : new Dictionary<string, object?>(properties)
+            {
+                [FullyQualifiedTestClassNameLabel] = testMethod.FullClassName,
+                [ManagedTypeLabel] = testMethod.ManagedTypeName,
+                [ManagedMethodLabel] = testMethod.ManagedMethodName,
+                [TestNameLabel] = testMethod.Name,
+            };
+
         _testResultFiles = [];
     }