From 2ac4067e29c9e5ad6410f38dbd4ba19c162c7c4c Mon Sep 17 00:00:00 2001
From: Cody Jones <cody.dw.jones@gmail.com>
Date: Sat, 13 Mar 2021 08:54:20 -0800
Subject: [PATCH] Return context error if that's why attempts are ceasing

---
 retry.go      | 8 ++++----
 retry_test.go | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/retry.go b/retry.go
index 510bb17..cb5a13d 100644
--- 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
index 15432cd..7e3253b 100644
--- 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")
 	}
 }