Ensure that fixtures also support sta threading (#2769)

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

@@ -3,6 +3,7 @@
 
 using System.Globalization;
 using System.Reflection;
+using System.Runtime.InteropServices;
 
 using Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.ObjectModel;
 
@@ -15,6 +16,13 @@ internal static class MethodRunner
     internal static TestFailedException? RunWithTimeoutAndCancellation(
         Action action, CancellationTokenSource cancellationTokenSource, int? timeout, MethodInfo methodInfo,
         string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
+        => RuntimeInformation.IsOSPlatform(OSPlatform.Windows) && Thread.CurrentThread.GetApartmentState() == ApartmentState.STA
+            ? RunWithTimeoutAndCancellationWithSTAThread(action, cancellationTokenSource, timeout, methodInfo, methodCancelledMessageFormat, methodTimedOutMessageFormat)
+            : RunWithTimeoutAndCancellationWithThreadPool(action, cancellationTokenSource, timeout, methodInfo, methodCancelledMessageFormat, methodTimedOutMessageFormat);
+
+    private static TestFailedException? RunWithTimeoutAndCancellationWithThreadPool(
+        Action action, CancellationTokenSource cancellationTokenSource, int? timeout, MethodInfo methodInfo,
+        string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
     {
         Exception? realException = null;
         Task? executionTask = null;
@@ -83,6 +91,91 @@ internal static class MethodRunner
                 throw realException;
             }
 
+            throw;
+        }
+    }
+
+#if NET6_0_OR_GREATER
+    [System.Runtime.Versioning.SupportedOSPlatform("windows")]
+#endif
+    private static TestFailedException? RunWithTimeoutAndCancellationWithSTAThread(
+        Action action, CancellationTokenSource cancellationTokenSource, int? timeout, MethodInfo methodInfo,
+        string methodCancelledMessageFormat, string methodTimedOutMessageFormat)
+    {
+        Exception? realException = null;
+        Thread executionThread = new(new ThreadStart(() =>
+        {
+            try
+            {
+                action();
+            }
+            catch (Exception ex)
+            {
+                realException = ex;
+                throw;
+            }
+        }))
+        {
+            IsBackground = true,
+            Name = "MSTest Fixture Execution Thread",
+        };
+
+        executionThread.SetApartmentState(ApartmentState.STA);
+        executionThread.Start();
+
+        try
+        {
+            // Creates a Task<bool> that represents the results of the execution thread.
+            var executionTask = Task.Run(
+                () =>
+                {
+                    if (timeout.HasValue)
+                    {
+                        return executionThread.Join(timeout.Value);
+                    }
+
+                    executionThread.Join();
+                    return true;
+                },
+                cancellationTokenSource.Token);
+            executionTask.Wait(cancellationTokenSource.Token);
+
+            // If the execution thread completes before the timeout, the task will return true, otherwise false.
+            if (executionTask.Result)
+            {
+                return realException is not null
+                    ? throw realException
+                    : null;
+            }
+
+            // Timed out. For cancellation, either OCE or AggregateException with TCE will be thrown.
+            return new(
+                UnitTestOutcome.Timeout,
+                string.Format(
+                    CultureInfo.InvariantCulture,
+                    methodTimedOutMessageFormat,
+                    methodInfo.DeclaringType!.FullName,
+                    methodInfo.Name));
+        }
+        catch (Exception ex) when
+
+            // This exception occurs when the cancellation happens before the task is actually started.
+            ((ex is TaskCanceledException tce && tce.CancellationToken == cancellationTokenSource.Token)
+            || (ex is OperationCanceledException oce && oce.CancellationToken == cancellationTokenSource.Token)
+            || (ex is AggregateException aggregateEx && aggregateEx.InnerExceptions.OfType<TaskCanceledException>().Any()))
+        {
+            return new(
+                UnitTestOutcome.Timeout,
+                string.Format(CultureInfo.InvariantCulture, methodCancelledMessageFormat, methodInfo.DeclaringType!.FullName, methodInfo.Name));
+        }
+        catch (Exception)
+        {
+            // We throw the real exception to have the original stack trace to elaborate up the chain.
+            if (realException is not null)
+            {
+                throw realException;
+            }
+
             throw;
         }
     }