Port Akka.Tests.Actor.Stash tests to async/await - ActorWithStashSpec

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

@@ -12,6 +12,7 @@
 using Akka.TestKit;
 using Akka.TestKit.TestActors;
 using Akka.Tests.TestUtils;
+using Akka.Tests.Util;
 using Xunit;
 
 namespace Akka.Tests.Actor.Stash
@@ -89,7 +90,7 @@ await EventFilter.DeadLetter<string>(s=>s=="message", source: stasher.Path.ToStr
         }
 
         [Fact]
-        public void An_actor_Must_process_stashed_messages_after_restart()
+        public async Task An_actor_Must_process_stashed_messages_after_restart()
         {
             SupervisorStrategy strategy = new OneForOneStrategy(2, TimeSpan.FromSeconds(1), e => Directive.Restart);
             var boss = ActorOf(() => new Supervisor(strategy));
@@ -98,7 +99,7 @@ public void An_actor_Must_process_stashed_messages_after_restart()
             var slaveProps = Props.Create(() => new SlaveActor(restartLatch, hasMsgLatch, "stashme"));
 
             //Send the props to supervisor, which will create an actor and return the ActorRef
-            var slave = boss.AskAndWait<IActorRef>(slaveProps, TestKitSettings.DefaultTimeout);
+            var slave = await boss.Ask<IActorRef>(slaveProps).WithTimeout(TestKitSettings.DefaultTimeout);
 
             //send a message that will be stashed
             slave.Tell("stashme");
@@ -121,7 +122,7 @@ public async Task An_actor_that_clears_the_stash_on_preRestart_Must_not_receive_
             var slaveProps = Props.Create(() => new ActorsThatClearsStashOnPreRestart(restartLatch));
 
             //Send the props to supervisor, which will create an actor and return the ActorRef
-            var slave = boss.AskAndWait<IActorRef>(slaveProps, TestKitSettings.DefaultTimeout);
+            var slave = await boss.Ask<IActorRef>(slaveProps).WithTimeout(TestKitSettings.DefaultTimeout);;
 
             //send messages that will be stashed
             slave.Tell("stashme 1");

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

@@ -30,5 +30,33 @@ public static async Task<bool> AwaitWithTimeout(this Task parentTask, TimeSpan t
                 }
             }
         }
+
+        public static async Task<T> WithTimeout<T>(this Task<T> parentTask, TimeSpan timeout)
+        {
+            using (var cts = new CancellationTokenSource())
+            {
+                try
+                {
+                    var delayed = Task.Delay(timeout, cts.Token);
+                    var returnedTask = await Task.WhenAny(delayed, parentTask);
+
+                    if (returnedTask != parentTask)
+                        throw new TaskCanceledException($"Task timed out after {timeout.TotalSeconds} seconds");
+
+                    if(returnedTask == parentTask && returnedTask.Exception != null)
+                    {
+                        var flattened = returnedTask.Exception.Flatten();
+                        ExceptionDispatchInfo.Capture(flattened.InnerException).Throw();
+                    }
+
+                    return parentTask.Result;
+                }
+                finally
+                {
+                    cts.Cancel();
+                }
+            }
+        }
+
     }
 }