fix tests.

tikv · Mar 22, 2017 · 7b66ebd · 7b66ebd
1 parent d68489d
commit 7b66ebd
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 22 deletions.
diff --git a/server/balancer_test.go b/server/balancer_test.go
@@ -616,28 +616,49 @@ func (s *testReplicaCheckerSuite) TestDistinctScore2(c *C) {
 }
 
 func checkAddPeer(c *C, bop Operator, storeID uint64) {
-	op := bop.(*regionOperator).Ops[0].(*changePeerOperator)
+	var op *changePeerOperator
+	switch bop.(type) {
+	case *changePeerOperator:
+		op = bop.(*changePeerOperator)
+	case *regionOperator:
+		op = bop.(*regionOperator).Ops[0].(*changePeerOperator)
+	}
 	c.Assert(op.ChangePeer.GetChangeType(), Equals, raftpb.ConfChangeType_AddNode)
 	c.Assert(op.ChangePeer.GetPeer().GetStoreId(), Equals, storeID)
 }
 
 func checkRemovePeer(c *C, bop Operator, storeID uint64) {
-	op := bop.(*regionOperator).Ops[0].(*changePeerOperator)
+	var op *changePeerOperator
+	switch bop.(type) {
+	case *changePeerOperator:
+		op = bop.(*changePeerOperator)
+	case *regionOperator:
+		op = bop.(*regionOperator).Ops[0].(*changePeerOperator)
+	}
 	c.Assert(op.ChangePeer.GetChangeType(), Equals, raftpb.ConfChangeType_RemoveNode)
 	c.Assert(op.ChangePeer.GetPeer().GetStoreId(), Equals, storeID)
 }
 
 func checkTransferPeer(c *C, bop Operator, sourceID, targetID uint64) {
-	op := bop.(*regionOperator).Ops[0].(*changePeerOperator)
-	c.Assert(op.ChangePeer.GetChangeType(), Equals, raftpb.ConfChangeType_AddNode)
-	c.Assert(op.ChangePeer.GetPeer().GetStoreId(), Equals, targetID)
-	op = bop.(*regionOperator).Ops[1].(*changePeerOperator)
-	c.Assert(op.ChangePeer.GetChangeType(), Equals, raftpb.ConfChangeType_RemoveNode)
-	c.Assert(op.ChangePeer.GetPeer().GetStoreId(), Equals, sourceID)
+	op := bop.(*regionOperator)
+	if len(op.Ops) == 2 {
+		checkAddPeer(c, op.Ops[0], targetID)
+		checkRemovePeer(c, op.Ops[1], sourceID)
+	} else {
+		checkAddPeer(c, op.Ops[0], targetID)
+		checkTransferLeader(c, op.Ops[1], sourceID, targetID)
+		checkRemovePeer(c, op.Ops[2], sourceID)
+	}
 }
 
 func checkTransferLeader(c *C, bop Operator, sourceID, targetID uint64) {
-	op := bop.(*regionOperator).Ops[0].(*transferLeaderOperator)
+	var op *transferLeaderOperator
+	switch bop.(type) {
+	case *transferLeaderOperator:
+		op = bop.(*transferLeaderOperator)
+	case *regionOperator:
+		op = bop.(*regionOperator).Ops[0].(*transferLeaderOperator)
+	}
 	c.Assert(op.OldLeader.GetStoreId(), Equals, sourceID)
 	c.Assert(op.NewLeader.GetStoreId(), Equals, targetID)
 }
diff --git a/server/coordinator_test.go b/server/coordinator_test.go
@@ -31,7 +31,7 @@ type testOperator struct {
 func newTestOperator(regionID uint64, kind ResourceKind) Operator {
 	region := newRegionInfo(&metapb.Region{Id: regionID}, nil)
 	op := &testOperator{RegionID: regionID, Kind: kind}
-	return newRegionOperator(region, op)
+	return newRegionOperator(region, kind, op)
 }
 
 func (op *testOperator) GetRegionID() uint64           { return op.RegionID }

diff --git a/server/operator.go b/server/operator.go
@@ -76,14 +76,15 @@ func (op *adminOperator) Do(region *regionInfo) (*pdpb.RegionHeartbeatResponse,
 }
 
 type regionOperator struct {
-	Region *regionInfo `json:"region"`
-	Start  time.Time   `json:"start"`
-	End    time.Time   `json:"end"`
-	Index  int         `json:"index"`
-	Ops    []Operator  `json:"ops"`
+	Region *regionInfo  `json:"region"`
+	Start  time.Time    `json:"start"`
+	End    time.Time    `json:"end"`
+	Index  int          `json:"index"`
+	Ops    []Operator   `json:"ops"`
+	Kind   ResourceKind `json:"kind"`
 }
 
-func newRegionOperator(region *regionInfo, ops ...Operator) *regionOperator {
+func newRegionOperator(region *regionInfo, kind ResourceKind, ops ...Operator) *regionOperator {
 	// Do some check here, just fatal because it must be bug.
 	if len(ops) == 0 {
 		log.Fatalf("[region %d] new region operator with no ops", region.GetId())
@@ -93,6 +94,7 @@ func newRegionOperator(region *regionInfo, ops ...Operator) *regionOperator {
 		Region: region,
 		Start:  time.Now(),
 		Ops:    ops,
+		Kind:   kind,
 	}
 }
 
@@ -105,7 +107,7 @@ func (op *regionOperator) GetRegionID() uint64 {
 }
 
 func (op *regionOperator) GetResourceKind() ResourceKind {
-	return op.Ops[0].GetResourceKind()
+	return op.Kind
 }
 
 func (op *regionOperator) Do(region *regionInfo) (*pdpb.RegionHeartbeatResponse, bool) {

diff --git a/server/scheduler.go b/server/scheduler.go
@@ -238,27 +238,27 @@ func (s *shuffleRegionScheduler) Schedule(cluster *clusterInfo) Operator {
 
 func newAddPeer(region *regionInfo, peer *metapb.Peer) Operator {
 	addPeer := newAddPeerOperator(region.GetId(), peer)
-	return newRegionOperator(region, addPeer)
+	return newRegionOperator(region, regionKind, addPeer)
 }
 
 func newRemovePeer(region *regionInfo, peer *metapb.Peer) Operator {
 	removePeer := newRemovePeerOperator(region.GetId(), peer)
-	return newRegionOperator(region, removePeer)
+	return newRegionOperator(region, regionKind, removePeer)
 }
 
 func newTransferPeer(region *regionInfo, oldPeer, newPeer *metapb.Peer) Operator {
 	addPeer := newAddPeerOperator(region.GetId(), newPeer)
 	removePeer := newRemovePeerOperator(region.GetId(), oldPeer)
 	if region.Leader != nil && region.Leader.GetId() == oldPeer.GetId() {
 		transferLeader := newTransferLeaderOperator(region.GetId(), region.Leader, newPeer)
-		return newRegionOperator(region, addPeer, transferLeader, removePeer)
+		return newRegionOperator(region, regionKind, addPeer, transferLeader, removePeer)
 	}
-	return newRegionOperator(region, addPeer, removePeer)
+	return newRegionOperator(region, regionKind, addPeer, removePeer)
 }
 
 func newTransferLeader(region *regionInfo, newLeader *metapb.Peer) Operator {
 	transferLeader := newTransferLeaderOperator(region.GetId(), region.Leader, newLeader)
-	return newRegionOperator(region, transferLeader)
+	return newRegionOperator(region, leaderKind, transferLeader)
 }
 
 // scheduleAddPeer schedules a new peer.