whisper: global message expiration tests, polishes

3 files changed, 49 insertions, 16 deletions

diff --git a/whisper/common_test.go b/whisper/common_test.go
index 4e221d6b1..76e6463e6 100644
--- a/whisper/common_test.go
+++ b/whisper/common_test.go
@@ -8,9 +8,8 @@ import (
 	"io/ioutil"
 	"math/rand"
 
-	"github.com/ethereum/go-ethereum/p2p"
-
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/discover"
 )
 
diff --git a/whisper/peer_test.go b/whisper/peer_test.go
index c9092dd23..5f24504da 100644
--- a/whisper/peer_test.go
+++ b/whisper/peer_test.go
@@ -38,6 +38,20 @@ func startTestPeer() *testPeer {
 	}
 }
 
+func startTestPeerInited() (*testPeer, error) {
+	peer := startTestPeer()
+
+	if err := p2p.ExpectMsg(peer.stream, statusCode, []uint64{protocolVersion}); err != nil {
+		peer.stream.Close()
+		return nil, err
+	}
+	if err := p2p.SendItems(peer.stream, statusCode, protocolVersion); err != nil {
+		peer.stream.Close()
+		return nil, err
+	}
+	return peer, nil
+}
+
 func TestPeerStatusMessage(t *testing.T) {
 	tester := startTestPeer()
 
@@ -102,15 +116,12 @@ func TestPeerHandshakeSuccess(t *testing.T) {
 
 func TestPeerSend(t *testing.T) {
 	// Start a tester and execute the handshake
-	tester := startTestPeer()
+	tester, err := startTestPeerInited()
+	if err != nil {
+		t.Fatalf("failed to start initialized peer: %v", err)
+	}
 	defer tester.stream.Close()
 
-	if err := p2p.ExpectMsg(tester.stream, statusCode, []uint64{protocolVersion}); err != nil {
-		t.Fatalf("status message mismatch: %v", err)
-	}
-	if err := p2p.SendItems(tester.stream, statusCode, protocolVersion); err != nil {
-		t.Fatalf("failed to send status: %v", err)
-	}
 	// Construct a message and inject into the tester
 	message := NewMessage([]byte("peer broadcast test message"))
 	envelope, err := message.Wrap(DefaultPoW, Options{
@@ -138,15 +149,12 @@ func TestPeerSend(t *testing.T) {
 
 func TestPeerDeliver(t *testing.T) {
 	// Start a tester and execute the handshake
-	tester := startTestPeer()
+	tester, err := startTestPeerInited()
+	if err != nil {
+		t.Fatalf("failed to start initialized peer: %v", err)
+	}
 	defer tester.stream.Close()
 
-	if err := p2p.ExpectMsg(tester.stream, statusCode, []uint64{protocolVersion}); err != nil {
-		t.Fatalf("status message mismatch: %v", err)
-	}
-	if err := p2p.SendItems(tester.stream, statusCode, protocolVersion); err != nil {
-		t.Fatalf("failed to send status: %v", err)
-	}
 	// Watch for all inbound messages
 	arrived := make(chan struct{}, 1)
 	tester.client.Watch(Filter{
diff --git a/whisper/whisper_test.go b/whisper/whisper_test.go
index df27a945e..c072ba26b 100644
--- a/whisper/whisper_test.go
+++ b/whisper/whisper_test.go
@@ -156,3 +156,29 @@ func testBroadcast(anonymous bool, t *testing.T) {
 		}
 	}
 }
+
+func TestMessageExpiration(t *testing.T) {
+	// Start the single node cluster and inject a dummy message
+	node := startTestCluster(1)[0]
+
+	message := NewMessage([]byte("expiring message"))
+	envelope, err := message.Wrap(DefaultPoW, Options{
+		TTL: time.Second,
+	})
+	if err != nil {
+		t.Fatalf("failed to wrap message: %v", err)
+	}
+	if err := node.Send(envelope); err != nil {
+		t.Fatalf("failed to inject message: %v", err)
+	}
+	// Check that the message is inside the cache
+	if _, ok := node.messages[envelope.Hash()]; !ok {
+		t.Fatalf("message not found in cache")
+	}
+	// Wait for expiration and check cache again
+	time.Sleep(time.Second)     // wait for expiration
+	time.Sleep(expirationTicks) // wait for cleanup cycle
+	if _, ok := node.messages[envelope.Hash()]; ok {
+		t.Fatalf("message not expired from cache")
+	}
+}