From 9d445d2fcf36f01971433bf9a1a70e41ce4c0e4a Mon Sep 17 00:00:00 2001
From: Yicheng Qin <qycqycqycqycqyc@gmail.com>
Date: Thu, 26 Feb 2015 00:13:35 -0800
Subject: [PATCH] rafthttp: add transport tests

---
 rafthttp/transport.go      | 10 +++----
 rafthttp/transport_test.go | 58 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/rafthttp/transport.go b/rafthttp/transport.go
index 5d931b2c9..1e2c36263 100644
--- a/rafthttp/transport.go
+++ b/rafthttp/transport.go
@@ -52,8 +52,8 @@ type transport struct {
 	serverStats  *stats.ServerStats
 	leaderStats  *stats.LeaderStats
 
-	mu     sync.RWMutex       // protect the peer map
-	peers  map[types.ID]*peer // remote peers
+	mu     sync.RWMutex      // protect the peer map
+	peers  map[types.ID]Peer // remote peers
 	errorc chan error
 }
 
@@ -65,7 +65,7 @@ func NewTransporter(rt http.RoundTripper, id, cid types.ID, r Raft, errorc chan
 		raft:         r,
 		serverStats:  ss,
 		leaderStats:  ls,
-		peers:        make(map[types.ID]*peer),
+		peers:        make(map[types.ID]Peer),
 		errorc:       errorc,
 	}
 }
@@ -181,12 +181,12 @@ type Pausable interface {
 // for testing
 func (t *transport) Pause() {
 	for _, p := range t.peers {
-		p.Pause()
+		p.(Pausable).Pause()
 	}
 }
 
 func (t *transport) Resume() {
 	for _, p := range t.peers {
-		p.Resume()
+		p.(Pausable).Resume()
 	}
 }
diff --git a/rafthttp/transport_test.go b/rafthttp/transport_test.go
index 9eb495cf8..87132f4ea 100644
--- a/rafthttp/transport_test.go
+++ b/rafthttp/transport_test.go
@@ -16,6 +16,7 @@ package rafthttp
 
 import (
 	"net/http"
+	"reflect"
 	"testing"
 	"time"
 
@@ -25,12 +26,51 @@ import (
 	"github.com/coreos/etcd/raft/raftpb"
 )
 
+// TestTransportSend tests that transport can send messages using correct
+// underlying peer, and drop local or unknown-target messages.
+func TestTransportSend(t *testing.T) {
+	ss := &stats.ServerStats{}
+	ss.Initialize()
+	peer1 := newFakePeer()
+	peer2 := newFakePeer()
+	tr := &transport{
+		serverStats: ss,
+		peers:       map[types.ID]Peer{types.ID(1): peer1, types.ID(2): peer2},
+	}
+	wmsgsIgnored := []raftpb.Message{
+		// bad local message
+		{Type: raftpb.MsgBeat},
+		// bad remote message
+		{Type: raftpb.MsgProp, To: 3},
+	}
+	wmsgsTo1 := []raftpb.Message{
+		// good message
+		{Type: raftpb.MsgProp, To: 1},
+		{Type: raftpb.MsgApp, To: 1},
+	}
+	wmsgsTo2 := []raftpb.Message{
+		// good message
+		{Type: raftpb.MsgProp, To: 2},
+		{Type: raftpb.MsgApp, To: 2},
+	}
+	tr.Send(wmsgsIgnored)
+	tr.Send(wmsgsTo1)
+	tr.Send(wmsgsTo2)
+
+	if !reflect.DeepEqual(peer1.msgs, wmsgsTo1) {
+		t.Errorf("msgs to peer 1 = %+v, want %+v", peer1.msgs, wmsgsTo1)
+	}
+	if !reflect.DeepEqual(peer2.msgs, wmsgsTo2) {
+		t.Errorf("msgs to peer 2 = %+v, want %+v", peer2.msgs, wmsgsTo2)
+	}
+}
+
 func TestTransportAdd(t *testing.T) {
 	ls := stats.NewLeaderStats("")
 	tr := &transport{
 		roundTripper: &roundTripperRecorder{},
 		leaderStats:  ls,
-		peers:        make(map[types.ID]*peer),
+		peers:        make(map[types.ID]Peer),
 	}
 	tr.AddPeer(1, []string{"http://a"})
 	defer tr.Stop()
@@ -55,7 +95,7 @@ func TestTransportRemove(t *testing.T) {
 	tr := &transport{
 		roundTripper: &roundTripperRecorder{},
 		leaderStats:  stats.NewLeaderStats(""),
-		peers:        make(map[types.ID]*peer),
+		peers:        make(map[types.ID]Peer),
 	}
 	tr.AddPeer(1, []string{"http://a"})
 	tr.RemovePeer(types.ID(1))
@@ -66,12 +106,24 @@ func TestTransportRemove(t *testing.T) {
 	}
 }
 
+func TestTransportUpdate(t *testing.T) {
+	peer := newFakePeer()
+	tr := &transport{
+		peers: map[types.ID]Peer{types.ID(1): peer},
+	}
+	u := "http://localhost:7001"
+	tr.UpdatePeer(types.ID(1), []string{u})
+	if w := "http://localhost:7001/raft"; peer.u != w {
+		t.Errorf("url = %s, want %s", peer.u, w)
+	}
+}
+
 func TestTransportErrorc(t *testing.T) {
 	errorc := make(chan error, 1)
 	tr := &transport{
 		roundTripper: newRespRoundTripper(http.StatusForbidden, nil),
 		leaderStats:  stats.NewLeaderStats(""),
-		peers:        make(map[types.ID]*peer),
+		peers:        make(map[types.ID]Peer),
 		errorc:       errorc,
 	}
 	tr.AddPeer(1, []string{"http://a"})