Allow AssemblyCleanup/ClassCleanup 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,18 +66,20 @@ 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)
         {
-            classInfo.ExecuteClassCleanup();
+            classInfo.ExecuteClassCleanup(testContext);
         }
 
         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/TestClassInfo.cs

@@ -491,12 +491,12 @@ UnitTestResult DoRun()
                 try
                 {
                     classCleanupMethod = ClassCleanupMethod;
-                    ClassCleanupException = classCleanupMethod is not null ? InvokeCleanupMethod(classCleanupMethod, BaseClassCleanupMethods.Count) : null;
+                    ClassCleanupException = classCleanupMethod is not null ? InvokeCleanupMethod(classCleanupMethod, BaseClassCleanupMethods.Count, null!) : null;
                     var baseClassCleanupQueue = new Queue<MethodInfo>(BaseClassCleanupMethods);
                     while (baseClassCleanupQueue.Count > 0 && ClassCleanupException is null)
                     {
                         classCleanupMethod = baseClassCleanupQueue.Dequeue();
-                        ClassCleanupException = InvokeCleanupMethod(classCleanupMethod, baseClassCleanupQueue.Count);
+                        ClassCleanupException = InvokeCleanupMethod(classCleanupMethod, baseClassCleanupQueue.Count, null!);
                     }
 
                     IsClassCleanupExecuted = ClassCleanupException is null;
@@ -548,7 +548,7 @@ UnitTestResult DoRun()
     /// This is a replacement for RunClassCleanup but as we are on a bug fix version, we do not want to change
     /// the public API, hence this method.
     /// </remarks>
-    internal void ExecuteClassCleanup()
+    internal void ExecuteClassCleanup(TestContext testContext)
     {
         if ((ClassCleanupMethod is null && BaseClassCleanupMethods.Count == 0)
             || IsClassCleanupExecuted)
@@ -575,7 +575,7 @@ internal void ExecuteClassCleanup()
                 {
                     if (!ReflectHelper.Instance.IsNonDerivedAttributeDefined<IgnoreAttribute>(classCleanupMethod.DeclaringType!, false))
                     {
-                        ClassCleanupException = InvokeCleanupMethod(classCleanupMethod, remainingCleanupCount: BaseClassCleanupMethods.Count);
+                        ClassCleanupException = InvokeCleanupMethod(classCleanupMethod, remainingCleanupCount: BaseClassCleanupMethods.Count, testContext);
                     }
                 }
 
@@ -586,7 +586,7 @@ internal void ExecuteClassCleanup()
                         classCleanupMethod = BaseClassCleanupMethods[i];
                         if (!ReflectHelper.Instance.IsNonDerivedAttributeDefined<IgnoreAttribute>(classCleanupMethod.DeclaringType!, false))
                         {
-                            ClassCleanupException = InvokeCleanupMethod(classCleanupMethod, remainingCleanupCount: BaseClassCleanupMethods.Count - 1 - i);
+                            ClassCleanupException = InvokeCleanupMethod(classCleanupMethod, remainingCleanupCount: BaseClassCleanupMethods.Count - 1 - i, testContext);
                             if (ClassCleanupException is not null)
                             {
                                 break;
@@ -707,7 +707,7 @@ void DoRun()
                 using LogMessageListener logListener = new(MSTestSettings.CurrentSettings.CaptureDebugTraces);
                 try
                 {
-                    ExecuteClassCleanup();
+                    ExecuteClassCleanup(testContext.Context);
                 }
                 finally
                 {
@@ -745,7 +745,7 @@ void DoRun()
         }
     }
 
-    private TestFailedException? InvokeCleanupMethod(MethodInfo methodInfo, int remainingCleanupCount)
+    private TestFailedException? InvokeCleanupMethod(MethodInfo methodInfo, int remainingCleanupCount, TestContext testContext)
     {
         TimeoutInfo? timeout = null;
         if (ClassCleanupMethodTimeoutMilliseconds.TryGetValue(methodInfo, out TimeoutInfo localTimeout))
@@ -754,7 +754,17 @@ void DoRun()
         }
 
         return FixtureMethodRunner.RunWithTimeoutAndCancellation(
-            () => methodInfo.InvokeAsSynchronousTask(null),
+            () =>
+            {
+                if (methodInfo.GetParameters().Length == 0)
+                {
+                    methodInfo.InvokeAsSynchronousTask(null);
+                }
+                else
+                {
+                    methodInfo.InvokeAsSynchronousTask(null, testContext);
+                }
+            },
             new CancellationTokenSource(),
             timeout,
             methodInfo,

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

@@ -260,13 +260,13 @@ private static void RunAssemblyCleanupIfNeeded(ITestContext testContext, ClassCl
                 IEnumerable<TestClassInfo> classInfoCache = typeCache.ClassInfoListWithExecutableCleanupMethods;
                 foreach (TestClassInfo classInfo in classInfoCache)
                 {
-                    classInfo.ExecuteClassCleanup();
+                    classInfo.ExecuteClassCleanup(testContext.Context);
                 }
 
                 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/PublicAPI.Unshipped.txt

@@ -1,3 +1,5 @@
 #nullable enable
 Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.Interface.ITestContext.DisplayMessage(Microsoft.VisualStudio.TestTools.UnitTesting.MessageLevel messageLevel, string! message) -> void
+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
 override Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.TestContextImplementation.DisplayMessage(Microsoft.VisualStudio.TestTools.UnitTesting.MessageLevel messageLevel, string! message) -> void
+*REMOVED*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 = [];
     }
 

src/Analyzers/MSTest.Analyzers/AssemblyCleanupShouldBeValidAnalyzer.cs

@@ -41,15 +41,23 @@ public override void Initialize(AnalysisContext context)
             {
                 INamedTypeSymbol? taskSymbol = context.Compilation.GetOrCreateTypeByMetadataName(WellKnownTypeNames.SystemThreadingTasksTask);
                 INamedTypeSymbol? valueTaskSymbol = context.Compilation.GetOrCreateTypeByMetadataName(WellKnownTypeNames.SystemThreadingTasksValueTask);
+                INamedTypeSymbol? testContextSymbol = context.Compilation.GetOrCreateTypeByMetadataName(WellKnownTypeNames.MicrosoftVisualStudioTestToolsUnitTestingTestContext);
                 bool canDiscoverInternals = context.Compilation.CanDiscoverInternals();
                 context.RegisterSymbolAction(
-                    context => AnalyzeSymbol(context, assemblyCleanupAttributeSymbol, testClassAttributeSymbol, taskSymbol, valueTaskSymbol, canDiscoverInternals),
+                    context => AnalyzeSymbol(context, assemblyCleanupAttributeSymbol, testClassAttributeSymbol, taskSymbol, valueTaskSymbol, testContextSymbol, canDiscoverInternals),
                     SymbolKind.Method);
             }
         });
     }
 
-    private static void AnalyzeSymbol(SymbolAnalysisContext context, INamedTypeSymbol assemblyCleanupAttributeSymbol, INamedTypeSymbol testClassAttributeSymbol, INamedTypeSymbol? taskSymbol, INamedTypeSymbol? valueTaskSymbol, bool canDiscoverInternals)
+    private static void AnalyzeSymbol(
+        SymbolAnalysisContext context,
+        INamedTypeSymbol assemblyCleanupAttributeSymbol,
+        INamedTypeSymbol testClassAttributeSymbol,
+        INamedTypeSymbol? taskSymbol,
+        INamedTypeSymbol? valueTaskSymbol,
+        INamedTypeSymbol? testContextSymbol,
+        bool canDiscoverInternals)
     {
         var methodSymbol = (IMethodSymbol)context.Symbol;
 
@@ -58,7 +66,7 @@ private static void AnalyzeSymbol(SymbolAnalysisContext context, INamedTypeSymbo
             return;
         }
 
-        if (!methodSymbol.HasValidFixtureMethodSignature(taskSymbol, valueTaskSymbol, canDiscoverInternals, shouldBeStatic: true, allowGenericType: false, testContextSymbol: null, testClassAttributeSymbol, fixtureAllowInheritedTestClass: false, out bool isFixable))
+        if (!methodSymbol.HasValidFixtureMethodSignature(taskSymbol, valueTaskSymbol, canDiscoverInternals, shouldBeStatic: true, allowGenericType: false, FixtureParameterMode.OptionalTestContext, testContextSymbol, testClassAttributeSymbol, fixtureAllowInheritedTestClass: false, out bool isFixable))
         {
             context.ReportDiagnostic(isFixable
                 ? methodSymbol.CreateDiagnostic(Rule, methodSymbol.Name)

src/Analyzers/MSTest.Analyzers/AssemblyInitializeShouldBeValidAnalyzer.cs

@@ -57,7 +57,7 @@ private static void AnalyzeSymbol(SymbolAnalysisContext context, INamedTypeSymbo
 
         if (methodSymbol.IsAssemblyInitializeMethod(assemblyInitializeAttributeSymbol)
             && !methodSymbol.HasValidFixtureMethodSignature(taskSymbol, valueTaskSymbol, canDiscoverInternals, shouldBeStatic: true,
-                allowGenericType: false, testContextSymbol, testClassAttributeSymbol, fixtureAllowInheritedTestClass: false, out bool isFixable))
+                allowGenericType: false, FixtureParameterMode.MustHaveTestContext, testContextSymbol, testClassAttributeSymbol, fixtureAllowInheritedTestClass: false, out bool isFixable))
         {
             context.ReportDiagnostic(isFixable
                 ? methodSymbol.CreateDiagnostic(Rule, methodSymbol.Name)