*: let TestEvictLeaderScheduler run in two modes

pkg/schedule/schedulers/evict_leader.go

@@ -103,13 +103,6 @@
 	return false, errs.ErrScheduleConfigNotExist.FastGenByArgs()
 }
 
-func (conf *evictLeaderSchedulerConfig) removeStore(id uint64) {
-	conf.Lock()
-	defer conf.Unlock()
-	// if the store is not existed, no need to resume leader transfer
-	_, _ = conf.removeStoreLocked(id)
-}
-
 func (conf *evictLeaderSchedulerConfig) resetStoreLocked(id uint64, keyRange []core.KeyRange) {
 	if err := conf.cluster.PauseLeaderTransfer(id); err != nil {
 		log.Error("pause leader transfer failed", zap.Uint64("store-id", id), errs.ZapError(err))
@@ -182,6 +175,12 @@
 	return true, nil
 }
 
+func (conf *evictLeaderSchedulerConfig) resumeLeaderTransferIfExist(id uint64) {
+	conf.RLock()
+	defer conf.RUnlock()
+	conf.cluster.ResumeLeaderTransfer(id)
+}
+
 func (conf *evictLeaderSchedulerConfig) update(id uint64, newRanges []core.KeyRange, batch int) error {
 	conf.Lock()
 	defer conf.Unlock()
@@ -415,7 +414,7 @@
 	batchFloat, ok := input["batch"].(float64)
 	if ok {
 		if batchFloat < 1 || batchFloat > 10 {
-			handler.config.removeStore(id)
+			handler.config.resumeLeaderTransferIfExist(id)
 			handler.rd.JSON(w, http.StatusBadRequest, "batch is invalid, it should be in [1, 10]")
 			return
 		}
@@ -425,7 +424,7 @@
 	ranges, ok := (input["ranges"]).([]string)
 	if ok {
 		if !inputHasStoreID {
-			handler.config.removeStore(id)
+			handler.config.resumeLeaderTransferIfExist(id)
 			handler.rd.JSON(w, http.StatusInternalServerError, errs.ErrSchedulerConfig.FastGenByArgs("id"))
 			return
 		}
@@ -435,11 +434,12 @@
 
 	newRanges, err = getKeyRanges(ranges)
 	if err != nil {
-		handler.config.removeStore(id)
+		handler.config.resumeLeaderTransferIfExist(id)
 		handler.rd.JSON(w, http.StatusInternalServerError, err.Error())
 		return
 	}
 
+	// StoreIDWithRanges is only changed in update function.
 	err = handler.config.update(id, newRanges, batch)
 	if err != nil {
 		handler.rd.JSON(w, http.StatusInternalServerError, err.Error())

pkg/schedule/schedulers/grant_leader.go

@@ -271,7 +271,9 @@
 
 	err := handler.config.buildWithArgs(args)
 	if err != nil {
-		_, _ = handler.config.removeStore(id)
+		handler.config.Lock()
+		handler.config.cluster.ResumeLeaderTransfer(id)
+		handler.config.Unlock()
 		handler.rd.JSON(w, http.StatusBadRequest, err.Error())
 		return
 	}

tests/server/api/region_test.go

@@ -91,7 +91,7 @@ func (suite *regionTestSuite) TearDownTest() {
 			return true
 		})
 	}
-	suite.env.RunFuncInTwoModes(cleanFunc)
+	suite.env.RunTestBasedOnMode(cleanFunc)
 }
 
 func (suite *regionTestSuite) TestSplitRegions() {

tests/server/api/rule_test.go

@@ -74,7 +74,7 @@ func (suite *ruleTestSuite) TearDownTest() {
 		err = tu.CheckPostJSON(tests.TestDialClient, urlPrefix+"/pd/api/v1/config/placement-rule", data, tu.StatusOK(re))
 		re.NoError(err)
 	}
-	suite.env.RunFuncInTwoModes(cleanFunc)
+	suite.env.RunTestBasedOnMode(cleanFunc)
 }
 
 func (suite *ruleTestSuite) TestSet() {

tests/testutil.go

@@ -371,16 +371,6 @@ func (s *SchedulingTestEnvironment) RunTestInAPIMode(test func(*TestCluster)) {
 	test(s.clusters[APIMode])
 }
 
-// RunFuncInTwoModes is to run func in two modes.
-func (s *SchedulingTestEnvironment) RunFuncInTwoModes(f func(*TestCluster)) {
-	if c, ok := s.clusters[PDMode]; ok {
-		f(c)
-	}
-	if c, ok := s.clusters[APIMode]; ok {
-		f(c)
-	}
-}
-
 // Cleanup is to cleanup the environment.
 func (s *SchedulingTestEnvironment) Cleanup() {
 	for _, cluster := range s.clusters {

tools/pd-ctl/tests/config/config_test.go

@@ -96,7 +96,7 @@ func (suite *configTestSuite) TearDownTest() {
 		err = testutil.CheckPostJSON(testDialClient, urlPrefix+"/pd/api/v1/config/placement-rule", data, testutil.StatusOK(re))
 		re.NoError(err)
 	}
-	suite.env.RunFuncInTwoModes(cleanFunc)
+	suite.env.RunTestBasedOnMode(cleanFunc)
 	suite.env.Cleanup()
 }
 

tools/pd-ctl/tests/hot/hot_test.go

@@ -71,7 +71,7 @@ func (suite *hotTestSuite) TearDownTest() {
 		}
 		hotStat.HotCache.CleanCache()
 	}
-	suite.env.RunFuncInTwoModes(cleanFunc)
+	suite.env.RunTestBasedOnMode(cleanFunc)
 }
 
 func (suite *hotTestSuite) TestHot() {

tools/pd-ctl/tests/scheduler/scheduler_test.go

@@ -94,7 +94,7 @@ func (suite *schedulerTestSuite) TearDownTest() {
 			}
 		}
 	}
-	suite.env.RunFuncInTwoModes(cleanFunc)
+	suite.env.RunTestBasedOnMode(cleanFunc)
 	suite.env.Cleanup()
 }
 
@@ -816,8 +816,7 @@ func (suite *schedulerTestSuite) checkSchedulerDiagnostic(cluster *pdTests.TestC
 }
 
 func (suite *schedulerTestSuite) TestEvictLeaderScheduler() {
-	// FIXME: API mode may have the problem
-	suite.env.RunTestInPDMode(suite.checkEvictLeaderScheduler)
+	suite.env.RunTestBasedOnMode(suite.checkEvictLeaderScheduler)
 }
 
 func (suite *schedulerTestSuite) checkEvictLeaderScheduler(cluster *pdTests.TestCluster) {
@@ -864,9 +863,14 @@ func (suite *schedulerTestSuite) checkEvictLeaderScheduler(cluster *pdTests.Test
 	output, err = tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "add", "evict-leader-scheduler", "1"}...)
 	re.NoError(err)
 	re.Contains(string(output), "Success!")
-	output, err = tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "remove", "evict-leader-scheduler-1"}...)
-	re.NoError(err)
-	re.Contains(string(output), "Success!")
+	testutil.Eventually(re, func() bool {
+		output, err = tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "remove", "evict-leader-scheduler-1"}...)
+		return err == nil && strings.Contains(string(output), "Success!")
+	})
+	testutil.Eventually(re, func() bool {
+		output, err = tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "show"}...)
+		return err == nil && !strings.Contains(string(output), "evict-leader-scheduler")
+	})
 }
 
 func mustExec(re *require.Assertions, cmd *cobra.Command, args []string, v any) string {