Move Akka.Tests.Util.TaskHelper to Akka.TestKit.Extensions.TaskExtensions

src/core/Akka.Tests/Util/TaskHelpers.cs → src/core/Akka.TestKit/Extensions/TaskExtensions.cs

@@ -3,13 +3,13 @@
 using System.Threading;
 using System.Threading.Tasks;
 
-namespace Akka.Tests.Util
+namespace Akka.TestKit.Extensions
 {
-    public static class TaskHelpers
+    public static class TaskExtensions
     {
-        public static async Task<bool> AwaitWithTimeout(this Task parentTask, TimeSpan timeout)
+        public static async Task<bool> AwaitWithTimeout(this Task parentTask, TimeSpan timeout, CancellationToken cancellationToken = default)
         {
-            using (var cts = new CancellationTokenSource())
+            using (var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken))
             {
                 try
                 {
@@ -19,7 +19,10 @@ public static async Task<bool> AwaitWithTimeout(this Task parentTask, TimeSpan t
                     if(returnedTask == parentTask && returnedTask.Exception != null)
                     {
                         var flattened = returnedTask.Exception.Flatten();
-                        ExceptionDispatchInfo.Capture(flattened.InnerException).Throw();
+                        if(flattened.InnerExceptions.Count == 1)
+                            ExceptionDispatchInfo.Capture(flattened.InnerExceptions[0]).Throw();
+                        else
+                            ExceptionDispatchInfo.Capture(returnedTask.Exception).Throw();
                     }
 
                     return parentTask.IsCompleted;
@@ -31,9 +34,9 @@ public static async Task<bool> AwaitWithTimeout(this Task parentTask, TimeSpan t
             }
         }
 
-        public static async Task<T> WithTimeout<T>(this Task<T> parentTask, TimeSpan timeout)
+        public static async Task<T> WithTimeout<T>(this Task<T> parentTask, TimeSpan timeout, CancellationToken cancellationToken = default)
         {
-            using (var cts = new CancellationTokenSource())
+            using (var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken))
             {
                 try
                 {
@@ -46,7 +49,10 @@ public static async Task<T> WithTimeout<T>(this Task<T> parentTask, TimeSpan tim
                     if(returnedTask == parentTask && returnedTask.Exception != null)
                     {
                         var flattened = returnedTask.Exception.Flatten();
-                        ExceptionDispatchInfo.Capture(flattened.InnerException).Throw();
+                        if(flattened.InnerExceptions.Count == 1)
+                            ExceptionDispatchInfo.Capture(flattened.InnerExceptions[0]).Throw();
+                        else
+                            ExceptionDispatchInfo.Capture(returnedTask.Exception).Throw();
                     }
 
                     return parentTask.Result;

src/core/Akka.Tests/Actor/ActorSystemSpec.cs

@@ -20,6 +20,7 @@
 using Akka.Configuration;
 using Akka.Dispatch;
 using Akka.Event;
+using Akka.TestKit.Extensions;
 using FluentAssertions.Execution;
 using Akka.Tests.Util;
 

src/core/Akka.Tests/Actor/CoordinatedShutdownSpec.cs

@@ -14,6 +14,7 @@
 using System.Linq;
 using System.Threading.Tasks;
 using Akka.Configuration;
+using Akka.TestKit.Extensions;
 using FluentAssertions;
 using Xunit;
 using static Akka.Actor.CoordinatedShutdown;

src/core/Akka.Tests/Actor/DeathWatchSpec.cs

@@ -16,6 +16,7 @@
 using Akka.Dispatch.SysMsg;
 using Akka.Event;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Akka.Tests.TestUtils;
 using Akka.Tests.Util;
 using Akka.Util.Internal;

src/core/Akka.Tests/Actor/InboxSpec.cs

@@ -13,6 +13,7 @@
 using Akka.Actor.Internal;
 using Akka.Event;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Akka.Tests.Util;
 using Xunit;
 

src/core/Akka.Tests/Actor/LocalActorRefProviderSpec.cs

@@ -11,6 +11,7 @@
 using Akka.Actor;
 using Akka.Actor.Internal;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Xunit;
 using Akka.TestKit.TestActors;
 using Akka.Tests.Util;

src/core/Akka.Tests/Actor/PatternSpec.cs

@@ -10,6 +10,7 @@
 using Akka.Actor;
 using Akka.Event;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Akka.Tests.Util;
 using Xunit;
 

src/core/Akka.Tests/Actor/Scheduler/SchedulerShutdownSpec.cs

@@ -13,6 +13,7 @@
 using Akka.Dispatch;
 using Akka.Event;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Akka.Tests.Util;
 using Akka.Util.Internal;
 using FluentAssertions;

src/core/Akka.Tests/Actor/Stash/ActorWithStashSpec.cs

@@ -10,6 +10,7 @@
 using Akka.Actor;
 using Akka.Actor.Internal;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Akka.TestKit.TestActors;
 using Akka.Tests.TestUtils;
 using Akka.Tests.Util;

src/core/Akka.Tests/Pattern/CircuitBreakerSpec.cs

@@ -14,6 +14,7 @@
 using System.Threading.Tasks;
 using Akka.Pattern;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Akka.Tests.Util;
 using Xunit;
 

src/core/Akka.Tests/Util/IndexSpec.cs

@@ -10,6 +10,7 @@
 using System.Linq;
 using System.Threading.Tasks;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Akka.Util;
 using Xunit;
 using FluentAssertions;

src/core/Akka.Tests/Util/Internal/InterlockedSpinTests.cs

@@ -9,6 +9,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Akka.TestKit;
+using Akka.TestKit.Extensions;
 using Akka.Util.Internal;
 using FluentAssertions.Extensions;
 using Xunit;