test: fix TestGetMergeRegionSizeAndCount test (#52241)

close #52095
pingcap · Apr 10, 2024 · 5fc42c9 · 5fc42c9
1 parent 2b6b78a
commit 5fc42c9
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 11 deletions.
diff --git a/br/pkg/conn/BUILD.bazel b/br/pkg/conn/BUILD.bazel
@@ -61,5 +61,6 @@ go_test(
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",
         "@org_uber_go_goleak//:goleak",
+        "@org_uber_go_multierr//:multierr",
     ],
 )
diff --git a/br/pkg/conn/conn.go b/br/pkg/conn/conn.go
@@ -88,16 +88,20 @@ func GetAllTiKVStoresWithRetry(ctx context.Context,
 		func() error {
 			stores, err = util.GetAllTiKVStores(ctx, pdClient, storeBehavior)
 			failpoint.Inject("hint-GetAllTiKVStores-error", func(val failpoint.Value) {
+				logutil.CL(ctx).Debug("failpoint hint-GetAllTiKVStores-error injected.")
 				if val.(bool) {
-					logutil.CL(ctx).Debug("failpoint hint-GetAllTiKVStores-error injected.")
 					err = status.Error(codes.Unknown, "Retryable error")
+				} else {
+					err = context.Canceled
 				}
 			})
 
 			failpoint.Inject("hint-GetAllTiKVStores-cancel", func(val failpoint.Value) {
+				logutil.CL(ctx).Debug("failpoint hint-GetAllTiKVStores-cancel injected.")
 				if val.(bool) {
-					logutil.CL(ctx).Debug("failpoint hint-GetAllTiKVStores-cancel injected.")
 					err = status.Error(codes.Canceled, "Cancel Retry")
+				} else {
+					err = context.Canceled
 				}
 			})
 

diff --git a/br/pkg/conn/conn_test.go b/br/pkg/conn/conn_test.go
@@ -19,14 +19,17 @@ import (
 	"github.com/pingcap/tidb/br/pkg/pdutil"
 	"github.com/pingcap/tidb/br/pkg/utils"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/multierr"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
 
 func TestGetAllTiKVStoresWithRetryCancel(t *testing.T) {
-	_ = failpoint.Enable("github.com/pingcap/tidb/br/pkg/conn/hint-GetAllTiKVStores-cancel", "return(true)")
+	err := failpoint.Enable("github.com/pingcap/tidb/br/pkg/conn/hint-GetAllTiKVStores-cancel", "1*return(true)->1*return(false)")
+	require.NoError(t, err)
 	defer func() {
-		_ = failpoint.Disable("github.com/pingcap/tidb/br/pkg/conn/hint-GetAllTiKVStores-cancel")
+		err = failpoint.Disable("github.com/pingcap/tidb/br/pkg/conn/hint-GetAllTiKVStores-cancel")
+		require.NoError(t, err)
 	}()
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -58,15 +61,19 @@ func TestGetAllTiKVStoresWithRetryCancel(t *testing.T) {
 		Stores: stores,
 	}
 
-	_, err := GetAllTiKVStoresWithRetry(ctx, fpdc, util.SkipTiFlash)
+	_, err = GetAllTiKVStoresWithRetry(ctx, fpdc, util.SkipTiFlash)
 	require.Error(t, err)
-	require.Equal(t, codes.Canceled, status.Code(errors.Cause(err)))
+	errs := multierr.Errors(err)
+	require.Equal(t, 2, len(errs))
+	require.Equal(t, codes.Canceled, status.Code(errors.Cause(errs[0])))
 }
 
 func TestGetAllTiKVStoresWithUnknown(t *testing.T) {
-	_ = failpoint.Enable("github.com/pingcap/tidb/br/pkg/conn/hint-GetAllTiKVStores-error", "return(true)")
+	err := failpoint.Enable("github.com/pingcap/tidb/br/pkg/conn/hint-GetAllTiKVStores-error", "1*return(true)->1*return(false)")
+	require.NoError(t, err)
 	defer func() {
-		_ = failpoint.Disable("github.com/pingcap/tidb/br/pkg/conn/hint-GetAllTiKVStores-error")
+		err = failpoint.Disable("github.com/pingcap/tidb/br/pkg/conn/hint-GetAllTiKVStores-error")
+		require.NoError(t, err)
 	}()
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -98,10 +105,13 @@ func TestGetAllTiKVStoresWithUnknown(t *testing.T) {
 		Stores: stores,
 	}
 
-	_, err := GetAllTiKVStoresWithRetry(ctx, fpdc, util.SkipTiFlash)
+	_, err = GetAllTiKVStoresWithRetry(ctx, fpdc, util.SkipTiFlash)
 	require.Error(t, err)
-	require.Equal(t, codes.Unknown, status.Code(errors.Cause(err)))
+	errs := multierr.Errors(err)
+	require.Equal(t, 2, len(errs))
+	require.Equal(t, codes.Unknown, status.Code(errors.Cause(errs[0])))
 }
+
 func TestCheckStoresAlive(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -405,14 +415,18 @@ func TestGetMergeRegionSizeAndCount(t *testing.T) {
 		},
 	}
 
-	ctx := context.Background()
+	pctx := context.Background()
 	for _, ca := range cases {
+		ctx, cancel := context.WithCancel(pctx)
 		pdCli := utils.FakePDClient{Stores: ca.stores}
 		require.Equal(t, len(ca.content), len(ca.stores))
 		count := 0
 		mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			switch strings.TrimSpace(r.URL.Path) {
 			case "/config":
+				if len(ca.content[count]) == 0 {
+					cancel()
+				}
 				_, _ = fmt.Fprint(w, ca.content[count])
 			default:
 				http.NotFoundHandler().ServeHTTP(w, r)