Fixed awaiting task

zeromq · May 24, 2024 · c4b4ba2 · c4b4ba2
1 parent 4d3c831
commit c4b4ba2
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/src/NetMQ.Tests/NetMQPollerTest.cs b/src/NetMQ.Tests/NetMQPollerTest.cs
@@ -984,7 +984,7 @@ public void TwoThreads()
 
                 TaskUtils.Wait(t1, TimeSpan.FromMilliseconds(1000));
                 TaskUtils.Wait(t2, TimeSpan.FromMilliseconds(1000));
-                TaskUtils.WaitAll(allTasks.ToArray(), TimeSpan.FromMilliseconds(1000));
+                TaskUtils.WaitAll(allTasks, TimeSpan.FromMilliseconds(1000));
 
                 Assert.Equal(100, count1);
                 Assert.Equal(100, count2);

diff --git a/src/NetMQ.Tests/TaskUtils.cs b/src/NetMQ.Tests/TaskUtils.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -30,24 +31,26 @@ internal static async Task PollUntil(Func<bool> condition, CancellationToken ct
             }
         }
 
-        internal static bool WaitAll(Task[] tasks, TimeSpan timeout)
+        internal static bool WaitAll(IEnumerable<Task> tasks, TimeSpan timeout)
         {
-            return PollUntil(() => tasks.All(t => t.IsCompleted), timeout).Status == TaskStatus.RanToCompletion;
+            PollUntil(() => tasks.All(t => t.IsCompleted), timeout).Wait();
+            return tasks.All(t => t.Status == TaskStatus.RanToCompletion);
         }
 
-        internal static bool WaitAll(Task[] tasks)
+        internal static void WaitAll(IEnumerable<Task> tasks)
         {
-            return WaitAll(tasks, TimeSpan.MaxValue);
+            PollUntil(() => tasks.All(t => t.IsCompleted), Timeout.InfiniteTimeSpan).Wait();
         }
 
         internal static bool Wait(Task task, TimeSpan timeout)
         {
-            return PollUntil(() => task.IsCompleted, timeout).Status == TaskStatus.RanToCompletion;
+            PollUntil(() => task.IsCompleted, timeout).Wait();
+            return task.Status == TaskStatus.RanToCompletion;
         }
 
-        internal static bool Wait(Task task)
+        internal static void Wait(Task task)
         {
-            return Wait(task, TimeSpan.MaxValue);
+            PollUntil(() => task.IsCompleted, Timeout.InfiniteTimeSpan).Wait();
         }
     }
 }