Extend mail response with a peer and refactor tests

whisperv6/mailserver_response.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/p2p/enode"
 )
 
 const (
@@ -40,22 +41,22 @@ func CreateMailServerRequestFailedPayload(requestID common.Hash, err error) []by
 // * request completed successfully
 // * request failed
 // If the payload is unknown/unparseable, it returns `nil`
-func CreateMailServerEvent(payload []byte) (*EnvelopeEvent, error) {
+func CreateMailServerEvent(nodeID enode.ID, payload []byte) (*EnvelopeEvent, error) {
 
 	if len(payload) < common.HashLength {
 		return nil, invalidResponseSizeError(len(payload))
 	}
 
-	event, err := tryCreateMailServerRequestFailedEvent(payload)
+	event, err := tryCreateMailServerRequestFailedEvent(nodeID, payload)
 
 	if err != nil || event != nil {
 		return event, err
 	}
 
-	return tryCreateMailServerRequestCompletedEvent(payload)
+	return tryCreateMailServerRequestCompletedEvent(nodeID, payload)
 }
 
-func tryCreateMailServerRequestFailedEvent(payload []byte) (*EnvelopeEvent, error) {
+func tryCreateMailServerRequestFailedEvent(nodeID enode.ID, payload []byte) (*EnvelopeEvent, error) {
 	if len(payload) < common.HashLength+len(mailServerFailedPayloadPrefix) {
 		return nil, nil
 	}
@@ -75,6 +76,7 @@ func tryCreateMailServerRequestFailedEvent(payload []byte) (*EnvelopeEvent, erro
 	errorMsg = string(remainder)
 
 	event := EnvelopeEvent{
+		Peer:  nodeID,
 		Hash:  requestID,
 		Event: EventMailServerRequestCompleted,
 		Data: &MailServerResponse{
@@ -86,7 +88,7 @@ func tryCreateMailServerRequestFailedEvent(payload []byte) (*EnvelopeEvent, erro
 
 }
 
-func tryCreateMailServerRequestCompletedEvent(payload []byte) (*EnvelopeEvent, error) {
+func tryCreateMailServerRequestCompletedEvent(nodeID enode.ID, payload []byte) (*EnvelopeEvent, error) {
 	// check if payload is
 	// - requestID or
 	// - requestID + lastEnvelopeHash or
@@ -115,6 +117,7 @@ func tryCreateMailServerRequestCompletedEvent(payload []byte) (*EnvelopeEvent, e
 	}
 
 	event := EnvelopeEvent{
+		Peer:  nodeID,
 		Hash:  requestID,
 		Event: EventMailServerRequestCompleted,
 		Data: &MailServerResponse{

whisperv6/mailserver_response_test.go

@@ -1,100 +1,56 @@
 package whisperv6
 
 import (
-	"bytes"
 	"encoding/binary"
-	"strings"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/stretchr/testify/require"
 	"github.com/syndtr/goleveldb/leveldb/errors"
 )
 
 func checkValidErrorPayload(t *testing.T, id []byte, errorMsg string) {
 	requestID := common.BytesToHash(id)
-
 	errPayload := CreateMailServerRequestFailedPayload(requestID, errors.New(errorMsg))
+	nid := enode.ID{1}
+	event, err := CreateMailServerEvent(nid, errPayload)
 
-	event, err := CreateMailServerEvent(errPayload)
+	require.NoError(t, err)
-
+	require.NotNil(t, event)
-	if err != nil {
+	require.Equal(t, nid, event.Peer)
-		t.Errorf("Unexpected error: %v", err)
+	require.Equal(t, requestID, event.Hash)
-		return
-	}
-
-	if event == nil {
-		t.Errorf("Could not parse payload: %v", errPayload)
-		return
-	}
-
-	if !bytes.Equal(event.Hash[:], requestID[:]) {
-		t.Errorf("Unexpected hash: %v, expected %v", event.Hash, requestID)
-		return
-	}
 
 	eventData, ok := event.Data.(*MailServerResponse)
 	if !ok {
-		t.Errorf("Unexpected data in event: %v, expected a MailServerResponse", event.Data)
+		require.FailNow(t, "Unexpected data in event: %v, expected a MailServerResponse", event.Data)
-		return
-	}
-
-	if strings.Compare(eventData.Error.Error(), errorMsg) != 0 {
-		t.Errorf("Unexpected error string: '%s', expected '%s'", eventData, errorMsg)
-		return
 	}
+	require.EqualError(t, eventData.Error, errorMsg)
 }
 
 func checkValidSuccessPayload(t *testing.T, id []byte, lastHash []byte, timestamp uint32, envHash []byte) {
 	requestID := common.BytesToHash(id)
 	lastEnvelopeHash := common.BytesToHash(lastHash)
-
 	timestampBytes := make([]byte, 4)
 	binary.LittleEndian.PutUint32(timestampBytes, timestamp)
-
 	envelopeHash := common.BytesToHash(envHash)
-
 	cursor := append(timestampBytes, envelopeHash[:]...)
-
 	successPayload := CreateMailServerRequestCompletedPayload(requestID, lastEnvelopeHash, cursor)
+	nid := enode.ID{1}
+	event, err := CreateMailServerEvent(nid, successPayload)
 
-	event, err := CreateMailServerEvent(successPayload)
+	require.NoError(t, err)
-
+	require.NotNil(t, event)
-	if err != nil {
+	require.Equal(t, nid, event.Peer)
-		t.Errorf("Unexpected error: %v", err)
+	require.Equal(t, requestID, event.Hash)
-		return
-	}
-
-	if event == nil {
-		t.Errorf("Could not parse payload: %v", successPayload)
-		return
-	}
-
-	if !bytes.Equal(event.Hash[:], requestID[:]) {
-		t.Errorf("Unexpected hash: %v, expected %v", event.Hash, requestID)
-		return
-	}
 
 	eventData, ok := event.Data.(*MailServerResponse)
 	if !ok {
-		t.Errorf("Unexpected data in event: %v, expected a MailServerResponse", event.Data)
+		require.FailNow(t, "Unexpected data in event: %v, expected a MailServerResponse", event.Data)
-		return
-	}
-
-	if !bytes.Equal(eventData.LastEnvelopeHash[:], lastEnvelopeHash[:]) {
-		t.Errorf("Unexpected LastEnvelopeHash: %v, expected %v",
-			eventData.LastEnvelopeHash, lastEnvelopeHash)
-		return
-	}
-
-	if !bytes.Equal(eventData.Cursor, cursor) {
-		t.Errorf("Unexpected cursor: %v, expected: %v", eventData.Cursor, cursor)
-		return
-	}
-
-	if eventData.Error != nil {
-		t.Errorf("Unexpected error: %v", eventData.Error)
-		return
 	}
+	require.Equal(t, lastEnvelopeHash, eventData.LastEnvelopeHash)
+	require.Equal(t, cursor, eventData.Cursor)
+	require.NoError(t, eventData.Error)
 }
 
 func TestCreateMailServerEvent(t *testing.T) {
@@ -113,18 +69,11 @@ func TestCreateMailServerEvent(t *testing.T) {
 	// invalid payloads
 
 	// too small
-	_, err := CreateMailServerEvent([]byte{0x00})
+	_, err := CreateMailServerEvent(enode.ID{}, []byte{0x00})
-	if err == nil {
+	require.Error(t, err)
-		t.Errorf("Expected an error, got nil")
-		return
-	}
 
 	// too big and not error payload
 	payloadTooBig := make([]byte, common.HashLength*2+cursorSize+100)
-	_, err = CreateMailServerEvent(payloadTooBig)
+	_, err = CreateMailServerEvent(enode.ID{}, payloadTooBig)
-	if err == nil {
+	require.Error(t, err)
-		t.Errorf("Expected an error, got nil")
-		return
-	}
-
 }

whisperv6/whisper.go

@@ -971,7 +971,7 @@ func (whisper *Whisper) runMessageLoop(p *Peer, rw p2p.MsgReadWriter) error {
 					return errors.New("invalid request response message")
 				}
 
-				event, err := CreateMailServerEvent(payload)
+				event, err := CreateMailServerEvent(p.peer.ID(), payload)
 
 				if err != nil {
 					log.Warn("error while parsing request complete code, peer will be disconnected", "peer", p.peer.ID(), "err", err)