Add tests on retryAfterFactory

src/Polly.Specs/Helpers/RateLimit/ResultClassWithRetryAfter.cs

@@ -0,0 +1,21 @@
+using System;
+
+namespace Polly.Specs.Helpers.RateLimit
+{
+    internal class ResultClassWithRetryAfter : ResultClass
+    {
+        public TimeSpan RetryAfter { get; }
+
+        public ResultClassWithRetryAfter(ResultPrimitive result)
+        : base(result)
+        {
+            RetryAfter = TimeSpan.Zero;
+        }
+
+        public ResultClassWithRetryAfter(TimeSpan retryAfter)
+        : base(ResultPrimitive.Undefined)
+        {
+            RetryAfter = retryAfter;
+        }
+    }
+}

src/Polly.Specs/RateLimit/AsyncRateLimitPolicyTResultSpecs.cs

@@ -2,13 +2,14 @@
 using System.Threading.Tasks;
 using Polly.RateLimit;
 using Polly.Specs.Helpers;
+using Polly.Specs.Helpers.RateLimit;
 using Polly.Utilities;
 using Xunit;
 
 namespace Polly.Specs.RateLimit
 {
     [Collection(Polly.Specs.Helpers.Constants.SystemClockDependentTestCollection)]
-    public class AsyncRateLimitPolicyTResultSpecs : RateLimitPolicySpecsBase, IDisposable
+    public class AsyncRateLimitPolicyTResultSpecs : RateLimitPolicyTResultSpecsBase, IDisposable
     {
         public void Dispose()
         {
@@ -17,21 +18,27 @@ public void Dispose()
 
         protected override IRateLimitPolicy GetPolicyViaSyntax(int numberOfExecutions, TimeSpan perTimeSpan)
         {
-            return Policy.RateLimitAsync<ResultClass>(numberOfExecutions, perTimeSpan);
+            return Policy.RateLimitAsync<ResultClassWithRetryAfter>(numberOfExecutions, perTimeSpan);
         }
 
         protected override IRateLimitPolicy GetPolicyViaSyntax(int numberOfExecutions, TimeSpan perTimeSpan, int maxBurst)
         {
-            return Policy.RateLimitAsync<ResultClass>(numberOfExecutions, perTimeSpan, maxBurst);
+            return Policy.RateLimitAsync<ResultClassWithRetryAfter>(numberOfExecutions, perTimeSpan, maxBurst);
+        }
+
+        protected override IRateLimitPolicy<TResult> GetPolicyViaSyntax<TResult>(int numberOfExecutions, TimeSpan perTimeSpan, int maxBurst,
+            Func<TimeSpan, Context, TResult> retryAfterFactory)
+        {
+            return Policy.RateLimitAsync<TResult>(numberOfExecutions, perTimeSpan, maxBurst, retryAfterFactory);
         }
 
         protected override (bool, TimeSpan) TryExecuteThroughPolicy(IRateLimitPolicy policy)
         {
-            if (policy is AsyncRateLimitPolicy<ResultClass> typedPolicy)
+            if (policy is AsyncRateLimitPolicy<ResultClassWithRetryAfter> typedPolicy)
             {
                 try
                 {
-                    typedPolicy.ExecuteAsync(() => Task.FromResult(new ResultClass(ResultPrimitive.Good))).GetAwaiter().GetResult();
+                    typedPolicy.ExecuteAsync(() => Task.FromResult(new ResultClassWithRetryAfter(ResultPrimitive.Good))).GetAwaiter().GetResult();
                     return (true, TimeSpan.Zero);
                 }
                 catch (RateLimitRejectedException e)
@@ -44,5 +51,17 @@ protected override (bool, TimeSpan) TryExecuteThroughPolicy(IRateLimitPolicy pol
                 throw new InvalidOperationException("Unexpected policy type in test construction.");
             }
         }
+
+        protected override TResult TryExecuteThroughPolicy<TResult>(IRateLimitPolicy<TResult> policy, TResult resultIfExecutionPermitted)
+        {
+            if (policy is AsyncRateLimitPolicy<TResult> typedPolicy)
+            {
+                return typedPolicy.ExecuteAsync(() => Task.FromResult(resultIfExecutionPermitted)).GetAwaiter().GetResult();
+            }
+            else
+            {
+                throw new InvalidOperationException("Unexpected policy type in test construction.");
+            }
+        }
     }
 }

src/Polly.Specs/RateLimit/RateLimitPolicyTResultSpecsBase.cs

@@ -0,0 +1,47 @@
+using System;
+using FluentAssertions;
+using Polly.RateLimit;
+using Polly.Specs.Helpers;
+using Polly.Specs.Helpers.RateLimit;
+using Xunit;
+
+namespace Polly.Specs.RateLimit
+{
+    public abstract class RateLimitPolicyTResultSpecsBase : RateLimitPolicySpecsBase
+    {
+        protected abstract IRateLimitPolicy<TResult> GetPolicyViaSyntax<TResult>(
+            int numberOfExecutions,
+            TimeSpan perTimeSpan,
+            int maxBurst,
+            Func<TimeSpan, Context, TResult> retryAfterFactory);
+
+        protected abstract TResult TryExecuteThroughPolicy<TResult>(IRateLimitPolicy<TResult> policy, TResult resultIfExecutionPermitted);
+
+        [Theory]
+        [InlineData(1)]
+        [InlineData(2)]
+        [InlineData(5)]
+        public void Ratelimiter_specifies_correct_wait_until_next_execution_by_custom_factory(int onePerSeconds)
+        {
+            FixClock();
+
+            // Arrange
+            TimeSpan onePer = TimeSpan.FromSeconds(onePerSeconds);
+            Func<TimeSpan, Context, ResultClassWithRetryAfter> retryAfterFactory = (t, ctx) => new ResultClassWithRetryAfter(t);
+            var rateLimiter = GetPolicyViaSyntax<ResultClassWithRetryAfter>(1, onePer, 1, retryAfterFactory);
+
+            // Assert - first execution after initialising should always be permitted.
+            ShouldPermitAnExecution(rateLimiter);
+
+            // Arrange
+            // (do nothing - time not advanced)
+
+            // Act - try another execution.
+            var resultExpectedBlocked = TryExecuteThroughPolicy(rateLimiter, new ResultClassWithRetryAfter(ResultPrimitive.Good));
+
+            // Assert - should be blocked - time not advanced. Result should be expressed per the retryAfterFactory.
+            resultExpectedBlocked.ResultCode.Should().NotBe(ResultPrimitive.Good);
+            resultExpectedBlocked.RetryAfter.Should().Be(onePer);
+        }
+    }
+}