From 0934e641129d8bac3c5975372b2ce492b81fd9db Mon Sep 17 00:00:00 2001
From: lhy1024 <admin@liudos.us>
Date: Mon, 19 Jun 2023 16:13:10 +0800
Subject: [PATCH] tests: reduce unnecessary time.sleep in keyspace group
 (#6632)

ref tikv/pd#6599

Signed-off-by: lhy1024 <admin@liudos.us>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
---
 tests/cluster.go                                     |  8 +++++---
 .../mcs/keyspace/tso_keyspace_group_test.go          | 12 +++++++++---
 tests/pdctl/keyspace/keyspace_group_test.go          |  3 ---
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/tests/cluster.go b/tests/cluster.go
index 7fa029487fe..fc949a7f0be 100644
--- a/tests/cluster.go
+++ b/tests/cluster.go
@@ -62,6 +62,8 @@ var (
 	WaitLeaderReturnDelay = 20 * time.Millisecond
 	// WaitLeaderCheckInterval represents the time interval of WaitLeader running check.
 	WaitLeaderCheckInterval = 500 * time.Millisecond
+	// WaitLeaderRetryTimes represents the maximum number of loops of WaitLeader.
+	WaitLeaderRetryTimes = 100
 )
 
 // TestServer is only for test.
@@ -440,7 +442,7 @@ func (s *TestServer) BootstrapCluster() error {
 // make a test know the PD leader has been elected as soon as possible.
 // If it exceeds the maximum number of loops, it will return nil.
 func (s *TestServer) WaitLeader() bool {
-	for i := 0; i < 100; i++ {
+	for i := 0; i < WaitLeaderRetryTimes; i++ {
 		if s.server.GetMember().IsLeader() {
 			return true
 		}
@@ -649,7 +651,7 @@ func (c *TestCluster) GetFollower() string {
 // If it exceeds the maximum number of loops, it will return an empty string.
 func (c *TestCluster) WaitLeader(ops ...WaitOption) string {
 	option := &WaitOp{
-		retryTimes:   100,
+		retryTimes:   WaitLeaderRetryTimes,
 		waitInterval: WaitLeaderCheckInterval,
 	}
 	for _, op := range ops {
@@ -716,7 +718,7 @@ func (c *TestCluster) ResignLeader() error {
 // If it exceeds the maximum number of loops, it will return an empty string.
 func (c *TestCluster) WaitAllocatorLeader(dcLocation string, ops ...WaitOption) string {
 	option := &WaitOp{
-		retryTimes:   100,
+		retryTimes:   WaitLeaderRetryTimes,
 		waitInterval: WaitLeaderCheckInterval,
 	}
 	for _, op := range ops {
diff --git a/tests/integrations/mcs/keyspace/tso_keyspace_group_test.go b/tests/integrations/mcs/keyspace/tso_keyspace_group_test.go
index a284f8a44ac..41bcba0e90b 100644
--- a/tests/integrations/mcs/keyspace/tso_keyspace_group_test.go
+++ b/tests/integrations/mcs/keyspace/tso_keyspace_group_test.go
@@ -24,6 +24,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/stretchr/testify/suite"
 	bs "github.com/tikv/pd/pkg/basicserver"
 	"github.com/tikv/pd/pkg/mcs/utils"
@@ -54,6 +55,7 @@ func TestKeyspaceGroupTestSuite(t *testing.T) {
 }
 
 func (suite *keyspaceGroupTestSuite) SetupTest() {
+	suite.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes", `return(true)`))
 	ctx, cancel := context.WithCancel(context.Background())
 	suite.ctx = ctx
 	cluster, err := tests.NewTestAPICluster(suite.ctx, 1)
@@ -77,6 +79,7 @@ func (suite *keyspaceGroupTestSuite) SetupTest() {
 func (suite *keyspaceGroupTestSuite) TearDownTest() {
 	suite.cleanupFunc()
 	suite.cluster.Destroy()
+	suite.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes"))
 }
 
 func (suite *keyspaceGroupTestSuite) TestAllocNodesUpdate() {
@@ -292,9 +295,12 @@ func (suite *keyspaceGroupTestSuite) TestDefaultKeyspaceGroup() {
 	mcs.WaitForPrimaryServing(suite.Require(), nodes)
 
 	// the default keyspace group is exist.
-	time.Sleep(2 * time.Second)
-	kg, code := suite.tryGetKeyspaceGroup(utils.DefaultKeyspaceGroupID)
-	suite.Equal(http.StatusOK, code)
+	var kg *endpoint.KeyspaceGroup
+	var code int
+	testutil.Eventually(suite.Require(), func() bool {
+		kg, code = suite.tryGetKeyspaceGroup(utils.DefaultKeyspaceGroupID)
+		return code == http.StatusOK && kg != nil
+	}, testutil.WithWaitFor(time.Second*1))
 	suite.Equal(utils.DefaultKeyspaceGroupID, kg.ID)
 	suite.Len(kg.Members, utils.DefaultKeyspaceGroupReplicaCount)
 	for _, member := range kg.Members {
diff --git a/tests/pdctl/keyspace/keyspace_group_test.go b/tests/pdctl/keyspace/keyspace_group_test.go
index 927ce803d30..0c44700f48e 100644
--- a/tests/pdctl/keyspace/keyspace_group_test.go
+++ b/tests/pdctl/keyspace/keyspace_group_test.go
@@ -21,7 +21,6 @@ import (
 	"strconv"
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/pingcap/failpoint"
 	"github.com/stretchr/testify/require"
@@ -117,7 +116,6 @@ func TestSplitKeyspaceGroup(t *testing.T) {
 	re.NoError(err)
 	cmd := pdctlCmd.GetRootCmd()
 
-	time.Sleep(2 * time.Second)
 	tc.WaitLeader()
 	leaderServer := tc.GetServer(tc.GetLeader())
 	re.NoError(leaderServer.BootstrapCluster())
@@ -157,7 +155,6 @@ func TestExternalAllocNodeWhenStart(t *testing.T) {
 
 	cmd := pdctlCmd.GetRootCmd()
 
-	time.Sleep(2 * time.Second)
 	tc.WaitLeader()
 	leaderServer := tc.GetServer(tc.GetLeader())
 	re.NoError(leaderServer.BootstrapCluster())