peer: Add startPeer test function

Signed-off-by: Ononiwu Maureen <maureen.ononiwu@outlook.com>

peer/brontide_test.go

@@ -1012,27 +1012,8 @@ func TestPeerCustomMessage(t *testing.T) {
 
 	remoteKey := alicePeer.PubKey()
 
-	// Set up the init sequence.
-	go func() {
-		// Read init message.
-		<-mockConn.writtenMessages
-
-		// Write the init reply message.
-		initReplyMsg := lnwire.NewInitMessage(
-			lnwire.NewRawFeatureVector(
-				lnwire.DataLossProtectRequired,
-			),
-			lnwire.NewRawFeatureVector(),
-		)
-		var b bytes.Buffer
-		_, err := lnwire.WriteMessage(&b, initReplyMsg, 0)
-		require.NoError(t, err)
-
-		mockConn.readMessages <- b.Bytes()
-	}()
-
-	// Start the peer.
-	require.NoError(t, alicePeer.Start())
+	// Start peer.
+	startPeer(t, mockConn, alicePeer)
 
 	// Send a custom message.
 	customMsg, err := lnwire.NewCustom(
@@ -1320,27 +1301,8 @@ func TestStartupWriteMessageRace(t *testing.T) {
 		close(sendPingDone)
 	}()
 
-	// Handle init messages.
-	go func() {
-		// Read init message.
-		<-mockConn.writtenMessages
-
-		// Write the init reply message.
-		initReplyMsg := lnwire.NewInitMessage(
-			lnwire.NewRawFeatureVector(
-				lnwire.DataLossProtectRequired,
-			),
-			lnwire.NewRawFeatureVector(),
-		)
-		var b bytes.Buffer
-		_, err = lnwire.WriteMessage(&b, initReplyMsg, 0)
-		require.NoError(t, err)
-
-		mockConn.readMessages <- b.Bytes()
-	}()
-
 	// Start the peer. No data race should occur.
-	require.NoError(t, peer.Start())
+	startPeer(t, mockConn, peer)
 
 	// Ensure messages were sent during startup.
 	<-sendPingDone

peer/test_utils.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	crand "crypto/rand"
 	"encoding/binary"
+	"errors"
 	"io"
 	"math/rand"
 	"net"
@@ -726,3 +727,79 @@ func createTestPeer(t *testing.T) *testPeerParams {
 		chanStatusMgr: chanStatusMgr,
 	}
 }
+
+func startPeer(t *testing.T, mockConn *mockMessageConn,
+	peer *Brontide, extraChecks ...func() error) {
+
+	errChan := make(chan error)
+	go func() {
+		select {
+		case rawMsg := <-mockConn.writtenMessages:
+			msgReader := bytes.NewReader(rawMsg)
+
+			nextMsg, err := lnwire.ReadMessage(msgReader, 0)
+			if err != nil {
+				errChan <- err
+
+				return
+			}
+
+			_, ok := nextMsg.(*lnwire.Init)
+			if !ok {
+				errChan <- errors.New("expected lnwire.Init " +
+					"type, got something else")
+
+				return
+			}
+
+		case <-time.After(timeout):
+			errChan <- errors.New("did not receive init message " +
+				"as expected")
+
+			return
+		}
+
+		// Write the init reply message.
+		initReplyMsg := lnwire.NewInitMessage(
+			lnwire.NewRawFeatureVector(
+				lnwire.DataLossProtectRequired,
+				lnwire.GossipQueriesOptional,
+			),
+			lnwire.NewRawFeatureVector(),
+		)
+
+		var b bytes.Buffer
+		_, err := lnwire.WriteMessage(&b, initReplyMsg, 0)
+		if err != nil {
+			errChan <- err
+
+			return
+		}
+
+		select {
+		case mockConn.readMessages <- b.Bytes():
+		case <-time.After(timeout):
+			errChan <- errors.New("could not reply init message " +
+				"as expected")
+
+			return
+		}
+
+		for _, f := range extraChecks {
+			err = f()
+
+			if err != nil {
+				errChan <- err
+
+				return
+			}
+		}
+
+		errChan <- nil
+	}()
+
+	require.NoError(t, peer.Start())
+
+	err := <-errChan
+	require.NoError(t, err)
+}