1 files changed, 118 insertions, 18 deletions
diff --git a/whisper/whisper_test.go b/whisper/whisper_test.go
index b29e34a5e..27c57eee1 100644
--- a/whisper/whisper_test.go
+++ b/whisper/whisper_test.go
@@ -4,35 +4,135 @@ 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"
 )
 
-func TestEvent(t *testing.T) {
-	res := make(chan *Message, 1)
-	whisper := New()
-	id := whisper.NewIdentity()
-	whisper.Watch(Filter{
-		To: &id.PublicKey,
+type testNode struct {
+	server *p2p.Server
+	client *Whisper
+}
+
+func startNodes(n int) ([]*testNode, error) {
+	cluster := make([]*testNode, 0, 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,
+		})
+	}
+	return cluster, nil
+}
+
+func stopNodes(cluster []*testNode) {
+	for _, node := range cluster {
+		node.server.Stop()
+	}
+}
+
+func TestSelfMessage(t *testing.T) {
+	cluster, err := startNodes(1)
+	if err != nil {
+		t.Fatalf("failed to boot test cluster: %v", err)
+	}
+	defer stopNodes(cluster)
+
+	client := cluster[0].client
+
+	// Start watching for self messages, signal any arrivals
+	self := client.NewIdentity()
+	done := make(chan struct{})
+
+	client.Watch(Filter{
+		To: &self.PublicKey,
 		Fn: func(msg *Message) {
-			res <- msg
+			close(done)
 		},
 	})
-
-	msg := NewMessage([]byte(fmt.Sprintf("Hello world. This is whisper-go. Incase you're wondering; the time is %v", time.Now())))
+	// Send a dummy message to oneself
+	msg := NewMessage([]byte("hello whisper"))
 	envelope, err := msg.Wrap(DefaultProofOfWork, Options{
+		From: self,
+		To:   &self.PublicKey,
 		TTL:  DefaultTimeToLive,
-		From: id,
-		To:   &id.PublicKey,
 	})
 	if err != nil {
-		fmt.Println(err)
-		t.FailNow()
+		t.Fatalf("failed to wrap message: %v", err)
+	}
+	// Dump the message into the system and wait for it to pop back out
+	if err := client.Send(envelope); err != nil {
+		t.Fatalf("failed to send self-message: %v", err)
 	}
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatalf("self-message receive timeout")
+	}
+}
 
-	tick := time.NewTicker(time.Second)
-	whisper.postEvent(envelope)
+func TestDirectMessage(t *testing.T) {
+	cluster, err := startNodes(2)
+	if err != nil {
+		t.Fatalf("failed to boot test cluster: %v", err)
+	}
+	defer stopNodes(cluster)
+
+	sender := cluster[0].client
+	senderId := sender.NewIdentity()
+
+	recipient := cluster[1].client
+	recipientId := recipient.NewIdentity()
+
+	// Watch for arriving messages on the recipient
+	done := make(chan struct{})
+	recipient.Watch(Filter{
+		To: &recipientId.PublicKey,
+		Fn: func(msg *Message) {
+			close(done)
+		},
+	})
+	// Send a dummy message from the sender
+	msg := NewMessage([]byte("hello whisper"))
+	envelope, err := msg.Wrap(DefaultProofOfWork, Options{
+		From: senderId,
+		To:   &recipientId.PublicKey,
+		TTL:  DefaultTimeToLive,
+	})
+	if err != nil {
+		t.Fatalf("failed to wrap message: %v", err)
+	}
+	if err := sender.Send(envelope); err != nil {
+		t.Fatalf("failed to send  direct message: %v", err)
+	}
+	// Wait for an arrival or a timeout
 	select {
-	case <-res:
-	case <-tick.C:
-		t.Error("did not receive message")
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatalf("direct message receive timeout")
 	}
 }