diff --git a/retry/backoffhandler.go b/retry/backoffhandler.go
index 8c09db55..33f866d0 100644
--- a/retry/backoffhandler.go
+++ b/retry/backoffhandler.go
@@ -91,10 +91,12 @@ func (b *BackoffHandler) Backoff(ctx context.Context) bool {
 
 // Sets a grace period within which the the backoff timer is maintained. After the grace
 // period expires, the number of retries & backoff duration is reset.
-func (b *BackoffHandler) SetGracePeriod() {
+func (b *BackoffHandler) SetGracePeriod() time.Duration {
 	maxTimeToWait := b.GetBaseTime() * 2 << (b.retries + 1)
 	timeToWait := time.Duration(rand.Int63n(maxTimeToWait.Nanoseconds()))
 	b.resetDeadline = Clock.Now().Add(timeToWait)
+
+	return timeToWait
 }
 
 func (b BackoffHandler) GetBaseTime() time.Duration {
diff --git a/retry/backoffhandler_test.go b/retry/backoffhandler_test.go
index 8a4d1fe3..93988515 100644
--- a/retry/backoffhandler_test.go
+++ b/retry/backoffhandler_test.go
@@ -57,9 +57,9 @@ func TestBackoffGracePeriod(t *testing.T) {
 		t.Fatalf("backoff failed immediately")
 	}
 	// the next call to Backoff would fail unless it's after the grace period
-	backoff.SetGracePeriod()
-	// advance time to after the grace period (~4 seconds) and see what happens
-	currentTime = currentTime.Add(time.Second * 5)
+	gracePeriod := backoff.SetGracePeriod()
+	// advance time to after the grace period, which at most will be 8 seconds, but we will advance +1 second.
+	currentTime = currentTime.Add(gracePeriod + time.Second)
 	if !backoff.Backoff(ctx) {
 		t.Fatalf("backoff failed after the grace period expired")
 	}