Become with arguments

OrleansContrib · Dec 19, 2019 · e114ad7 · e114ad7
1 parent 77b26da
commit e114ad7
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 22 deletions.
diff --git a/Source/Orleankka.Runtime/Behaviors/Behavior.cs b/Source/Orleankka.Runtime/Behaviors/Behavior.cs
@@ -82,52 +82,68 @@ public Task<object> Receive(object message)
             return Current.Receive(message);
         }
 
-        public Task Become(Receive behavior)
+        public Task Become(Receive behavior) => HandleBecome(behavior, Behaviors.Become.Message);
+        public Task Become<TArg>(Receive behavior, TArg arg) => HandleBecome(behavior, new Become<TArg>(arg));
+
+        Task HandleBecome(Receive behavior, Become message)
         {
             Requires.NotNull(behavior, nameof(behavior));
-            var name = behavior.Method.Name;
+
+            var state = states.Find(behavior.Method.Name) 
+                        ?? new State(behavior.Method.Name, behavior);
 
-            var configured = states.Find(name);
-            return Become(configured ?? new State(name, behavior));
+            return HandleBecome(state, message);
         }
 
-        public Task Become(string behavior)
+        public Task Become(string behavior) => HandleBecome(behavior, Behaviors.Become.Message);
+        public Task Become<TArg>(string behavior, TArg arg) => HandleBecome(behavior, new Become<TArg>(arg));
+
+        Task HandleBecome(string behavior, Become message)
         {
             Requires.NotNull(behavior, nameof(behavior));
 
-            return Become(State(behavior));
+            return HandleBecome(State(behavior), message);
         }
 
-        public Task BecomeStacked(Receive behavior)
+        public Task BecomeStacked(Receive behavior) => HandleBecomeStacked(behavior, Behaviors.Become.Message);
+        public Task BecomeStacked<TArg>(Receive behavior, TArg arg) => HandleBecomeStacked(behavior, new Become<TArg>(arg));
+
+        Task HandleBecomeStacked(Receive behavior, Become message)
         {
             Requires.NotNull(behavior, nameof(behavior));
-            var name = behavior.Method.Name;
+
+            var state = states.Find(behavior.Method.Name) 
+                        ?? new State(behavior.Method.Name, behavior);
 
-            var configured = states.Find(name);
-            return BecomeStacked(configured ?? new State(name, behavior));
+            return HandleBecomeStacked(state, message);
         }
 
-        public Task BecomeStacked(string behavior)
+        public Task BecomeStacked(string behavior) => HandleBecomeStacked(behavior, Behaviors.Become.Message);
+        public Task BecomeStacked<TArg>(string behavior, TArg arg) => HandleBecomeStacked(behavior, new Become<TArg>(arg));
+
+        Task HandleBecomeStacked(string behavior, Become message)
         {
             Requires.NotNull(behavior, nameof(behavior));
-            return BecomeStacked(State(behavior));
+
+            return HandleBecomeStacked(State(behavior), message);
         }
 
-        async Task BecomeStacked(State next)
+        async Task HandleBecomeStacked(State next, Become message)
         {
             history.Push(Current);
-            await Become(next);
+
+            await HandleBecome(next, message);
         }
 
         public async Task Unbecome()
         {
             if (history.Count == 0)
                 throw new InvalidOperationException("The previous behavior has not been recorded. Use BecomeStacked method to stack behaviors");
 
-            await Become(history.Pop());
+            await HandleBecome(history.Pop(), Behaviors.Become.Message);
         }
 
-        async Task Become(State next)
+        async Task HandleBecome(State next, Become message)
         {
             if (!Initialized())
                 throw new InvalidOperationException("Initial behavior should be set before calling Become");
@@ -146,7 +162,7 @@ async Task Become(State next)
             Previous = Current;
             Current = next;
 
-            await next.HandleBecome(transition);
+            await next.HandleBecome(transition, message);
             await next.HandleActivate(transition);
 
             transition = null;

diff --git a/Source/Orleankka.Runtime/Behaviors/BehaviorMessages.cs b/Source/Orleankka.Runtime/Behaviors/BehaviorMessages.cs
@@ -3,12 +3,21 @@
     public interface BehaviorMessage
     {}
 
-    public sealed class Become : BehaviorMessage, LifecycleMessage
+    public class Become : BehaviorMessage, LifecycleMessage
     {
-        Become(){}
+        protected Become(){}
+
         public static readonly Become Message = new Become();
     }
 
+    public sealed class Become<TArg> : Become
+    {
+        public readonly TArg Argument;
+
+        internal Become(TArg argument) => 
+            Argument = argument;
+    }
+
     public sealed class Unbecome : BehaviorMessage, LifecycleMessage
     {
         Unbecome(){}

diff --git a/Source/Orleankka.Runtime/Behaviors/State.cs b/Source/Orleankka.Runtime/Behaviors/State.cs
@@ -47,15 +47,15 @@ internal async Task<object> Receive(object message)
             return result;
         }
 
-        internal async Task HandleBecome(Transition transition)
+        internal async Task HandleBecome(Transition transition, Become message)
         {
             if (IsSuperOf(transition.From))
                 return;
 
             if (Super != null)
-                await Super.HandleBecome(transition);
+                await Super.HandleBecome(transition, message);
 
-            await CallBehavior(Become.Message);
+            await CallBehavior(message);
         }
 
         internal async Task HandleUnbecome(Transition transition)

diff --git a/Tests/Orleankka.Tests/Features/Actor_behaviors/Switchable_behaviors.cs b/Tests/Orleankka.Tests/Features/Actor_behaviors/Switchable_behaviors.cs
@@ -231,6 +231,31 @@ async Task<object> AttemptBecomeDuring<T>(Behavior b, string other, object messa
 
                 Assert.ThrowsAsync<InvalidOperationException>(async () => await behavior.Become("B"));
             }
+
+            [Test]
+            public async Task When_become_with_arguments()
+            {
+                string passedArg = null;
+
+                Task<object> Receive(object message)
+                {
+                    if (message is Become<string> m)
+                        passedArg = m.Argument;
+
+                    return TaskResult.Done;
+                };
+
+                Behavior behavior = new BehaviorTester(events)
+                    .State("A", Receive)
+                    .State("B", Receive)
+                    .Initial("A");
+
+                await behavior.Become("B", "arg1");
+                Assert.That(passedArg, Is.EqualTo("arg1"));
+
+                await behavior.BecomeStacked("A", "arg2");
+                Assert.That(passedArg, Is.EqualTo("arg2"));
+            }
         }
     }
 }