Flow execution context across fixture methods when using timeout

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

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Diagnostics.CodeAnalysis;
@@ -8,6 +8,7 @@
 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.TestTools.UnitTesting;
 
 using UnitTestOutcome = Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel.UnitTestOutcome;
@@ -140,11 +141,12 @@ public void RunAssemblyInitialize(TestContext testContext)
                 {
                     try
                     {
-                        AssemblyInitializationException = MethodRunner.RunWithTimeoutAndCancellation(
+                        AssemblyInitializationException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
                             () => AssemblyInitializeMethod.InvokeAsSynchronousTask(null, testContext),
                             testContext.CancellationTokenSource,
                             AssemblyInitializeMethodTimeoutMilliseconds,
                             AssemblyInitializeMethod,
+                            new AssemblyExecutionContextScope(isCleanup: false),
                             Resource.AssemblyInitializeWasCancelled,
                             Resource.AssemblyInitializeTimedOut);
                     }
@@ -213,11 +215,12 @@ public void RunAssemblyInitialize(TestContext testContext)
         {
             try
             {
-                assemblyCleanupException = MethodRunner.RunWithTimeoutAndCancellation(
+                assemblyCleanupException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
                      () => AssemblyCleanupMethod.InvokeAsSynchronousTask(null),
                      new CancellationTokenSource(),
                      AssemblyCleanupMethodTimeoutMilliseconds,
                      AssemblyCleanupMethod,
+                     new AssemblyExecutionContextScope(isCleanup: true),
                      Resource.AssemblyCleanupWasCancelled,
                      Resource.AssemblyCleanupTimedOut);
             }
@@ -269,11 +272,12 @@ internal void ExecuteAssemblyCleanup()
         {
             try
             {
-                assemblyCleanupException = MethodRunner.RunWithTimeoutAndCancellation(
+                assemblyCleanupException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
                      () => AssemblyCleanupMethod.InvokeAsSynchronousTask(null),
                      new CancellationTokenSource(),
                      AssemblyCleanupMethodTimeoutMilliseconds,
                      AssemblyCleanupMethod,
+                     new AssemblyExecutionContextScope(isCleanup: true),
                      Resource.AssemblyCleanupWasCancelled,
                      Resource.AssemblyCleanupTimedOut);
             }

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

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Diagnostics.CodeAnalysis;
@@ -8,6 +8,7 @@
 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.TestTools.UnitTesting;
 
 using ObjectModelUnitTestOutcome = Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel.UnitTestOutcome;
@@ -363,11 +364,12 @@ public void RunClassInitialize(TestContext testContext)
             timeout = localTimeout;
         }
 
-        return MethodRunner.RunWithTimeoutAndCancellation(
+        return FixtureMethodRunner.RunWithTimeoutAndCancellation(
             () => methodInfo.InvokeAsSynchronousTask(null, testContext),
             testContext.CancellationTokenSource,
             timeout,
             methodInfo,
+            new ClassExecutionContextScope(ClassType),
             Resource.ClassInitializeWasCancelled,
             Resource.ClassInitializeTimedOut);
     }
@@ -404,12 +406,12 @@ public void RunClassInitialize(TestContext testContext)
                 try
                 {
                     classCleanupMethod = ClassCleanupMethod;
-                    ClassCleanupException = classCleanupMethod is not null ? InvokeCleanupMethod(classCleanupMethod) : null;
+                    ClassCleanupException = classCleanupMethod is not null ? InvokeCleanupMethod(classCleanupMethod, BaseClassCleanupMethodsStack.Count) : null;
                     var baseClassCleanupQueue = new Queue<MethodInfo>(BaseClassCleanupMethodsStack);
                     while (baseClassCleanupQueue.Count > 0 && ClassCleanupException is null)
                     {
                         classCleanupMethod = baseClassCleanupQueue.Dequeue();
-                        ClassCleanupException = classCleanupMethod is not null ? InvokeCleanupMethod(classCleanupMethod) : null;
+                        ClassCleanupException = classCleanupMethod is not null ? InvokeCleanupMethod(classCleanupMethod, baseClassCleanupQueue.Count) : null;
                     }
 
                     IsClassCleanupExecuted = ClassCleanupException is null;
@@ -482,12 +484,16 @@ internal void ExecuteClassCleanup()
             try
             {
                 classCleanupMethod = ClassCleanupMethod;
-                ClassCleanupException = classCleanupMethod is not null ? InvokeCleanupMethod(classCleanupMethod) : null;
+                ClassCleanupException = classCleanupMethod is not null
+                    ? InvokeCleanupMethod(classCleanupMethod, BaseClassCleanupMethodsStack.Count)
+                    : null;
                 var baseClassCleanupQueue = new Queue<MethodInfo>(BaseClassCleanupMethodsStack);
                 while (baseClassCleanupQueue.Count > 0 && ClassCleanupException is null)
                 {
                     classCleanupMethod = baseClassCleanupQueue.Dequeue();
-                    ClassCleanupException = classCleanupMethod is not null ? InvokeCleanupMethod(classCleanupMethod) : null;
+                    ClassCleanupException = classCleanupMethod is not null
+                        ? InvokeCleanupMethod(classCleanupMethod, baseClassCleanupQueue.Count)
+                        : null;
                 }
 
                 IsClassCleanupExecuted = ClassCleanupException is null;
@@ -535,19 +541,20 @@ internal void ExecuteClassCleanup()
         throw testFailedException;
     }
 
-    private TestFailedException? InvokeCleanupMethod(MethodInfo methodInfo)
+    private TestFailedException? InvokeCleanupMethod(MethodInfo methodInfo, int remainingCleanupCount)
     {
         int? timeout = null;
         if (ClassCleanupMethodTimeoutMilliseconds.TryGetValue(methodInfo, out int localTimeout))
         {
             timeout = localTimeout;
         }
 
-        return MethodRunner.RunWithTimeoutAndCancellation(
+        return FixtureMethodRunner.RunWithTimeoutAndCancellation(
             () => methodInfo.InvokeAsSynchronousTask(null),
             new CancellationTokenSource(),
             timeout,
             methodInfo,
+            new ClassExecutionContextScope(ClassType, remainingCleanupCount),
             Resource.ClassCleanupWasCancelled,
             Resource.ClassCleanupTimedOut);
     }

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

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Diagnostics;
@@ -10,6 +10,7 @@
 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.TestTools.UnitTesting;
 
 using ObjectModelUnitTestOutcome = Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel.UnitTestOutcome;
@@ -239,7 +240,17 @@ private TestResult ExecuteInternal(object?[]? arguments)
                     if (RunTestInitializeMethod(classInstance, result))
                     {
                         hasTestInitializePassed = true;
-                        TestMethod.InvokeAsSynchronousTask(classInstance, arguments);
+                        if (IsTimeoutSet)
+                        {
+                            ExecutionContextService.RunActionOnContext(
+                                () => TestMethod.InvokeAsSynchronousTask(classInstance, arguments),
+                                new InstanceExecutionContextScope(classInstance, Parent.ClassType));
+                        }
+                        else
+                        {
+                            TestMethod.InvokeAsSynchronousTask(classInstance, arguments);
+                        }
+
                         result.Outcome = UTF.UnitTestOutcome.Passed;
                     }
                 }
@@ -448,12 +459,16 @@ private void RunTestCleanupMethod(object classInstance, TestResult result)
             {
                 // Test cleanups are called in the order of discovery
                 // Current TestClass -> Parent -> Grandparent
-                testCleanupException = testCleanupMethod is not null ? InvokeCleanupMethod(testCleanupMethod, classInstance) : null;
+                testCleanupException = testCleanupMethod is not null
+                    ? InvokeCleanupMethod(testCleanupMethod, classInstance, Parent.BaseTestCleanupMethodsQueue.Count)
+                    : null;
                 var baseTestCleanupQueue = new Queue<MethodInfo>(Parent.BaseTestCleanupMethodsQueue);
                 while (baseTestCleanupQueue.Count > 0 && testCleanupException is null)
                 {
                     testCleanupMethod = baseTestCleanupQueue.Dequeue();
-                    testCleanupException = testCleanupMethod is not null ? InvokeCleanupMethod(testCleanupMethod, classInstance) : null;
+                    testCleanupException = testCleanupMethod is not null
+                        ? InvokeCleanupMethod(testCleanupMethod, classInstance, baseTestCleanupQueue.Count)
+                        : null;
                 }
             }
             finally
@@ -635,28 +650,30 @@ private bool RunTestInitializeMethod(object classInstance, TestResult result)
             timeout = localTimeout;
         }
 
-        return MethodRunner.RunWithTimeoutAndCancellation(
+        return FixtureMethodRunner.RunWithTimeoutAndCancellation(
             () => methodInfo.InvokeAsSynchronousTask(classInstance, null),
             new CancellationTokenSource(),
             timeout,
             methodInfo,
+            new InstanceExecutionContextScope(classInstance, Parent.ClassType),
             Resource.TestInitializeWasCancelled,
             Resource.TestInitializeTimedOut);
     }
 
-    private TestFailedException? InvokeCleanupMethod(MethodInfo methodInfo, object classInstance)
+    private TestFailedException? InvokeCleanupMethod(MethodInfo methodInfo, object classInstance, int remainingCleanupCount)
     {
         int? timeout = null;
         if (Parent.TestCleanupMethodTimeoutMilliseconds.TryGetValue(methodInfo, out int localTimeout))
         {
             timeout = localTimeout;
         }
 
-        return MethodRunner.RunWithTimeoutAndCancellation(
+        return FixtureMethodRunner.RunWithTimeoutAndCancellation(
             () => methodInfo.InvokeAsSynchronousTask(classInstance, null),
             new CancellationTokenSource(),
             timeout,
             methodInfo,
+            new InstanceExecutionContextScope(classInstance, Parent.ClassType, remainingCleanupCount),
             Resource.TestCleanupWasCancelled,
             Resource.TestCleanupTimedOut);
     }

src/Adapter/MSTest.TestAdapter/Helpers/MethodRunner.cs → src/Adapter/MSTest.TestAdapter/Helpers/FixtureMethodRunner.cs

@@ -6,17 +6,20 @@
 using System.Runtime.InteropServices;
 
 using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel;
+using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices;
 
 using UnitTestOutcome = Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel.UnitTestOutcome;
 
 namespace Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Helpers;
 
-internal static class MethodRunner
+internal static class FixtureMethodRunner
 {
     internal static TestFailedException? RunWithTimeoutAndCancellation(
         Action action, CancellationTokenSource cancellationTokenSource, int? timeout, MethodInfo methodInfo,
-        string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
+        IExecutionContextScope executionContextScope, string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
     {
+        // If no timeout is specified, we can run the action directly. This avoids any overhead of creating a task/thread and
+        // ensures that the execution context is preserved (as we run the action on the current thread).
         if (timeout is null)
         {
             action();
@@ -25,13 +28,13 @@ internal static class MethodRunner
 
         // We need to start a thread to handle "cancellation" and "timeout" scenarios.
         return RuntimeInformation.IsOSPlatform(OSPlatform.Windows) && Thread.CurrentThread.GetApartmentState() == ApartmentState.STA
-                ? RunWithTimeoutAndCancellationWithSTAThread(action, cancellationTokenSource, timeout, methodInfo, methodCancelledMessageFormat, methodTimedOutMessageFormat)
-                : RunWithTimeoutAndCancellationWithThreadPool(action, cancellationTokenSource, timeout, methodInfo, methodCancelledMessageFormat, methodTimedOutMessageFormat);
+            ? RunWithTimeoutAndCancellationWithSTAThread(action, cancellationTokenSource, timeout, methodInfo, executionContextScope, methodCancelledMessageFormat, methodTimedOutMessageFormat)
+            : RunWithTimeoutAndCancellationWithThreadPool(action, cancellationTokenSource, timeout, methodInfo, executionContextScope, methodCancelledMessageFormat, methodTimedOutMessageFormat);
     }
 
     private static TestFailedException? RunWithTimeoutAndCancellationWithThreadPool(
         Action action, CancellationTokenSource cancellationTokenSource, int? timeout, MethodInfo methodInfo,
-        string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
+        IExecutionContextScope executionContextScope, string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
     {
         Exception? realException = null;
         Task? executionTask = null;
@@ -42,7 +45,7 @@ internal static class MethodRunner
                 {
                     try
                     {
-                        action();
+                        ExecutionContextService.RunActionOnContext(action, executionContextScope);
                     }
                     catch (Exception ex)
                     {
@@ -109,14 +112,14 @@ internal static class MethodRunner
 #endif
     private static TestFailedException? RunWithTimeoutAndCancellationWithSTAThread(
         Action action, CancellationTokenSource cancellationTokenSource, int? timeout, MethodInfo methodInfo,
-        string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
+        IExecutionContextScope executionContextScope, string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
     {
         Exception? realException = null;
         Thread executionThread = new(new ThreadStart(() =>
         {
             try
             {
-                action();
+                ExecutionContextService.RunActionOnContext(action, executionContextScope);
             }
             catch (Exception ex)
             {

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

@@ -0,0 +1,14 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices;
+
+internal readonly struct AssemblyExecutionContextScope : IExecutionContextScope
+{
+    public AssemblyExecutionContextScope(bool isCleanup)
+    {
+        IsCleanup = isCleanup;
+    }
+
+    public bool IsCleanup { get; }
+}

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

@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices;
+
+internal readonly struct ClassExecutionContextScope : IExecutionContextScope
+{
+    public ClassExecutionContextScope(Type type)
+    {
+        Type = type;
+        IsCleanup = false;
+        RemainingCleanupCount = 0;
+    }
+
+    public ClassExecutionContextScope(Type type, int remainingCleanupCount)
+    {
+        Type = type;
+        IsCleanup = true;
+        RemainingCleanupCount = remainingCleanupCount;
+    }
+
+    public Type Type { get; }
+
+    public bool IsCleanup { get; }
+
+    public int RemainingCleanupCount { get; }
+
+    public override readonly int GetHashCode() => Type.GetHashCode();
+
+    public override readonly bool Equals(object? obj) => Type.Equals(obj);
+}