Fix test runner crash

microsoft · May 7, 2024 · fb9ed53 · fb9ed53
1 parent aaf2c14
commit fb9ed53
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/test/StreamJsonRpc.Tests/JsonRpcTests.cs b/test/StreamJsonRpc.Tests/JsonRpcTests.cs
@@ -2969,6 +2969,11 @@ public void JoinableTaskFactory_IntegrationClientSideOnly_ManyConnections_Unique
         // Set up a main thread and JoinableTaskContext.
         JoinableTaskContext jtc = new();
 
+        // Track our async work so our test doesn't exit before our UI thread requests do,
+        // or the test process will crash.
+        JoinableTaskCollection jtCollection = jtc.CreateCollection();
+        JoinableTaskFactory jtf = jtc.CreateFactory(jtCollection);
+
         // Configure the client (only) to understand JTF.
         this.clientRpc.AllowModificationWhileListening = true;
         this.clientRpc.JoinableTaskFactory = jtc.Factory;
@@ -2997,7 +3002,7 @@ public void JoinableTaskFactory_IntegrationClientSideOnly_ManyConnections_Unique
         const string CallbackMethodName = "ClientNeedsMainThread";
         alternateClientRpc.AddLocalRpcMethod(CallbackMethodName, new Func<Task>(async delegate
         {
-            await jtc.Factory.SwitchToMainThreadAsync();
+            await jtf.SwitchToMainThreadAsync();
         }));
 
         alternateServerRpc.StartListening();
@@ -3007,6 +3012,9 @@ public void JoinableTaskFactory_IntegrationClientSideOnly_ManyConnections_Unique
         {
             await Assert.ThrowsAsync<OperationCanceledException>(() => this.clientRpc.InvokeWithCancellationAsync(nameof(this.server.CallbackOnAnotherConnection), new object?[] { CallbackMethodName }, this.TimeoutToken).WithCancellation(ExpectedTimeoutToken));
         });
+
+        // Drain any UI thread requests before exiting the test.
+        jtc.Factory.Run(jtCollection.JoinTillEmptyAsync);
     }
 
     [Fact]