Implement BeforeInject callback

Closes Polly-Contrib#65
iCodeSometime · Jan 20, 2022 · a7b5090 · a7b5090
1 parent a18b123
commit a7b5090
Show file tree

Hide file tree

Showing 14 changed files with 219 additions and 10 deletions.
diff --git a/src/Polly.Contrib.Simmy.Specs/Behavior/InjectBehaviourAsyncWithOptionsSpecs.cs b/src/Polly.Contrib.Simmy.Specs/Behavior/InjectBehaviourAsyncWithOptionsSpecs.cs
@@ -109,6 +109,59 @@ public void Should_inject_behaviour_before_executing_user_delegate()
             injectedBehaviourExecuted.Should().BeTrue();
         }
 
+        #region BeforeInject
+        [Fact]
+        public async Task Should_call_before_inject_callback_before_injecting_behavior()
+        {
+            var beforeInjectExecuted = false;
+            var injectedBehaviourExecuted = false;
+
+            var policy = MonkeyPolicy.InjectBehaviourAsync(with =>
+                with.Behaviour(async () =>
+                {
+                    beforeInjectExecuted.Should().BeTrue();
+                    injectedBehaviourExecuted = true;
+                })
+                .BeforeInject(async (context, cancellation) =>
+                {
+                    injectedBehaviourExecuted.Should().BeFalse();
+                    beforeInjectExecuted = true;
+                })
+                .InjectionRate(0.6)
+                .Enabled()
+            );
+
+            await policy.ExecuteAsync(() => Task.CompletedTask);
+
+            beforeInjectExecuted.Should().BeTrue();
+            injectedBehaviourExecuted.Should().BeTrue();
+        }
+
+        public async Task Should_not_call_before_inject_callback_if_not_injecting()
+        {
+            var beforeInjectExecuted = false;
+            var behaviorExecuted = false;
+
+            var policy = MonkeyPolicy.InjectBehaviourAsync(with =>
+                with.Behaviour(async () =>
+                {
+                    behaviorExecuted = true;
+                })
+                .BeforeInject(async (context, cancellation) =>
+                {
+                    beforeInjectExecuted = true;
+                })
+                .InjectionRate(0.4)
+                .Enabled()
+            );
+
+            await policy.ExecuteAsync(() => Task.CompletedTask);
+
+            beforeInjectExecuted.Should().BeFalse();
+            behaviorExecuted.Should().BeFalse();
+        }
+        #endregion
+
         #region invalid threshold on configuration and execution time
 
         [Fact]

diff --git a/src/Polly.Contrib.Simmy.Specs/Behavior/InjectBehaviourWithOptionsSpecs.cs b/src/Polly.Contrib.Simmy.Specs/Behavior/InjectBehaviourWithOptionsSpecs.cs
@@ -91,6 +91,59 @@ public void Should_inject_behaviour_before_executing_user_delegate()
             injectedBehaviourExecuted.Should().BeTrue();
         }
 
+        #region BeforeInject
+        [Fact]
+        public void Should_call_before_inject_callback_before_injecting_behavior()
+        {
+            var beforeInjectExecuted = false;
+            var injectedBehaviourExecuted = false;
+
+            var policy = MonkeyPolicy.InjectBehaviour(with =>
+                with.Behaviour(() =>
+                {
+                    beforeInjectExecuted.Should().BeTrue();
+                    injectedBehaviourExecuted = true;
+                })
+                .BeforeInject((context, cancellation) =>
+                {
+                    injectedBehaviourExecuted.Should().BeFalse();
+                    beforeInjectExecuted = true;
+                })
+                .InjectionRate(0.6)
+                .Enabled()
+            );
+
+            policy.Execute(() => { });
+
+            beforeInjectExecuted.Should().BeTrue();
+            injectedBehaviourExecuted.Should().BeTrue();
+        }
+
+        public void Should_not_call_before_inject_callback_if_not_injecting()
+        {
+            var beforeInjectExecuted = false;
+            var behaviorExecuted = false;
+
+            var policy = MonkeyPolicy.InjectBehaviour(with =>
+                with.Behaviour(() =>
+                {
+                    behaviorExecuted = true;
+                })
+                .BeforeInject((context, cancellation) =>
+                {
+                    beforeInjectExecuted = true;
+                })
+                .InjectionRate(0.4)
+                .Enabled()
+            );
+
+            policy.Execute(() => { });
+
+            beforeInjectExecuted.Should().BeFalse();
+            behaviorExecuted.Should().BeFalse();
+        }
+        #endregion
+
         #region invalid threshold on configuration and execution time
 
         [Fact]

