Improve ExecutionContext propagation

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

@@ -68,9 +68,10 @@ internal static void ForceCleanup(TypeCache typeCache, IDictionary<string, objec
         using var writer = new ThreadSafeStringWriter(CultureInfo.InvariantCulture, "context");
         TestContext testContext = new TestContextImplementation(null, writer, sourceLevelParameters, logger);
         IEnumerable<TestClassInfo> classInfoCache = typeCache.ClassInfoListWithExecutableCleanupMethods;
+        LogMessageListener? listener = null;
         foreach (TestClassInfo classInfo in classInfoCache)
         {
-            TestFailedException? ex = classInfo.ExecuteClassCleanup(testContext);
+            TestFailedException? ex = classInfo.ExecuteClassCleanup(testContext, out listener);
             if (ex is not null)
             {
                 throw ex;
@@ -80,7 +81,7 @@ internal static void ForceCleanup(TypeCache typeCache, IDictionary<string, objec
         IEnumerable<TestAssemblyInfo> assemblyInfoCache = typeCache.AssemblyInfoListWithExecutableCleanupMethods;
         foreach (TestAssemblyInfo assemblyInfo in assemblyInfoCache)
         {
-            TestFailedException? ex = assemblyInfo.ExecuteAssemblyCleanup(testContext);
+            TestFailedException? ex = assemblyInfo.ExecuteAssemblyCleanup(testContext, ref listener);
             if (ex is not null)
             {
                 throw ex;

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

@@ -4,7 +4,6 @@
 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.TestTools.UnitTesting.UnitTestOutcome;
@@ -107,14 +106,21 @@ internal set
     /// </summary>
     internal Assembly Assembly { get; }
 
+    internal ExecutionContext? ExecutionContext { get; set; }
+
     /// <summary>
     /// Runs assembly initialize method.
     /// </summary>
     /// <param name="testContext"> The test context. </param>
     /// <exception cref="TestFailedException"> Throws a test failed exception if the initialization method throws an exception. </exception>
-    [SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "Requirement is to handle all kinds of user exceptions and message appropriately.")]
     public void RunAssemblyInitialize(TestContext testContext)
+        => RunAssemblyInitialize(testContext, out _);
+
+    [SuppressMessage("Microsoft.Design", "CA1031:DoNotCatchGeneralExceptionTypes", Justification = "Requirement is to handle all kinds of user exceptions and message appropriately.")]
+    internal void RunAssemblyInitialize(TestContext testContext, out LogMessageListener? logListener)
     {
+        logListener = null;
+
         // No assembly initialize => nothing to do.
         if (AssemblyInitializeMethod == null)
         {
@@ -140,14 +146,23 @@ public void RunAssemblyInitialize(TestContext testContext)
                 // Perform a check again.
                 if (!IsAssemblyInitializeExecuted)
                 {
+                    LogMessageListener? logMessageListener = null;
+
                     try
                     {
                         AssemblyInitializationException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
-                            () => AssemblyInitializeMethod.InvokeAsSynchronousTask(null, testContext),
+                            () =>
+                            {
+                                logMessageListener = new LogMessageListener(MSTestSettings.CurrentSettings.CaptureDebugTraces);
+                                AssemblyInitializeMethod.InvokeAsSynchronousTask(null, testContext);
+                                // **After** we have executed the assembly initialize, we save the current context.
+                                // This context will contain async locals set by the assembly initialize method.
+                                ExecutionContext = ExecutionContext.Capture();
+                            },
                             testContext.CancellationTokenSource,
                             AssemblyInitializeMethodTimeoutMilliseconds,
                             AssemblyInitializeMethod,
-                            new AssemblyExecutionContextScope(isCleanup: false),
+                            executionContext: ExecutionContext,
                             Resource.AssemblyInitializeWasCancelled,
                             Resource.AssemblyInitializeTimedOut);
                     }
@@ -158,6 +173,7 @@ public void RunAssemblyInitialize(TestContext testContext)
                     finally
                     {
                         IsAssemblyInitializeExecuted = true;
+                        logListener = logMessageListener;
                     }
                 }
             }
@@ -216,11 +232,15 @@ public void RunAssemblyInitialize(TestContext testContext)
             try
             {
                 AssemblyCleanupException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
-                     () => AssemblyCleanupMethod.InvokeAsSynchronousTask(null),
+                     () =>
+                     {
+                         AssemblyCleanupMethod.InvokeAsSynchronousTask(null);
+                         ExecutionContext = ExecutionContext.Capture();
+                     },
                      new CancellationTokenSource(),
                      AssemblyCleanupMethodTimeoutMilliseconds,
                      AssemblyCleanupMethod,
-                     new AssemblyExecutionContextScope(isCleanup: true),
+                     ExecutionContext,
                      Resource.AssemblyCleanupWasCancelled,
                      Resource.AssemblyCleanupTimedOut);
             }
@@ -260,7 +280,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 TestFailedException? ExecuteAssemblyCleanup(TestContext testContext)
+    internal TestFailedException? ExecuteAssemblyCleanup(TestContext testContext, ref LogMessageListener? logListener)
     {
         if (AssemblyCleanupMethod == null)
         {
@@ -269,11 +289,13 @@ public void RunAssemblyInitialize(TestContext testContext)
 
         lock (_assemblyInfoExecuteSyncObject)
         {
+            LogMessageListener? logMessageListener = logListener;
             try
             {
                 AssemblyCleanupException = FixtureMethodRunner.RunWithTimeoutAndCancellation(
                      () =>
                      {
+                         logMessageListener ??= new LogMessageListener(MSTestSettings.CurrentSettings.CaptureDebugTraces);
                          if (AssemblyCleanupMethod.GetParameters().Length == 0)
                          {
                              AssemblyCleanupMethod.InvokeAsSynchronousTask(null);
@@ -282,18 +304,24 @@ public void RunAssemblyInitialize(TestContext testContext)
                          {
                              AssemblyCleanupMethod.InvokeAsSynchronousTask(null, testContext);
                          }
+
+                         ExecutionContext = ExecutionContext.Capture();
                      },
                      testContext.CancellationTokenSource,
                      AssemblyCleanupMethodTimeoutMilliseconds,
                      AssemblyCleanupMethod,
-                     new AssemblyExecutionContextScope(isCleanup: true),
+                     ExecutionContext,
                      Resource.AssemblyCleanupWasCancelled,
                      Resource.AssemblyCleanupTimedOut);
             }
             catch (Exception ex)
             {
                 AssemblyCleanupException = ex;
             }
+            finally
+            {
+                logListener = logMessageListener;
+            }
         }
 
         // If assemblyCleanup was successful, then don't do anything