From 756e6da77cb04de6d070bd5dc9d2bd9e7cda071a Mon Sep 17 00:00:00 2001
From: PaGrom <pagrom23@gmail.com>
Date: Mon, 7 Sep 2020 02:21:20 +0300
Subject: [PATCH] Add CancellationToken to IfState condition

---
 Telegrom.StateMachine/Builder/GeneratedState.cs | 10 +++++-----
 Telegrom.StateMachine/Builder/IfState.cs        | 11 ++++++-----
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/Telegrom.StateMachine/Builder/GeneratedState.cs b/Telegrom.StateMachine/Builder/GeneratedState.cs
index fe7cd60..d906084 100644
--- a/Telegrom.StateMachine/Builder/GeneratedState.cs
+++ b/Telegrom.StateMachine/Builder/GeneratedState.cs
@@ -34,7 +34,7 @@ public async Task OnEnter(CancellationToken cancellationToken)
 
             await SaveOutputStateAttributesAsync(cancellationToken);
 
-            await MoveNextAsync(NextStateKind.AfterOnEnter);
+            await MoveNextAsync(NextStateKind.AfterOnEnter, cancellationToken);
         }
 
         public async Task Handle(CancellationToken cancellationToken)
@@ -43,7 +43,7 @@ public async Task Handle(CancellationToken cancellationToken)
 
             await SaveOutputStateAttributesAsync(cancellationToken);
 
-            await MoveNextAsync(NextStateKind.AfterHandle);
+            await MoveNextAsync(NextStateKind.AfterHandle, cancellationToken);
         }
 
         public async Task OnExit(CancellationToken cancellationToken)
@@ -52,10 +52,10 @@ public async Task OnExit(CancellationToken cancellationToken)
 
             await SaveOutputStateAttributesAsync(cancellationToken);
 
-            await MoveNextAsync(NextStateKind.AfterOnExit);
+            await MoveNextAsync(NextStateKind.AfterOnExit, cancellationToken);
         }
 
-        private async Task MoveNextAsync(NextStateKind nextStateKind)
+        private async Task MoveNextAsync(NextStateKind nextStateKind, CancellationToken cancellationToken)
         {
             if (_stateNode.NextStateKind != nextStateKind)
             {
@@ -66,7 +66,7 @@ private async Task MoveNextAsync(NextStateKind nextStateKind)
 
             foreach (var ifState in _stateNode.IfStates)
             {
-                if (await ifState.Condition(_stateContext))
+                if (await ifState.Condition(_stateContext, cancellationToken))
                 {
                     _stateContext.StateMachineContext.MoveTo(ifState.StateNode.StateName);
                     break;
diff --git a/Telegrom.StateMachine/Builder/IfState.cs b/Telegrom.StateMachine/Builder/IfState.cs
index 69c7cc9..6dfaaf2 100644
--- a/Telegrom.StateMachine/Builder/IfState.cs
+++ b/Telegrom.StateMachine/Builder/IfState.cs
@@ -1,14 +1,15 @@
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Telegrom.StateMachine.Builder
 {
     public sealed class IfState
     {
-        internal Func<IStateContext, Task<bool>> Condition { get; }
+        internal Func<IStateContext, CancellationToken, Task<bool>> Condition { get; }
         internal StateNode StateNode { get; }
 
-        public IfState(Func<IStateContext, Task<bool>> condition, Type stateType, string stateName = null)
+        public IfState(Func<IStateContext, CancellationToken, Task<bool>> condition, Type stateType, string stateName = null)
         {
             if (!typeof(IState).IsAssignableFrom(stateType))
             {
@@ -25,17 +26,17 @@ public IfState(Func<IStateContext, Task<bool>> condition, Type stateType, string
         }
 
         public IfState(Func<IStateContext, bool> condition, Type stateType, string stateName = null)
-            : this(context => new Task<bool>(() => condition(context)), stateType, stateName)
+            : this((context, ctk) => Task.Run(() => condition(context), ctk), stateType, stateName)
         { }
 
-        public IfState(Func<IStateContext, Task<bool>> condition, StateNode stateNode)
+        public IfState(Func<IStateContext, CancellationToken, Task<bool>> condition, StateNode stateNode)
         {
             Condition = condition;
             StateNode = stateNode;
         }
 
         public IfState(Func<IStateContext, bool> condition, StateNode stateNode)
-            : this(context => new Task<bool>(() => condition(context)), stateNode)
+            : this((context, ctk) => Task.Run(() => condition(context), ctk), stateNode)
         { }
     }
 }