diff --git a/src/Polly.Contrib.Simmy/AsyncMonkeyEngine.cs b/src/Polly.Contrib.Simmy/AsyncMonkeyEngine.cs
@@ -41,10 +41,15 @@ internal static class AsyncMonkeyEngine
             Func<Context, CancellationToken, Task> injectedBehaviour,
             Func<Context, CancellationToken, Task<Double>> injectionRate,
             Func<Context, CancellationToken, Task<bool>> enabled,
+            Func<Context, CancellationToken, Task> beforeInjectCallback,
             bool continueOnCapturedContext)
         {
             if (await ShouldInjectAsync(context, cancellationToken, injectionRate, enabled, continueOnCapturedContext).ConfigureAwait(continueOnCapturedContext))
             {
+                if (beforeInjectCallback != null)
+                {
+                    await beforeInjectCallback(context, cancellationToken);
+                }
                 await injectedBehaviour(context, cancellationToken).ConfigureAwait(continueOnCapturedContext);
             }
 
@@ -60,6 +65,7 @@ internal static class AsyncMonkeyEngine
             Func<Context, CancellationToken, Task<Exception>> injectedException,
             Func<Context, CancellationToken, Task<Double>> injectionRate,
             Func<Context, CancellationToken, Task<bool>> enabled,
+            Func<Context, CancellationToken, Task> beforeInjectCallback,
             bool continueOnCapturedContext)
         {
             if (await ShouldInjectAsync(context, cancellationToken, injectionRate, enabled, continueOnCapturedContext).ConfigureAwait(continueOnCapturedContext))
@@ -71,6 +77,10 @@ internal static class AsyncMonkeyEngine
 
                 if (exception != null)
                 {
+                    if (beforeInjectCallback != null)
+                    {
+                        await beforeInjectCallback(context, cancellationToken);
+                    }
                     throw exception;
                 }
             }
@@ -85,10 +95,15 @@ internal static class AsyncMonkeyEngine
             Func<Context, CancellationToken, Task<TResult>> injectedResult,
             Func<Context, CancellationToken, Task<Double>> injectionRate,
             Func<Context, CancellationToken, Task<bool>> enabled,
+            Func<Context, CancellationToken, Task> beforeInjectCallback,
             bool continueOnCapturedContext)
         {
             if (await ShouldInjectAsync(context, cancellationToken, injectionRate, enabled, continueOnCapturedContext).ConfigureAwait(continueOnCapturedContext))
             {
+                if (beforeInjectCallback != null)
+                {
+                    await beforeInjectCallback(context, cancellationToken);
+                }
                 return await injectedResult(context, cancellationToken).ConfigureAwait(continueOnCapturedContext);
             }
 

diff --git a/src/Polly.Contrib.Simmy/Behavior/AsyncInjectBehaviourPolicy.cs b/src/Polly.Contrib.Simmy/Behavior/AsyncInjectBehaviourPolicy.cs
@@ -10,6 +10,7 @@ namespace Polly.Contrib.Simmy.Behavior
     public class AsyncInjectBehaviourPolicy : AsyncMonkeyPolicy
     {
         private readonly Func<Context, CancellationToken, Task> _behaviour;
+        private readonly Func<Context, CancellationToken, Task> _beforeInjectCallback;
 
         [Obsolete]
         internal AsyncInjectBehaviourPolicy(Func<Context, CancellationToken, Task> behaviour, Func<Context, CancellationToken, Task<Double>> injectionRate, Func<Context, CancellationToken, Task<bool>> enabled)
@@ -22,6 +23,7 @@ internal AsyncInjectBehaviourPolicy(InjectBehaviourAsyncOptions options)
             : base(options.InjectionRate, options.Enabled)
         {
             _behaviour = options.BehaviourInternal ?? throw new ArgumentNullException(nameof(options.BehaviourInternal));
+            _beforeInjectCallback = options.BeforeInjectCallback;
         }
 
         /// <inheritdoc/>
@@ -35,6 +37,7 @@ internal AsyncInjectBehaviourPolicy(InjectBehaviourAsyncOptions options)
                 _behaviour,
                 InjectionRate,
                 Enabled,
+                _beforeInjectCallback,
                 continueOnCapturedContext);
         }
     }
