Return context error if that's why attempts are ceasing

Rican7 · Mar 13, 2021 · 2ac4067 · 2ac4067
1 parent 2c55b7d
commit 2ac4067
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/retry.go b/retry.go
@@ -33,14 +33,14 @@ func Retry(action Action, strategies ...strategy.Strategy) error {
 func RetryWithContext(ctx context.Context, action ActionWithContext, strategies ...strategy.Strategy) error {
 	var err error
 
-	if ctx.Err() != nil {
-		return ctx.Err()
-	}
-
 	for attempt := uint(0); (0 == attempt || nil != err && nil == ctx.Err()) && shouldAttempt(attempt, sleepFunc(ctx), strategies...); attempt++ {
 		err = action(ctx, attempt)
 	}
 
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
 	return err
 }
 

diff --git a/retry_test.go b/retry_test.go
@@ -103,8 +103,8 @@ func TestRetryWithContextSleepIsInterrupted(t *testing.T) {
 		t.Errorf("expected the action to be tried once, not %d times", numCalls)
 	}
 
-	if expectedErr != err {
-		t.Error("expected to receive the error returned by the action")
+	if context.DeadlineExceeded != err {
+		t.Error("expected a context error")
 	}
 }