Async TestInitialize and TestCleanup for WinUI throws

[Marv51]

Describe the bug

Async TestInitialize and TestCleanup work for plain tests, but for WinUI tests, test execution does not get invoked on the right thread.

Version used

    <PackageReference Include="Microsoft.TestPlatform.TestHost" Version="4.2.2" />
    <PackageReference Include="MSTest.TestAdapter" Version="4.2.2" />
    <PackageReference Include="MSTest.TestFramework" Version="4.2.2" />

Steps To Reproduce

[TestClass]
public class QuickTest
{
    [TestInitialize]
    public Task Setup()
    {
        return Task.Delay(10);
    }

    [UITestMethod]
    public void TestMethod2()
    {
        var grid = new Grid();
        Assert.AreEqual(0, grid.ActualWidth);
    }
}

Expected behavior

Test passes

Actual behavior

Exception: COMException: The application called an interface that was marshalled for a different thread. (0x8001010E)

    ExceptionHelpers.<ThrowExceptionForHR>g__Throw|41_0(Int32 hr)
    ExceptionHelpers.ThrowExceptionForHR(Int32 hr)
    _IGridFactoryMethods.CreateInstance(IObjectReference _obj, Object baseInterface, IntPtr& innerInterface)
    Grid.ctor()
    QuickTest.TestMethod2()
    MethodBaseInvoker.InterpretedInvoke_Method(Object obj, IntPtr* args)
    MethodBaseInvoker.InvokeWithNoArgs(Object obj, BindingFlags invokeAttr)

Additional context

I setup the WinUI project from the template and everything else works as expected.

The use case for me is that I am aiming to render some WinUI components during the test, and want to clean up the test window in TestInitialize and TestCleanup.