From 579068d0b13446fc1ea93149b06d712c3546cad9 Mon Sep 17 00:00:00 2001
From: Sam Harwell <Sam.Harwell@microsoft.com>
Date: Tue, 4 Feb 2020 11:24:21 -0800
Subject: [PATCH] Fix incorrect SynchronizationContext use in TestSweeper

---
 test/Microsoft.ML.Sweeper.Tests/TestSweeper.cs | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/test/Microsoft.ML.Sweeper.Tests/TestSweeper.cs b/test/Microsoft.ML.Sweeper.Tests/TestSweeper.cs
index ec54cd116c..c1b20cd6bb 100644
--- a/test/Microsoft.ML.Sweeper.Tests/TestSweeper.cs
+++ b/test/Microsoft.ML.Sweeper.Tests/TestSweeper.cs
@@ -175,7 +175,7 @@ public void TestSimpleSweeperAsync()
         }
 
         [Fact]
-        public void TestDeterministicSweeperAsyncCancellation()
+        public async Task TestDeterministicSweeperAsyncCancellation()
         {
             var random = new Random(42);
             var env = new MLContext(42);
@@ -215,7 +215,7 @@ public void TestDeterministicSweeperAsyncCancellation()
             // Cancel after the first barrier and check if the number of registered actions
             // is indeed 2 * batchSize.
             sweeper.Cancel();
-            Task.WaitAll(tasks.ToArray());
+            await Task.WhenAll(tasks);
             foreach (var task in tasks)
             {
                 if (task.Result != null)
@@ -225,7 +225,7 @@ public void TestDeterministicSweeperAsyncCancellation()
         }
 
         [Fact]
-        public void TestDeterministicSweeperAsync()
+        public async Task TestDeterministicSweeperAsync()
         {
             var random = new Random(42);
             var env = new MLContext(42);
@@ -289,8 +289,7 @@ public void TestDeterministicSweeperAsync()
             foreach (var run in results)
                 sweeper.Update(run.Key, run.Value);
 
-            Task.WaitAll(tasks);
-            tasks.All(t => t.IsCompleted);
+            await Task.WhenAll(tasks);
         }
 
         [Fact]
@@ -327,7 +326,7 @@ public void TestDeterministicSweeperAsyncParallel()
             int[] sleeps = new int[sweeps];
             for (int i = 0; i < sleeps.Length; i++)
                 sleeps[i] = random.Next(10, 100);
-            var r = Parallel.For(0, sweeps, options, (int i) =>
+            var r = Task.Run(() => Parallel.For(0, sweeps, options, (int i) =>
             {
                 var task = sweeper.Propose();
                 task.Wait();
@@ -340,7 +339,7 @@ public void TestDeterministicSweeperAsyncParallel()
                 sweeper.Update(paramWithId.Id, result);
                 lock (mlock)
                     paramSets.Add(paramWithId.ParameterSet);
-            });
+            }));
             Assert.True(paramSets.Count <= sweeps);
             CheckAsyncSweeperResult(paramSets);
         }