@@ -46,6 +49,7 @@ internal AsyncInjectBehaviourPolicy(InjectBehaviourAsyncOptions options)
     public class AsyncInjectBehaviourPolicy<TResult> : AsyncMonkeyPolicy<TResult>
     {
         private readonly Func<Context, CancellationToken, Task> _behaviour;
+        private readonly Func<Context, CancellationToken, Task> _beforeInjectCallback;
 
         [Obsolete]
         internal AsyncInjectBehaviourPolicy(Func<Context, CancellationToken, Task> behaviour, Func<Context, CancellationToken, Task<Double>> injectionRate, Func<Context, CancellationToken, Task<bool>> enabled)
@@ -58,6 +62,7 @@ internal AsyncInjectBehaviourPolicy(InjectBehaviourAsyncOptions options)
             : base(options.InjectionRate, options.Enabled)
         {
             _behaviour = options.BehaviourInternal ?? throw new ArgumentNullException(nameof(options.BehaviourInternal));
+            _beforeInjectCallback = options.BeforeInjectCallback;
         }
 
         /// <inheritdoc/>
@@ -70,6 +75,7 @@ protected override Task<TResult> ImplementationAsync(Func<Context, CancellationT
                 _behaviour,
                 InjectionRate,
                 Enabled,
+                _beforeInjectCallback,
                 continueOnCapturedContext);
         }
     }

diff --git a/src/Polly.Contrib.Simmy/Behavior/InjectBehaviourPolicy.cs b/src/Polly.Contrib.Simmy/Behavior/InjectBehaviourPolicy.cs
@@ -9,6 +9,7 @@ namespace Polly.Contrib.Simmy.Behavior
     public class InjectBehaviourPolicy : Simmy.MonkeyPolicy
     {
         private readonly Action<Context, CancellationToken> _behaviour;
+        private readonly Action<Context, CancellationToken> _beforeInjectCallback;
 
         [Obsolete]
         internal InjectBehaviourPolicy(Action<Context, CancellationToken> behaviour, Func<Context, CancellationToken, double> injectionRate, Func<Context, CancellationToken, bool> enabled) 
@@ -21,6 +22,7 @@ internal InjectBehaviourPolicy(InjectBehaviourOptions options)
             : base(options.InjectionRate, options.Enabled)
         {
             _behaviour = options.BehaviourInternal ?? throw new ArgumentNullException(nameof(options.BehaviourInternal));
+            _beforeInjectCallback = options.BeforeInjectCallback;
         }
 
         /// <inheritdoc/>
@@ -32,7 +34,8 @@ protected override TResult Implementation<TResult>(Func<Context, CancellationTok
                 cancellationToken,
                 (ctx, ct) => _behaviour(ctx, ct),
                 InjectionRate,
-                Enabled);
+                Enabled,
+                _beforeInjectCallback);
         }
     }
     /// <summary>
@@ -42,6 +45,7 @@ protected override TResult Implementation<TResult>(Func<Context, CancellationTok
     public class InjectBehaviourPolicy<TResult> : MonkeyPolicy<TResult>
     {
         private readonly Action<Context, CancellationToken> _behaviour;
+        private readonly Action<Context, CancellationToken> _beforeInjectCallback;
 
         [Obsolete]
         internal InjectBehaviourPolicy(Action<Context, CancellationToken> behaviour, Func<Context, CancellationToken, double> injectionRate, Func<Context, CancellationToken, bool> enabled)
@@ -54,6 +58,7 @@ internal InjectBehaviourPolicy(InjectBehaviourOptions options)
             : base(options.InjectionRate, options.Enabled)
         {
             _behaviour = options.BehaviourInternal ?? throw new ArgumentNullException(nameof(options.BehaviourInternal));
+            _beforeInjectCallback = options.BeforeInjectCallback;
         }
 
         /// <inheritdoc/>
@@ -65,7 +70,8 @@ protected override TResult Implementation(Func<Context, CancellationToken, TResu
                 cancellationToken,
                 (ctx, ct) => _behaviour(ctx, ct),
                 InjectionRate,
-                Enabled);
+                Enabled,
+                _beforeInjectCallback);
         }
     }
 }
diff --git a/src/Polly.Contrib.Simmy/InjectOptionsAsyncBase.cs b/src/Polly.Contrib.Simmy/InjectOptionsAsyncBase.cs
@@ -18,5 +18,10 @@ public abstract class InjectOptionsAsyncBase
         /// Lambda to check if this policy is enabled in current context
         /// </summary>
         internal Func<Context, CancellationToken, Task<bool>> Enabled { get; set; }
+
+        /// <summary>
+        /// Lambda to call immediately before injecting
+        /// </summary>
+        internal virtual Func<Context, CancellationToken, Task> BeforeInjectCallback { get; set; }
     }
 }
diff --git a/src/Polly.Contrib.Simmy/InjectOptionsAsyncBaseExtensions.cs b/src/Polly.Contrib.Simmy/InjectOptionsAsyncBaseExtensions.cs
@@ -60,5 +60,17 @@ public static InjectOptionsAsyncBase InjectionRate(this InjectOptionsAsyncBase o
             options.InjectionRate = injectionRateProvider;
             return options;
         }
