raft: Separate test methods for vote and pre-vote tests

raft/raft_test.go

@@ -288,155 +288,173 @@ func TestProgressPaused(t *testing.T) {
 }
 
 func TestLeaderElection(t *testing.T) {
-	for i, preVote := range []bool{false, true} {
-		var cfg func(*Config)
-		if preVote {
-			cfg = preVoteConfig
-		}
-		tests := []struct {
-			*network
-			state   StateType
-			expTerm uint64
-		}{
-			{newNetworkWithConfig(cfg, nil, nil, nil), StateLeader, 1},
-			{newNetworkWithConfig(cfg, nil, nil, nopStepper), StateLeader, 1},
-			{newNetworkWithConfig(cfg, nil, nopStepper, nopStepper), StateCandidate, 1},
-			{newNetworkWithConfig(cfg, nil, nopStepper, nopStepper, nil), StateCandidate, 1},
-			{newNetworkWithConfig(cfg, nil, nopStepper, nopStepper, nil, nil), StateLeader, 1},
-
-			// three logs further along than 0, but in the same term so rejections
-			// are returned instead of the votes being ignored.
-			{newNetworkWithConfig(cfg, nil, ents(1), ents(1), ents(1, 1), nil), StateFollower, 1},
-
-			// logs converge
-			{newNetworkWithConfig(cfg, ents(1), nil, ents(2), ents(1), nil), StateLeader, 2},
-		}
-
-		for j, tt := range tests {
-			tt.send(pb.Message{From: 1, To: 1, Type: pb.MsgHup})
-			sm := tt.network.peers[1].(*raft)
-			var expState StateType
-			var expTerm uint64
-			if tt.state == StateCandidate && preVote {
-				// In pre-vote mode, an election that fails to complete
-				// leaves the node in pre-candidate state without advancing
-				// the term.
-				expState = StatePreCandidate
-				expTerm = 0
-			} else {
-				expState = tt.state
-				expTerm = tt.expTerm
-			}
-			if sm.state != expState {
-				t.Errorf("#%d.%d: state = %s, want %s", i, j, sm.state, expState)
-			}
-			if g := sm.Term; g != expTerm {
-				t.Errorf("#%d.%d: term = %d, want %d", i, j, g, expTerm)
-			}
+	testLeaderElection(t, false)
+}
+
+func TestLeaderElectionPreVote(t *testing.T) {
+	testLeaderElection(t, true)
+}
+
+func testLeaderElection(t *testing.T, preVote bool) {
+	var cfg func(*Config)
+	if preVote {
+		cfg = preVoteConfig
+	}
+	tests := []struct {
+		*network
+		state   StateType
+		expTerm uint64
+	}{
+		{newNetworkWithConfig(cfg, nil, nil, nil), StateLeader, 1},
+		{newNetworkWithConfig(cfg, nil, nil, nopStepper), StateLeader, 1},
+		{newNetworkWithConfig(cfg, nil, nopStepper, nopStepper), StateCandidate, 1},
+		{newNetworkWithConfig(cfg, nil, nopStepper, nopStepper, nil), StateCandidate, 1},
+		{newNetworkWithConfig(cfg, nil, nopStepper, nopStepper, nil, nil), StateLeader, 1},
+
+		// three logs further along than 0, but in the same term so rejections
+		// are returned instead of the votes being ignored.
+		{newNetworkWithConfig(cfg, nil, ents(1), ents(1), ents(1, 1), nil), StateFollower, 1},
+
+		// logs converge
+		{newNetworkWithConfig(cfg, ents(1), nil, ents(2), ents(1), nil), StateLeader, 2},
+	}
+
+	for i, tt := range tests {
+		tt.send(pb.Message{From: 1, To: 1, Type: pb.MsgHup})
+		sm := tt.network.peers[1].(*raft)
+		var expState StateType
+		var expTerm uint64
+		if tt.state == StateCandidate && preVote {
+			// In pre-vote mode, an election that fails to complete
+			// leaves the node in pre-candidate state without advancing
+			// the term.
+			expState = StatePreCandidate
+			expTerm = 0
+		} else {
+			expState = tt.state
+			expTerm = tt.expTerm
+		}
+		if sm.state != expState {
+			t.Errorf("#%d: state = %s, want %s", i, sm.state, expState)
+		}
+		if g := sm.Term; g != expTerm {
+			t.Errorf("#%d: term = %d, want %d", i, g, expTerm)
 		}
 	}
 }
 
-// TestLeaderCycle verifies that each node in a cluster can campaign
+func TestLeaderCycle(t *testing.T) {
+	testLeaderCycle(t, false)
+}
+
+func TestLeaderCyclePreVote(t *testing.T) {
+	testLeaderCycle(t, true)
+}
+
+// testLeaderCycle verifies that each node in a cluster can campaign
 // and be elected in turn. This ensures that elections (including
 // pre-vote) work when not starting from a clean slate (as they do in
 // TestLeaderElection)
-func TestLeaderCycle(t *testing.T) {
-	for _, preVote := range []bool{false, true} {
-		var cfg func(*Config)
-		if preVote {
-			cfg = preVoteConfig
-		}
-		n := newNetworkWithConfig(cfg, nil, nil, nil)
-		for campaignerID := uint64(1); campaignerID <= 3; campaignerID++ {
-			n.send(pb.Message{From: campaignerID, To: campaignerID, Type: pb.MsgHup})
-
-			for _, peer := range n.peers {
-				sm := peer.(*raft)
-				if sm.id == campaignerID && sm.state != StateLeader {
-					t.Errorf("preVote=%v: campaigning node %d state = %v, want StateLeader",
-						preVote, sm.id, sm.state)
-				} else if sm.id != campaignerID && sm.state != StateFollower {
-					t.Errorf("preVote=%v: after campaign of node %d, "+
-						"node %d had state = %v, want StateFollower",
-						preVote, campaignerID, sm.id, sm.state)
-				}
+func testLeaderCycle(t *testing.T, preVote bool) {
+	var cfg func(*Config)
+	if preVote {
+		cfg = preVoteConfig
+	}
+	n := newNetworkWithConfig(cfg, nil, nil, nil)
+	for campaignerID := uint64(1); campaignerID <= 3; campaignerID++ {
+		n.send(pb.Message{From: campaignerID, To: campaignerID, Type: pb.MsgHup})
+
+		for _, peer := range n.peers {
+			sm := peer.(*raft)
+			if sm.id == campaignerID && sm.state != StateLeader {
+				t.Errorf("preVote=%v: campaigning node %d state = %v, want StateLeader",
+					preVote, sm.id, sm.state)
+			} else if sm.id != campaignerID && sm.state != StateFollower {
+				t.Errorf("preVote=%v: after campaign of node %d, "+
+					"node %d had state = %v, want StateFollower",
+					preVote, campaignerID, sm.id, sm.state)
 			}
 		}
 	}
 }
 
 func TestVoteFromAnyState(t *testing.T) {
-	for _, vt := range []pb.MessageType{pb.MsgVote, pb.MsgPreVote} {
-		for st := StateType(0); st < numStates; st++ {
-			r := newTestRaft(1, []uint64{1, 2, 3}, 10, 1, NewMemoryStorage())
-			r.Term = 1
+	testVoteFromAnyState(t, pb.MsgVote)
+}
 
-			switch st {
-			case StateFollower:
-				r.becomeFollower(r.Term, 3)
-			case StatePreCandidate:
-				r.becomePreCandidate()
-			case StateCandidate:
-				r.becomeCandidate()
-			case StateLeader:
-				r.becomeCandidate()
-				r.becomeLeader()
+func TestPreVoteFromAnyState(t *testing.T) {
+	testVoteFromAnyState(t, pb.MsgPreVote)
+}
+
+func testVoteFromAnyState(t *testing.T, vt pb.MessageType) {
+	for st := StateType(0); st < numStates; st++ {
+		r := newTestRaft(1, []uint64{1, 2, 3}, 10, 1, NewMemoryStorage())
+		r.Term = 1
+
+		switch st {
+		case StateFollower:
+			r.becomeFollower(r.Term, 3)
+		case StatePreCandidate:
+			r.becomePreCandidate()
+		case StateCandidate:
+			r.becomeCandidate()
+		case StateLeader:
+			r.becomeCandidate()
+			r.becomeLeader()
+		}
+
+		// Note that setting our state above may have advanced r.Term
+		// past its initial value.
+		origTerm := r.Term
+		newTerm := r.Term + 1
+
+		msg := pb.Message{
+			From:    2,
+			To:      1,
+			Type:    vt,
+			Term:    newTerm,
+			LogTerm: newTerm,
+			Index:   42,
+		}
+		if err := r.Step(msg); err != nil {
+			t.Errorf("%s,%s: Step failed: %s", vt, st, err)
+		}
+		if len(r.msgs) != 1 {
+			t.Errorf("%s,%s: %d response messages, want 1: %+v", vt, st, len(r.msgs), r.msgs)
+		} else {
+			resp := r.msgs[0]
+			if resp.Type != voteRespMsgType(vt) {
+				t.Errorf("%s,%s: response message is %s, want %s",
+					vt, st, resp.Type, voteRespMsgType(vt))
+			}
+			if resp.Reject {
+				t.Errorf("%s,%s: unexpected rejection", vt, st)
 			}
+		}
 
-			// Note that setting our state above may have advanced r.Term
-			// past its initial value.
-			origTerm := r.Term
-			newTerm := r.Term + 1
-
-			msg := pb.Message{
-				From:    2,
-				To:      1,
-				Type:    vt,
-				Term:    newTerm,
-				LogTerm: newTerm,
-				Index:   42,
+		// If this was a real vote, we reset our state and term.
+		if vt == pb.MsgVote {
+			if r.state != StateFollower {
+				t.Errorf("%s,%s: state %s, want %s", vt, StateFollower, r.state, st)
 			}
-			if err := r.Step(msg); err != nil {
-				t.Errorf("%s,%s: Step failed: %s", vt, st, err)
+			if r.Term != newTerm {
+				t.Errorf("%s,%s: term %d, want %d", vt, st, r.Term, newTerm)
 			}
-			if len(r.msgs) != 1 {
-				t.Errorf("%s,%s: %d response messages, want 1: %+v", vt, st, len(r.msgs), r.msgs)
-			} else {
-				resp := r.msgs[0]
-				if resp.Type != voteRespMsgType(vt) {
-					t.Errorf("%s,%s: response message is %s, want %s",
-						vt, st, resp.Type, voteRespMsgType(vt))
-				}
-				if resp.Reject {
-					t.Errorf("%s,%s: unexpected rejection", vt, st)
-				}
+			if r.Vote != 2 {
+				t.Errorf("%s,%s: vote %d, want 2", vt, st, r.Vote)
 			}
-
-			// If this was a real vote, we reset our state and term.
-			if vt == pb.MsgVote {
-				if r.state != StateFollower {
-					t.Errorf("%s,%s: state %s, want %s", vt, StateFollower, r.state, st)
-				}
-				if r.Term != newTerm {
-					t.Errorf("%s,%s: term %d, want %d", vt, st, r.Term, newTerm)
-				}
-				if r.Vote != 2 {
-					t.Errorf("%s,%s: vote %d, want 2", vt, st, r.Vote)
-				}
-			} else {
-				// In a prevote, nothing changes.
-				if r.state != st {
-					t.Errorf("%s,%s: state %s, want %s", vt, st, r.state, st)
-				}
-				if r.Term != origTerm {
-					t.Errorf("%s,%s: term %d, want %d", vt, st, r.Term, origTerm)
-				}
-				// if st == StateFollower or StatePreCandidate, r hasn't voted yet.
-				// In StateCandidate or StateLeader, it's voted for itself.
-				if r.Vote != None && r.Vote != 1 {
-					t.Errorf("%s,%s: vote %d, want %d or 1", vt, st, r.Vote, None)
-				}
+		} else {
+			// In a prevote, nothing changes.
+			if r.state != st {
+				t.Errorf("%s,%s: state %s, want %s", vt, st, r.state, st)
+			}
+			if r.Term != origTerm {
+				t.Errorf("%s,%s: term %d, want %d", vt, st, r.Term, origTerm)
+			}
+			// if st == StateFollower or StatePreCandidate, r hasn't voted yet.
+			// In StateCandidate or StateLeader, it's voted for itself.
+			if r.Vote != None && r.Vote != 1 {
+				t.Errorf("%s,%s: vote %d, want %d or 1", vt, st, r.Vote, None)
 			}
 		}
 	}
@@ -1225,71 +1243,77 @@ func TestMsgAppRespWaitReset(t *testing.T) {
 }
 
 func TestRecvMsgVote(t *testing.T) {
-	for i, msgType := range []pb.MessageType{pb.MsgVote, pb.MsgPreVote} {
-		tests := []struct {
-			state   StateType
-			i, term uint64
-			voteFor uint64
-			wreject bool
-		}{
-			{StateFollower, 0, 0, None, true},
-			{StateFollower, 0, 1, None, true},
-			{StateFollower, 0, 2, None, true},
-			{StateFollower, 0, 3, None, false},
-
-			{StateFollower, 1, 0, None, true},
-			{StateFollower, 1, 1, None, true},
-			{StateFollower, 1, 2, None, true},
-			{StateFollower, 1, 3, None, false},
-
-			{StateFollower, 2, 0, None, true},
-			{StateFollower, 2, 1, None, true},
-			{StateFollower, 2, 2, None, false},
-			{StateFollower, 2, 3, None, false},
-
-			{StateFollower, 3, 0, None, true},
-			{StateFollower, 3, 1, None, true},
-			{StateFollower, 3, 2, None, false},
-			{StateFollower, 3, 3, None, false},
-
-			{StateFollower, 3, 2, 2, false},
-			{StateFollower, 3, 2, 1, true},
-
-			{StateLeader, 3, 3, 1, true},
-			{StatePreCandidate, 3, 3, 1, true},
-			{StateCandidate, 3, 3, 1, true},
-		}
-
-		for j, tt := range tests {
-			sm := newTestRaft(1, []uint64{1}, 10, 1, NewMemoryStorage())
-			sm.state = tt.state
-			switch tt.state {
-			case StateFollower:
-				sm.step = stepFollower
-			case StateCandidate, StatePreCandidate:
-				sm.step = stepCandidate
-			case StateLeader:
-				sm.step = stepLeader
-			}
-			sm.Vote = tt.voteFor
-			sm.raftLog = &raftLog{
-				storage:  &MemoryStorage{ents: []pb.Entry{{}, {Index: 1, Term: 2}, {Index: 2, Term: 2}}},
-				unstable: unstable{offset: 3},
-			}
+	testRecvMsgVote(t, pb.MsgVote)
+}
+
+func testRecvMsgPreVote(t *testing.T) {
+	testRecvMsgVote(t, pb.MsgPreVote)
+}
 
-			sm.Step(pb.Message{Type: msgType, From: 2, Index: tt.i, LogTerm: tt.term})
+func testRecvMsgVote(t *testing.T, msgType pb.MessageType) {
+	tests := []struct {
+		state   StateType
+		i, term uint64
+		voteFor uint64
+		wreject bool
+	}{
+		{StateFollower, 0, 0, None, true},
+		{StateFollower, 0, 1, None, true},
+		{StateFollower, 0, 2, None, true},
+		{StateFollower, 0, 3, None, false},
 
-			msgs := sm.readMessages()
-			if g := len(msgs); g != 1 {
-				t.Fatalf("#%d.%d: len(msgs) = %d, want 1", i, j, g)
-				continue
-			}
-			if g := msgs[0].Type; g != voteRespMsgType(msgType) {
-				t.Errorf("#%d.%d, m.Type = %v, want %v", i, j, g, voteRespMsgType(msgType))
-			}
-			if g := msgs[0].Reject; g != tt.wreject {
-				t.Errorf("#%d.%d, m.Reject = %v, want %v", i, j, g, tt.wreject)
-			}
+		{StateFollower, 1, 0, None, true},
+		{StateFollower, 1, 1, None, true},
+		{StateFollower, 1, 2, None, true},
+		{StateFollower, 1, 3, None, false},
+
+		{StateFollower, 2, 0, None, true},
+		{StateFollower, 2, 1, None, true},
+		{StateFollower, 2, 2, None, false},
+		{StateFollower, 2, 3, None, false},
+
+		{StateFollower, 3, 0, None, true},
+		{StateFollower, 3, 1, None, true},
+		{StateFollower, 3, 2, None, false},
+		{StateFollower, 3, 3, None, false},
+
+		{StateFollower, 3, 2, 2, false},
+		{StateFollower, 3, 2, 1, true},
+
+		{StateLeader, 3, 3, 1, true},
+		{StatePreCandidate, 3, 3, 1, true},
+		{StateCandidate, 3, 3, 1, true},
+	}
+
+	for i, tt := range tests {
+		sm := newTestRaft(1, []uint64{1}, 10, 1, NewMemoryStorage())
+		sm.state = tt.state
+		switch tt.state {
+		case StateFollower:
+			sm.step = stepFollower
+		case StateCandidate, StatePreCandidate:
+			sm.step = stepCandidate
+		case StateLeader:
+			sm.step = stepLeader
+		}
+		sm.Vote = tt.voteFor
+		sm.raftLog = &raftLog{
+			storage:  &MemoryStorage{ents: []pb.Entry{{}, {Index: 1, Term: 2}, {Index: 2, Term: 2}}},
+			unstable: unstable{offset: 3},
+		}
+
+		sm.Step(pb.Message{Type: msgType, From: 2, Index: tt.i, LogTerm: tt.term})
+
+		msgs := sm.readMessages()
+		if g := len(msgs); g != 1 {
+			t.Fatalf("#%d: len(msgs) = %d, want 1", i, g)
+			continue
+		}
+		if g := msgs[0].Type; g != voteRespMsgType(msgType) {
+			t.Errorf("#%d, m.Type = %v, want %v", i, g, voteRespMsgType(msgType))
+		}
+		if g := msgs[0].Reject; g != tt.wreject {
+			t.Errorf("#%d, m.Reject = %v, want %v", i, g, tt.wreject)
 		}
 	}
 }
@@ -2396,27 +2420,33 @@ func TestRaftNodes(t *testing.T) {
 }
 
 func TestCampaignWhileLeader(t *testing.T) {
-	for i, preVote := range []bool{false, true} {
-		cfg := newTestConfig(1, []uint64{1}, 5, 1, NewMemoryStorage())
-		cfg.PreVote = preVote
-		r := newRaft(cfg)
-		if r.state != StateFollower {
-			t.Errorf("#%d: expected new node to be follower but got %s", i, r.state)
-		}
-		// We don't call campaign() directly because it comes after the check
-		// for our current state.
-		r.Step(pb.Message{From: 1, To: 1, Type: pb.MsgHup})
-		if r.state != StateLeader {
-			t.Errorf("#%d: expected single-node election to become leader but got %s", i, r.state)
-		}
-		term := r.Term
-		r.Step(pb.Message{From: 1, To: 1, Type: pb.MsgHup})
-		if r.state != StateLeader {
-			t.Errorf("#%d: expected to remain leader but got %s", i, r.state)
-		}
-		if r.Term != term {
-			t.Errorf("#%d: expected to remain in term %v but got %v", i, term, r.Term)
-		}
+	testCampaignWhileLeader(t, false)
+}
+
+func TestPreCampaignWhileLeader(t *testing.T) {
+	testCampaignWhileLeader(t, true)
+}
+
+func testCampaignWhileLeader(t *testing.T, preVote bool) {
+	cfg := newTestConfig(1, []uint64{1}, 5, 1, NewMemoryStorage())
+	cfg.PreVote = preVote
+	r := newRaft(cfg)
+	if r.state != StateFollower {
+		t.Errorf("expected new node to be follower but got %s", r.state)
+	}
+	// We don't call campaign() directly because it comes after the check
+	// for our current state.
+	r.Step(pb.Message{From: 1, To: 1, Type: pb.MsgHup})
+	if r.state != StateLeader {
+		t.Errorf("expected single-node election to become leader but got %s", r.state)
+	}
+	term := r.Term
+	r.Step(pb.Message{From: 1, To: 1, Type: pb.MsgHup})
+	if r.state != StateLeader {
+		t.Errorf("expected to remain leader but got %s", r.state)
+	}
+	if r.Term != term {
+		t.Errorf("expected to remain in term %v but got %v", term, r.Term)
 	}
 }