Search code examples
c#unit-testingxunitpollyretry-logic

Unit test sleep durations polly wait retry policy


I wanted to write a unit test to see if the execution is put to sleep for the specified duration. I have come across the SystemClock which is part of Polly.Utilities, but I'm looking for the implementation similar to Polly unit test as mentioned here WaitAndRetrySpecs, which looks like this

 [Fact]
 public void Should_sleep_for_the_specified_duration_each_retry_when_specified_exception_thrown_same_number_of_times_as_there_are_sleep_durations()
 {
      var totalTimeSlept = 0;

      var policy = Policy
            .Handle<DivideByZeroException>()
            .WaitAndRetry(new[]
            {
               1.Seconds(),
               2.Seconds(),
               3.Seconds()
            });

      SystemClock.Sleep = span => totalTimeSlept += span.Seconds;

      policy.RaiseException<DivideByZeroException>(3);

      totalTimeSlept.Should()
                      .Be(1 + 2 + 3);
  }

Currently my policy looks like

var customPolicy = Policy
                     .Handle<SqlException>(x => IsTransientError(x))
                     .WaitAndRetryAsync(
                      3,
                      (retryAttempt) => getSleepDurationByRetryAtempt(retryAttempt)
                     );

I want to test overall time slept for policy. For each retry attempts [1,2,3] sleep durations are [1,2,3]. After all 3 retries the total sleep duration should be 1 + 2 + 3 = 6. This test very similar to Polly specs mentioned in the link above.

Question: How do I write the unit test for the customPolicy to test the total sleep duration similar to the polly specs. I want to see a implementation or directions to write the unit test.


Solution

  • You can achieve this by utilizing the onRetry function.

    To make it simple let me define the IsTransientError and GetSleepDurationByRetryAttempt methods like this:

    public TimeSpan GetSleepDurationByRetryAttempt(int attempt) => TimeSpan.FromSeconds(attempt);
    public bool IsTransientError(SqlException ex) => true;
    

    BTW you can shorten your policy definition by avoiding (unnecessary) anonymous lambdas:

    var customPolicy = Policy
        .Handle<SqlException>(IsTransientError)
        .WaitAndRetryAsync(3, GetSleepDurationByRetryAttempt)
    

    So, back to the onRetry. There is an overload which have the following signature: Action<Exception, TimeSpan, Context>. Here the second parameter is the sleep duration.

    All we need to do is to provide a function here, which accumulates the sleep durations.

    var totalSleepDuration = TimeSpan.Zero;
    ...
    onRetry: (ex, duration, ctx) => { totalSleepDuration = totalSleepDuration.Add(duration); }
    

    Let's put all these together:

    [Fact]
    public async Task GivenACustomSleepDurationProvider_WhenIUseItInARetryPolicy_ThenTheAccumulatedDurationIsAsExpected()
    {
        //Arrange
        var totalSleepDuration = TimeSpan.Zero;
        var customPolicy = Policy
            .Handle<SqlException>(IsTransientError)
            .WaitAndRetryAsync(3, GetSleepDurationByRetryAttempt,
                onRetry: (ex, duration, ctx) => { totalSleepDuration = totalSleepDuration.Add(duration); }
            );
    
        //Act
        Func<Task> actionWithRetry = async() => await customPolicy.ExecuteAsync(() => throw new SqlException());
        
        //Assert
        _ = await Assert.ThrowsAsync<SqlException>(actionWithRetry);
        Assert.Equal(6, totalSleepDuration.Seconds);
    }
    

    UPDATE #1: Reduce delays and introduce theory

    Depending on your requirements it might make sense to run this same test case with different parameters. That's where Theory and InlineData can help you:

    [Theory]
    [InlineData(3, 600)]
    [InlineData(4, 1000)]
    [InlineData(5, 1500)]
    public async Task GivenACustomSleepDurationProvider_WhenIUseItInARetryPolicy_ThenTheAccumulatedDurationIsAsExpected(int retryCount, int expectedTotalSleepInMs)
    {
        //Arrange
        var totalSleepDuration = TimeSpan.Zero;
        var customPolicy = Policy
            .Handle<SqlException>(IsTransientError)
            .WaitAndRetryAsync(retryCount, GetSleepDurationByRetryAttempt,
                onRetry: (ex, duration, ctx) => { totalSleepDuration = totalSleepDuration.Add(duration); }
            );
    
        //Act
        Func<Task> actionWithRetry = async () => await customPolicy.ExecuteAsync(() => throw new SqlException());
    
        //Assert
        _ = await Assert.ThrowsAsync<SqlException>(actionWithRetry);
        Assert.Equal(TimeSpan.FromMilliseconds(expectedTotalSleepInMs), totalSleepDuration);
    }
    
    public static TimeSpan GetSleepDurationByRetryAttempt(int attempt) => TimeSpan.FromMilliseconds(attempt * 100);
    

    UPDATE #2: Passing the TimeSpan via the Context

    In order to make the transfer and retrieval of the TimeSpan a bit more type-safe we can create two extension methods for this:

    public static class ContextExtensions
    {
        private const string Accumulator = "DurationAccumulator";
    
        public static Context SetAccumulator(this Context context, TimeSpan durationAccumulator)
        {
            context[Accumulator] = durationAccumulator;
            return context;
        }
    
        public static TimeSpan? GetAccumulator(this Context context)
        {
            if (!context.TryGetValue(Accumulator, out var ts))
                return null;
    
            if (ts is TimeSpan accumulator) 
                return accumulator;
    
            return null;
        }
    }
    

    We can also extract the Policy creation logic:

    private AsyncPolicy GetCustomPolicy(int retryCount)
        => Policy
            .Handle<SqlException>(IsTransientError)
            .WaitAndRetryAsync(retryCount, GetSleepDurationByRetryAttempt,
                onRetry: (ex, duration, ctx) =>
                {
                    var totalSleepDuration = ctx.GetAccumulator();
                    if (!totalSleepDuration.HasValue) return;
                    totalSleepDuration = totalSleepDuration.Value.Add(duration);
                    ctx.SetAccumulator(totalSleepDuration.Value);
                });
    

    Now let's put all these together (once again):

    [Theory]
    [InlineData(3, 600)]
    [InlineData(4, 1000)]
    [InlineData(5, 1500)]
    public async Task GivenACustomSleepDurationProvider_WhenIUseItInARetryPolicy_ThenTheAccumulatedDurationIsAsExpected(
            int retryCount, int expectedTotalSleepInMs)
    {
        //Arrange
        var totalSleepDuration = TimeSpan.Zero;
        var customPolicy = GetCustomPolicy(retryCount);
        var context = new Context().SetAccumulator(totalSleepDuration);
    
        //Act
        Func<Task> actionWithRetry = async () => await customPolicy.ExecuteAsync(ctx => throw new SqlException(), context);
    
        //Assert
        _ = await Assert.ThrowsAsync<SqlException>(actionWithRetry);
        var accumulator = context.GetAccumulator();
        Assert.NotNull(accumulator);
        Assert.Equal(TimeSpan.FromMilliseconds(expectedTotalSleepInMs), accumulator.Value);
    }