+
+        /// <summary>
+        /// Configure a callback to be run immediately before injecting chaos.
+        /// </summary>
+        /// <param name="options">The configuration object.</param>
+        /// <param name="beforeInjectCallback">A delegate to be run immediately before injecting chaos.</param>
+        /// <returns></returns>
+        public static InjectOptionsAsyncBase BeforeInject(this InjectOptionsAsyncBase options, Func<Context, CancellationToken, Task> beforeInjectCallback)
+        {
+            options.BeforeInjectCallback = beforeInjectCallback;
+            return options;
+        }
     }
 }
diff --git a/src/Polly.Contrib.Simmy/InjectOptionsBase.cs b/src/Polly.Contrib.Simmy/InjectOptionsBase.cs
@@ -17,5 +17,10 @@ public abstract class InjectOptionsBase
         /// Lambda to check if this policy is enabled in current context
         /// </summary>
         internal Func<Context, CancellationToken, bool> Enabled { get; set; }
+
+        /// <summary>
+        /// Lambda to call immediately before injecting
+        /// </summary>
+        internal Action<Context, CancellationToken> BeforeInjectCallback { get; set; }
     }
 }
diff --git a/src/Polly.Contrib.Simmy/InjectOptionsBaseExtensions.cs b/src/Polly.Contrib.Simmy/InjectOptionsBaseExtensions.cs
@@ -59,5 +59,17 @@ public static InjectOptionsBase InjectionRate(this InjectOptionsBase options, Fu
             options.InjectionRate = injectionRateProvider;
             return options;
         }
+
+        /// <summary>
+        /// Configure a callback to be run immediately before injecting chaos.
+        /// </summary>
+        /// <param name="options">The configuration object.</param>
+        /// <param name="beforeInjectCallback">A delegate to be run immediately before injecting chaos.</param>
+        /// <returns></returns>
+        public static InjectOptionsBase BeforeInject(this InjectOptionsBase options, Action<Context, CancellationToken> beforeInjectCallback)
+        {
+            options.BeforeInjectCallback = beforeInjectCallback;
+            return options;
+        }
     }
 }
diff --git a/src/Polly.Contrib.Simmy/Latency/AsyncInjectLatencyPolicy.cs b/src/Polly.Contrib.Simmy/Latency/AsyncInjectLatencyPolicy.cs
@@ -11,6 +11,7 @@ namespace Polly.Contrib.Simmy.Latency
     public class AsyncInjectLatencyPolicy : AsyncMonkeyPolicy
     {
         private readonly Func<Context, CancellationToken, Task<TimeSpan>> _latencyProvider;
+        private readonly Func<Context, CancellationToken, Task> _beforeInjectCallback;
 
         [Obsolete]
         internal AsyncInjectLatencyPolicy(
@@ -26,6 +27,7 @@ internal AsyncInjectLatencyPolicy(InjectLatencyAsyncOptions options)
             : base(options.InjectionRate, options.Enabled)
         {
             _latencyProvider = options.LatencyInternal ?? throw new ArgumentNullException(nameof(options.LatencyInternal));
+            _beforeInjectCallback = options.BeforeInjectCallback;
         }
 
         /// <inheritdoc/>
@@ -52,6 +54,7 @@ internal AsyncInjectLatencyPolicy(InjectLatencyAsyncOptions options)
                 },
                 InjectionRate,
                 Enabled,
+                _beforeInjectCallback,
                 continueOnCapturedContext);
         }
     }
@@ -62,6 +65,7 @@ internal AsyncInjectLatencyPolicy(InjectLatencyAsyncOptions options)
     public class AsyncInjectLatencyPolicy<TResult> : AsyncMonkeyPolicy<TResult>
     {
         private readonly Func<Context, CancellationToken, Task<TimeSpan>> _latencyProvider;
+        private readonly Func<Context, CancellationToken, Task> _beforeInjectCallback;
 
         [Obsolete]
         internal AsyncInjectLatencyPolicy(
@@ -77,6 +81,7 @@ internal AsyncInjectLatencyPolicy(InjectLatencyAsyncOptions options)
             : base(options.InjectionRate, options.Enabled)
         {
             _latencyProvider = options.LatencyInternal ?? throw new ArgumentNullException(nameof(options.LatencyInternal));
+            _beforeInjectCallback = options.BeforeInjectCallback;
         }
 
         /// <inheritdoc/>
@@ -103,6 +108,7 @@ internal AsyncInjectLatencyPolicy(InjectLatencyAsyncOptions options)
                 },
                 InjectionRate,
                 Enabled,
+                _beforeInjectCallback,
                 continueOnCapturedContext);
         }
     }