add more test for balance range

Signed-off-by: 童剑 <[email protected]>

Author: bufferflies

pkg/schedule/schedulers/balance_leader_test.go

@@ -35,6 +35,7 @@ import (
 	"github.com/tikv/pd/pkg/schedule/plan"
 	"github.com/tikv/pd/pkg/schedule/types"
 	"github.com/tikv/pd/pkg/storage"
+	"github.com/tikv/pd/pkg/utils/keyutil"
 	"github.com/tikv/pd/pkg/utils/operatorutil"
 	"github.com/tikv/pd/pkg/versioninfo"
 )
@@ -449,6 +450,13 @@ func (suite *balanceLeaderRangeSchedulerTestSuite) TestSingleRangeBalance() {
 	re.NoError(err)
 	ops, _ = lb.Schedule(suite.tc, false)
 	re.Empty(ops)
+
+	kye := keyutil.NewKeyRange("a", "g")
+	suite.tc.Append([]keyutil.KeyRange{kye})
+	lb, err = CreateScheduler(types.BalanceLeaderScheduler, suite.oc, storage.NewStorageWithMemoryBackend(), ConfigSliceDecoder(types.BalanceLeaderScheduler, []string{"", ""}))
+	re.NoError(err)
+	ops, _ = lb.Schedule(suite.tc, false)
+	re.Empty(ops)
 }
 
 func (suite *balanceLeaderRangeSchedulerTestSuite) TestMultiRangeBalance() {

pkg/schedule/schedulers/balance_range.go

@@ -394,12 +394,16 @@ func (s *balanceRangeScheduler) IsScheduleAllowed(cluster sche.SchedulerCluster)
 			if err := s.conf.begin(index); err != nil {
 				return false
 			}
+			km := cluster.GetKeyRangeManager()
+			km.Append(job.Ranges)
 		}
 		// todo: add other conditions such as the diff of the score between the source and target store.
 		if time.Since(*job.Start) > job.Timeout {
 			if err := s.conf.finish(index); err != nil {
 				return false
 			}
+			km := cluster.GetKeyRangeManager()
+			km.Delete(job.Ranges)
 			balanceRangeExpiredCounter.Inc()
 		}
 	}

pkg/schedule/schedulers/balance_range_test.go

@@ -61,6 +61,14 @@ func TestTIKVEngine(t *testing.T) {
 	scheduler, err := CreateScheduler(types.BalanceRangeScheduler, oc, storage.NewStorageWithMemoryBackend(),
 		ConfigSliceDecoder(types.BalanceRangeScheduler,
 			[]string{"leader-scatter", "tikv", "1h", "test", "100", "200"}))
+	re.True(scheduler.IsScheduleAllowed(tc))
+	km := tc.GetKeyRangeManager()
+	kr := keyutil.NewKeyRange("", "")
+	ranges := km.GetNonOverlappingKeyRanges(&kr)
+	re.Len(ranges, 2)
+	re.Equal(ranges[0], keyutil.NewKeyRange("", "100"))
+	re.Equal(ranges[1], keyutil.NewKeyRange("200", ""))
+
 	re.NoError(err)
 	ops, _ := scheduler.Schedule(tc, true)
 	re.Empty(ops)

pkg/schedule/schedulers/balance_region_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/tikv/pd/pkg/schedule/plan"
 	"github.com/tikv/pd/pkg/schedule/types"
 	"github.com/tikv/pd/pkg/storage"
+	"github.com/tikv/pd/pkg/utils/keyutil"
 	"github.com/tikv/pd/pkg/utils/operatorutil"
 	"github.com/tikv/pd/pkg/versioninfo"
 )
@@ -654,9 +655,16 @@ func TestBalanceRegionEmptyRegion(t *testing.T) {
 		core.SetApproximateKeys(1),
 	)
 	tc.PutRegion(region)
+
+	tc.Append([]keyutil.KeyRange{keyutil.NewKeyRange("a", "b")})
 	operators, _ := sb.Schedule(tc, false)
+	re.Empty(operators)
+	tc.Delete([]keyutil.KeyRange{keyutil.NewKeyRange("a", "b")})
+
+	operators, _ = sb.Schedule(tc, false)
 	re.NotEmpty(operators)
 
+	// reject the empty regions if the cluster has more regions.
 	for i := uint64(10); i < 111; i++ {
 		tc.PutRegionStores(i, 1, 3, 4)
 	}