feat: retry with ctx deadline (#3626)

core/fx/retry_test.go

@@ -98,19 +98,51 @@ func TestRetryWithInterval(t *testing.T) {
 }
 
 func TestRetryCtx(t *testing.T) {
-	assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {
-		if retryCount == 0 {
-			return errors.New("any")
-		}
-		time.Sleep(time.Millisecond * 150)
-		return nil
-	}, WithTimeout(time.Millisecond*250), WithInterval(time.Millisecond*150)))
-
-	assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {
-		if retryCount == 1 {
+	t.Run("with timeout", func(t *testing.T) {
+		assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {
+			if retryCount == 0 {
+				return errors.New("any")
+			}
+			time.Sleep(time.Millisecond * 150)
 			return nil
-		}
-		time.Sleep(time.Millisecond * 150)
-		return errors.New("any ")
-	}, WithTimeout(time.Millisecond*250), WithInterval(time.Millisecond*150)))
+		}, WithTimeout(time.Millisecond*250), WithInterval(time.Millisecond*150)))
+
+		assert.NotNil(t, DoWithRetryCtx(context.Background(), func(ctx context.Context, retryCount int) error {
+			if retryCount == 1 {
+				return nil
+			}
+			time.Sleep(time.Millisecond * 150)
+			return errors.New("any ")
+		}, WithTimeout(time.Millisecond*250), WithInterval(time.Millisecond*150)))
+	})
+
+	t.Run("with deadline exceeded", func(t *testing.T) {
+		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Millisecond*250))
+		defer cancel()
+
+		var times int
+		assert.Error(t, DoWithRetryCtx(ctx, func(ctx context.Context, retryCount int) error {
+			times++
+			time.Sleep(time.Millisecond * 150)
+			return errors.New("any")
+		}, WithInterval(time.Millisecond*150)))
+		assert.Equal(t, 1, times)
+	})
+
+	t.Run("with deadline not exceeded", func(t *testing.T) {
+		ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(time.Millisecond*250))
+		defer cancel()
+
+		var times int
+		assert.NoError(t, DoWithRetryCtx(ctx, func(ctx context.Context, retryCount int) error {
+			times++
+			if times == defaultRetryTimes {
+				return nil
+			}
+
+			time.Sleep(time.Millisecond * 50)
+			return errors.New("any")
+		}))
+		assert.Equal(t, defaultRetryTimes, times)
+	})
 }