Regression: Custom SynchronizationContext set in TestInitialize is removed

[jbe2277]

Describe the bug

It's common to use a custom SynchronizationContext (e.g. UnitTestSynchronizationContext) for unit testing of client applications (e.g. MAUI, WPF, WinForms). This allows to simulate the same async behavior in the unit test environment as we have in the application.

We initialize the custom SynchronizationContext in a TestInitialize method and dispose it in TestCleanup. This worked fine for more than 10 years (even with legacy MSTest v1).

Starting with MSTest v4 this approach does not work anymore. The SynchronizationContext.Current is now null in the test method.

• Last working version: 3.9.3
• Not working version: 4.02

Steps To Reproduce

1. Clone repo MSTestSynchronizationContextIssue
2. It contains 3 projects:
   • ❌ UnitTest.MSTest: Uses MSTest 4.0.2. This one fails because the SynchonizationContext is set to null before calling the test method.
   • ✅ UnitTestOld.MSTest: Uses MSTest 3.9.3. This one passes as it works as expected.
   • ✅ UnitTest.XUnit: Uses the latest Xunit v3 with the same approach. This one passes as it works as expected.

[Youssef1313]

This is most likely because we moved away from blocking threads using GetAwaiter().GetResult() in multiple places, to be fully async instead. In that case we are controlled by AsyncTaskMethodBuilder. I would guess the previous synchronization context gets restored here:

https://github.com/dotnet/runtime/blob/8fcb70a35c893ad7620e0d6ac4f7ea812af0dca6/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilderCore.cs#L42-L47

There are three potential fixes I can think of:

• Detect that the fixture we are calling into isn't awaitable beforehand, and go into a fully synchronous code path instead of calling into an async method that will complete synchronously anyways. Avoiding the async method will avoid the state machine completely and will make sure the synchronization context is preserved.

  • Extremely niche scenario: what if the user defines an async TestInitialize with a custom async method builder (via AsyncMethodBuilderAttribute) that doesn't reset the sync context? In that case, our own async calls will reset it.

• Manually capture the SynchronizationContext exactly after we call the fixture, and restore it back just before we call the test method.

  • Any possible memory leaks to be aware of that needs special handling?

• Use a custom async method builder ourselves in our async methods for executing fixtures. Something like the following:

  public struct SynchronizationContextPreservingAsyncTaskMethodBuilder
  {
      private AsyncTaskMethodBuilder _inner;
  
      public static SynchronizationContextPreservingAsyncTaskMethodBuilder Create()
          => new() { _inner = AsyncTaskMethodBuilder.Create() };
  
      public Task Task
          => _inner.Task;
  
      public void SetResult()
          => _inner.SetResult();
  
      public void SetException(Exception ex)
          => _inner.SetException(ex);
  
      public void SetStateMachine(IAsyncStateMachine stateMachine)
          => _inner.SetStateMachine(stateMachine);
  
      public void Start<TStateMachine>(ref TStateMachine stateMachine)
          where TStateMachine : IAsyncStateMachine
      {
          // Start is the whole reason why we have this custom builder.
          // BCL implementation restores back SynchronizationContext.
          // See https://github.com/dotnet/runtime/blob/c591f971241e7074f8a31ccde744aec9794e2500/src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilderCore.cs#L45-L46
          // We want to avoid restoring the SynchronizationContext for the case when the task completes synchronously on the same thread.
          // This allows TestInitialize to set SynchronizationContext, and lets us be still fully async in our implementation.
          // But then TestMethod can see the correct SynchronizationContext in the case of TestInitialize completing synchronously.
          stateMachine.MoveNext();
      }
  
      public void AwaitOnCompleted<TAwaiter, TStateMachine>(
          ref TAwaiter awaiter, ref TStateMachine stateMachine)
          where TAwaiter : INotifyCompletion
          where TStateMachine : IAsyncStateMachine
      {
          _inner.AwaitOnCompleted(ref awaiter, ref stateMachine);
      }
  
      public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
          ref TAwaiter awaiter, ref TStateMachine stateMachine)
          where TAwaiter : ICriticalNotifyCompletion
          where TStateMachine : IAsyncStateMachine
      {
          _inner.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine);
      }
  }

  • Is it safe?
  • Is it going to work correctly on all .NET versions, including .NET Framework?