raft: set leader id in stepFollower

Follower has already set its leader ID from
previous append messages from the leader, but
to be consistent,  this adds a line to set its
leader id from leader snapshot message.

raft/raft.go

@@ -827,6 +827,7 @@ func stepFollower(r *raft, m pb.Message) {
 		r.handleHeartbeat(m)
 	case pb.MsgSnap:
 		r.electionElapsed = 0
+		r.lead = m.From
 		r.handleSnapshot(m)
 	case pb.MsgVote:
 		if (r.Vote == None || r.Vote == m.From) && r.raftLog.isUpToDate(m.Index, m.LogTerm) {

raft/raft_test.go

@@ -1901,6 +1901,10 @@ func TestRestoreFromSnapMsg(t *testing.T) {
 	sm := newTestRaft(2, []uint64{1, 2}, 10, 1, NewMemoryStorage())
 	sm.Step(m)
 
+	if sm.lead != uint64(1) {
+		t.Errorf("sm.lead = %d, want 1", sm.lead)
+	}
+
 	// TODO(bdarnell): what should this test?
 }