tests: make TestFallbackTSOConsistency stable (#7408)

close #7407 Signed-off-by: lhy1024 <admin@liudos.us>
tikv · Nov 22, 2023 · 0e2dad1 · 0e2dad1
1 parent 57fb020
commit 0e2dad1
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/tests/integrations/tso/consistency_test.go b/tests/integrations/tso/consistency_test.go
@@ -27,7 +27,7 @@ import (
 	tso "github.com/tikv/pd/pkg/mcs/tso/server"
 	tsopkg "github.com/tikv/pd/pkg/tso"
 	"github.com/tikv/pd/pkg/utils/tempurl"
-	pd "github.com/tikv/pd/pkg/utils/testutil"
+	tu "github.com/tikv/pd/pkg/utils/testutil"
 	"github.com/tikv/pd/pkg/utils/tsoutil"
 	"github.com/tikv/pd/tests"
 	"google.golang.org/grpc"
@@ -82,7 +82,7 @@ func (suite *tsoConsistencyTestSuite) SetupSuite() {
 	suite.pdLeaderServer = suite.cluster.GetServer(leaderName)
 	backendEndpoints := suite.pdLeaderServer.GetAddr()
 	if suite.legacy {
-		suite.pdClient = pd.MustNewGrpcClient(re, backendEndpoints)
+		suite.pdClient = tu.MustNewGrpcClient(re, backendEndpoints)
 	} else {
 		suite.tsoServer, suite.tsoServerCleanup = tests.StartSingleTSOTestServer(suite.ctx, re, backendEndpoints, tempurl.Alloc())
 		suite.tsoClientConn, suite.tsoClient = tso.MustNewGrpcClient(re, suite.tsoServer.GetAddr())
@@ -127,12 +127,15 @@ func (suite *tsoConsistencyTestSuite) request(ctx context.Context, count uint32)
 		DcLocation: tsopkg.GlobalDCLocation,
 		Count:      count,
 	}
-	tsoClient, err := suite.tsoClient.Tso(ctx)
-	re.NoError(err)
-	defer tsoClient.CloseSend()
-	re.NoError(tsoClient.Send(req))
-	resp, err := tsoClient.Recv()
-	re.NoError(err)
+	var resp *tsopb.TsoResponse
+	tu.Eventually(re, func() bool {
+		tsoClient, err := suite.tsoClient.Tso(ctx)
+		re.NoError(err)
+		defer tsoClient.CloseSend()
+		re.NoError(tsoClient.Send(req))
+		resp, err = tsoClient.Recv()
+		return err == nil && resp != nil
+	})
 	return checkAndReturnTimestampResponse(re, resp)
 }
 

diff --git a/tests/integrations/tso/server_test.go b/tests/integrations/tso/server_test.go
@@ -26,7 +26,7 @@ import (
 	tso "github.com/tikv/pd/pkg/mcs/tso/server"
 	tsopkg "github.com/tikv/pd/pkg/tso"
 	"github.com/tikv/pd/pkg/utils/tempurl"
-	pd "github.com/tikv/pd/pkg/utils/testutil"
+	tu "github.com/tikv/pd/pkg/utils/testutil"
 	"github.com/tikv/pd/tests"
 	"google.golang.org/grpc"
 )
@@ -80,7 +80,7 @@ func (suite *tsoServerTestSuite) SetupSuite() {
 	suite.pdLeaderServer = suite.cluster.GetServer(leaderName)
 	backendEndpoints := suite.pdLeaderServer.GetAddr()
 	if suite.legacy {
-		suite.pdClient = pd.MustNewGrpcClient(re, backendEndpoints)
+		suite.pdClient = tu.MustNewGrpcClient(re, backendEndpoints)
 	} else {
 		suite.tsoServer, suite.tsoServerCleanup = tests.StartSingleTSOTestServer(suite.ctx, re, backendEndpoints, tempurl.Alloc())
 		suite.tsoClientConn, suite.tsoClient = tso.MustNewGrpcClient(re, suite.tsoServer.GetAddr())