whisper: mock tests to use simulated peers

2 files changed, 55 insertions, 71 deletions

diff --git a/whisper/common_test.go b/whisper/common_test.go
index 8c12f8aeb..4e221d6b1 100644
--- a/whisper/common_test.go
+++ b/whisper/common_test.go
@@ -3,7 +3,9 @@
 package whisper
 
 import (
+	"bytes"
 	"fmt"
+	"io/ioutil"
 	"math/rand"
 
 	"github.com/ethereum/go-ethereum/p2p"
@@ -36,3 +38,33 @@ func whisperCaps() []p2p.Cap {
 		},
 	}
 }
+
+// bufMsgPipe creates a buffered message pipe between two endpoints.
+func bufMsgPipe() (*p2p.MsgPipeRW, *p2p.MsgPipeRW) {
+	A, midA := p2p.MsgPipe()
+	midB, B := p2p.MsgPipe()
+
+	go copyMsgPipe(midA, midB)
+	go copyMsgPipe(midB, midA)
+
+	return A, B
+}
+
+// copyMsgPipe copies messages from the src pipe to the dest.
+func copyMsgPipe(dst, src *p2p.MsgPipeRW) {
+	defer dst.Close()
+	for {
+		msg, err := src.ReadMsg()
+		if err != nil {
+			return
+		}
+		data, err := ioutil.ReadAll(msg.Payload)
+		if err != nil {
+			return
+		}
+		msg.Payload = bytes.NewReader(data)
+		if err := dst.WriteMsg(msg); err != nil {
+			return
+		}
+	}
+}
diff --git a/whisper/whisper_test.go b/whisper/whisper_test.go
index 97644fd27..df27a945e 100644
--- a/whisper/whisper_test.go
+++ b/whisper/whisper_test.go
@@ -1,75 +1,36 @@
 package whisper
 
 import (
-	"fmt"
 	"testing"
 	"time"
 
-	"github.com/ethereum/go-ethereum/common"
-	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/p2p"
-	"github.com/ethereum/go-ethereum/p2p/nat"
 )
 
-type testNode struct {
-	server *p2p.Server
-	client *Whisper
-}
-
-func startNodes(n int) ([]*testNode, error) {
-	// Start up the cluster of nodes
-	cluster := make([]*testNode, 0, n)
+func startTestCluster(n int) []*Whisper {
+	// Create the batch of simulated peers
+	nodes := make([]*p2p.Peer, n)
 	for i := 0; i < n; i++ {
-		shh := New()
-
-		// Generate the node identity
-		key, err := crypto.GenerateKey()
-		if err != nil {
-			return nil, err
-		}
-		name := common.MakeName(fmt.Sprintf("whisper-go-test-%d", i), "1.0")
-
-		// Create an Ethereum server to communicate through
-		server := &p2p.Server{
-			PrivateKey: key,
-			MaxPeers:   10,
-			Name:       name,
-			Protocols:  []p2p.Protocol{shh.Protocol()},
-			ListenAddr: fmt.Sprintf(":%d", 30300+i),
-			NAT:        nat.Any(),
-		}
-		if err := server.Start(); err != nil {
-			return nil, err
-		}
-		// Peer online, store and iterate
-		cluster = append(cluster, &testNode{
-			server: server,
-			client: shh,
-		})
+		nodes[i] = p2p.NewPeer(randomNodeID(), randomNodeName(), whisperCaps())
 	}
-	// Manually wire together the cluster nodes
-	root := cluster[0].server.Self()
-	for _, node := range cluster[1:] {
-		node.server.SuggestPeer(root)
+	whispers := make([]*Whisper, n)
+	for i := 0; i < n; i++ {
+		whispers[i] = New()
+		whispers[i].Start()
 	}
-	return cluster, nil
-}
+	// Wire all the peers to the root one
+	for i := 1; i < n; i++ {
+		src, dst := bufMsgPipe()
 
-func stopNodes(cluster []*testNode) {
-	for _, node := range cluster {
-		node.server.Stop()
+		go whispers[0].handlePeer(nodes[i], src)
+		go whispers[i].handlePeer(nodes[0], dst)
 	}
+	return whispers
 }
 
 func TestSelfMessage(t *testing.T) {
 	// Start the single node cluster
-	cluster, err := startNodes(1)
-	if err != nil {
-		t.Fatalf("failed to boot test cluster: %v", err)
-	}
-	defer stopNodes(cluster)
-
-	client := cluster[0].client
+	client := startTestCluster(1)[0]
 
 	// Start watching for self messages, signal any arrivals
 	self := client.NewIdentity()
@@ -104,16 +65,12 @@ func TestSelfMessage(t *testing.T) {
 
 func TestDirectMessage(t *testing.T) {
 	// Start the sender-recipient cluster
-	cluster, err := startNodes(2)
-	if err != nil {
-		t.Fatalf("failed to boot test cluster: %v", err)
-	}
-	defer stopNodes(cluster)
+	cluster := startTestCluster(2)
 
-	sender := cluster[0].client
+	sender := cluster[0]
 	senderId := sender.NewIdentity()
 
-	recipient := cluster[1].client
+	recipient := cluster[1]
 	recipientId := recipient.NewIdentity()
 
 	// Watch for arriving messages on the recipient
@@ -155,18 +112,13 @@ func TestIdentifiedBroadcast(t *testing.T) {
 
 func testBroadcast(anonymous bool, t *testing.T) {
 	// Start the single sender multi recipient cluster
-	cluster, err := startNodes(3)
-	if err != nil {
-		t.Fatalf("failed to boot test cluster: %v", err)
-	}
-	defer stopNodes(cluster)
+	cluster := startTestCluster(3)
 
-	sender := cluster[0].client
-	targets := make([]*Whisper, len(cluster)-1)
-	for i, node := range cluster[1:] {
-		targets[i] = node.client
+	sender := cluster[1]
+	targets := cluster[1:]
+	for _, target := range targets {
 		if !anonymous {
-			targets[i].NewIdentity()
+			target.NewIdentity()
 		}
 	}
 	// Watch for arriving messages on the recipients