itest+lntemp: refactor testCustomMessage

lntemp/rpc/lnd.go

@@ -648,3 +648,37 @@ func (h *HarnessRPC) SubscribeChannelEvents() ChannelEventsClient {
 
 	return client
 }
+
+type CustomMessageClient lnrpc.Lightning_SubscribeCustomMessagesClient
+
+// SubscribeCustomMessages creates a subscription client for custom messages.
+func (h *HarnessRPC) SubscribeCustomMessages() (CustomMessageClient,
+	context.CancelFunc) {
+
+	ctxt, cancel := context.WithCancel(h.runCtx)
+
+	req := &lnrpc.SubscribeCustomMessagesRequest{}
+
+	// SubscribeCustomMessages needs to have the context alive for the
+	// entire test case as the returned client will be used for send and
+	// receive events stream. Thus we use runCtx here instead of a timeout
+	// context.
+	stream, err := h.LN.SubscribeCustomMessages(ctxt, req)
+	h.NoError(err, "SubscribeCustomMessages")
+
+	return stream, cancel
+}
+
+// SendCustomMessage makes a RPC call to the node's SendCustomMessage and
+// returns the response.
+func (h *HarnessRPC) SendCustomMessage(
+	req *lnrpc.SendCustomMessageRequest) *lnrpc.SendCustomMessageResponse {
+
+	ctxt, cancel := context.WithTimeout(h.runCtx, DefaultTimeout)
+	defer cancel()
+
+	resp, err := h.LN.SendCustomMessage(ctxt, req)
+	h.NoError(err, "SendCustomMessage")
+
+	return resp
+}

lntest/itest/list_on_test.go

@@ -497,4 +497,8 @@ var allTestCasesTemp = []*lntemp.TestCase{
 		Name:     "open channel fee policy",
 		TestFunc: testOpenChannelUpdateFeePolicy,
 	},
+	{
+		Name:     "custom message",
+		TestFunc: testCustomMessage,
+	},
 }

lntest/itest/lnd_custom_message.go

@@ -1,13 +1,11 @@
 package itest
 
 import (
-	"context"
 	"fmt"
-	"sync"
 	"time"
 
 	"github.com/lightningnetwork/lnd/lnrpc"
-	"github.com/lightningnetwork/lnd/lntest"
+	"github.com/lightningnetwork/lnd/lntemp"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/stretchr/testify/require"
 )
@@ -15,16 +13,8 @@ import (
 // testCustomMessage tests sending and receiving of overridden custom message
 // types (within the message type range usually reserved for protocol messages)
 // via the send and subscribe custom message APIs.
-func testCustomMessage(net *lntest.NetworkHarness, t *harnessTest) {
-	ctx, cancel := context.WithCancel(context.Background())
-	var wg sync.WaitGroup
-
-	// At the end of our test, cancel our context and wait for all
-	// goroutines to exit.
-	defer func() {
-		cancel()
-		wg.Wait()
-	}()
+func testCustomMessage(ht *lntemp.HarnessTest) {
+	alice, bob := ht.Alice, ht.Bob
 
 	var (
 		overrideType1  uint32 = 554
@@ -34,28 +24,19 @@ func testCustomMessage(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// Update Alice to accept custom protocol messages with type 1 but do
 	// not allow Bob to handle them yet.
-	net.Alice.Cfg.ExtraArgs = append(
-		net.Alice.Cfg.ExtraArgs,
+	extraArgs := []string{
 		fmt.Sprintf(msgOverrideArg, overrideType1),
-	)
-	require.NoError(t.t, net.RestartNode(net.Alice, nil, nil))
-
-	// Wait for Alice's server to be active after the restart before we
-	// try to subscribe to our message stream.
-	require.NoError(t.t, net.Alice.WaitUntilServerActive())
+	}
+	ht.RestartNodeWithExtraArgs(alice, extraArgs)
 
 	// Subscribe Alice to custom messages before we send any, so that we
 	// don't miss any.
-	msgClient, err := net.Alice.LightningClient.SubscribeCustomMessages(
-		ctx, &lnrpc.SubscribeCustomMessagesRequest{},
-	)
-	require.NoError(t.t, err, "alice could not subscribe")
+	msgClient, cancel := alice.RPC.SubscribeCustomMessages()
+	defer cancel()
 
 	// Create a channel to receive custom messages on.
 	messages := make(chan *lnrpc.CustomMessage)
-	wg.Add(1)
 	go func() {
-		defer wg.Done()
 		for {
 			// If we fail to receive, just exit. The test should
 			// fail elsewhere if it doesn't get a message that it
@@ -69,120 +50,108 @@ func testCustomMessage(net *lntest.NetworkHarness, t *harnessTest) {
 			// test is shutting down.
 			select {
 			case messages <- msg:
-			case <-ctx.Done():
+			case <-ht.Context().Done():
 				return
 			}
 		}
 	}()
 
 	// Connect alice and bob so that they can exchange messages.
-	net.EnsureConnected(t.t, net.Alice, net.Bob)
+	ht.EnsureConnected(alice, bob)
 
 	// Create a custom message that is within our allowed range.
 	msgType := uint32(lnwire.CustomTypeStart + 1)
 	msgData := []byte{1, 2, 3}
 
 	// Send it from Bob to Alice.
-	ctxt, _ := context.WithTimeout(ctx, defaultTimeout)
-	_, err = net.Bob.LightningClient.SendCustomMessage(
-		ctxt, &lnrpc.SendCustomMessageRequest{
-			Peer: net.Alice.PubKey[:],
+	bobMsg := &lnrpc.SendCustomMessageRequest{
+		Peer: alice.PubKey[:],
 		Type: msgType,
 		Data: msgData,
-		},
-	)
-	require.NoError(t.t, err, "bob could not send")
+	}
+	bob.RPC.SendCustomMessage(bobMsg)
 
-	// Wait for Alice to receive the message. It should come through because
-	// it is within our allowed range.
+	// Wait for Alice to receive the message. It should come through
+	// because it is within our allowed range.
 	select {
 	case msg := <-messages:
-		// Check our type and data and (sanity) check the peer we got it
-		// from.
-		require.Equal(t.t, msgType, msg.Type, "first msg type wrong")
-		require.Equal(t.t, msgData, msg.Data, "first msg data wrong")
-		require.Equal(t.t, net.Bob.PubKey[:], msg.Peer, "first msg "+
+		// Check our type and data and (sanity) check the peer we got
+		// it from.
+		require.Equal(ht, msgType, msg.Type, "first msg type wrong")
+		require.Equal(ht, msgData, msg.Data, "first msg data wrong")
+		require.Equal(ht, bob.PubKey[:], msg.Peer, "first msg "+
 			"peer wrong")
 
 	case <-time.After(defaultTimeout):
-		t.t.Fatalf("alice did not receive first custom message: %v",
+		ht.Fatalf("alice did not receive first custom message: %v",
 			msgType)
 	}
 
 	// Try to send a message from Bob to Alice which has a message type
 	// outside of the custom type range and assert that it fails.
-	ctxt, _ = context.WithTimeout(ctx, defaultTimeout)
-	_, err = net.Bob.LightningClient.SendCustomMessage(
-		ctxt, &lnrpc.SendCustomMessageRequest{
-			Peer: net.Alice.PubKey[:],
+	bobMsg = &lnrpc.SendCustomMessageRequest{
+		Peer: alice.PubKey[:],
 		Type: overrideType1,
 		Data: msgData,
-		},
-	)
-	require.Error(t.t, err, "bob should not be able to send type 1")
+	}
+	_, err := bob.RPC.LN.SendCustomMessage(ht.Context(), bobMsg)
+	require.Error(ht, err, "bob should not be able to send type 1")
 
 	// Now, restart Bob with the ability to send two different custom
 	// protocol messages.
-	net.Bob.Cfg.ExtraArgs = append(
-		net.Bob.Cfg.ExtraArgs,
+	extraArgs = []string{
 		fmt.Sprintf(msgOverrideArg, overrideType1),
 		fmt.Sprintf(msgOverrideArg, overrideType2),
-	)
-	require.NoError(t.t, net.RestartNode(net.Bob, nil, nil))
+	}
+	ht.RestartNodeWithExtraArgs(bob, extraArgs)
 
 	// Make sure Bob and Alice are connected after his restart.
-	net.EnsureConnected(t.t, net.Alice, net.Bob)
+	ht.EnsureConnected(alice, bob)
 
 	// Send a message from Bob to Alice with a type that Bob is allowed to
 	// send, but Alice will not handle as a custom message.
-	ctxt, _ = context.WithTimeout(ctx, defaultTimeout)
-	_, err = net.Bob.LightningClient.SendCustomMessage(
-		ctxt, &lnrpc.SendCustomMessageRequest{
-			Peer: net.Alice.PubKey[:],
+	bobMsg = &lnrpc.SendCustomMessageRequest{
+		Peer: alice.PubKey[:],
 		Type: overrideType2,
 		Data: msgData,
-		},
-	)
-	require.NoError(t.t, err, "bob should be able to send type 2")
+	}
+	bob.RPC.SendCustomMessage(bobMsg)
 
 	// Do a quick check that Alice did not receive this message in her
-	// stream. Note that this is an instant check, so could miss the message
-	// being received. We'll also check below that she didn't get it, this
-	// is just a sanity check.
+	// stream. Note that this is an instant check, so could miss the
+	// message being received. We'll also check below that she didn't get
+	// it, this is just a sanity check.
 	select {
 	case msg := <-messages:
-		t.t.Fatalf("unexpected message: %v", msg)
+		ht.Fatalf("unexpected message: %v", msg)
 	default:
 	}
 
 	// Finally, send a custom message with a type that Bob is allowed to
 	// send and Alice is configured to receive.
-	ctxt, _ = context.WithTimeout(ctx, defaultTimeout)
-	_, err = net.Bob.LightningClient.SendCustomMessage(
-		ctxt, &lnrpc.SendCustomMessageRequest{
-			Peer: net.Alice.PubKey[:],
+	bobMsg = &lnrpc.SendCustomMessageRequest{
+		Peer: alice.PubKey[:],
 		Type: overrideType1,
 		Data: msgData,
-		},
-	)
-	require.NoError(t.t, err, "bob should be able to send type 1")
+	}
+	bob.RPC.SendCustomMessage(bobMsg)
 
 	// Wait to receive a message from Bob. This check serves to ensure that
 	// our message type 1 was delivered, and assert that the preceding one
-	// was not (we could have missed it in our check above). When we receive
-	// the second message, we know that the first one did not go through,
-	// because we expect our messages to deliver in order.
+	// was not (we could have missed it in our check above). When we
+	// receive the second message, we know that the first one did not go
+	// through, because we expect our messages to deliver in order.
 	select {
 	case msg := <-messages:
-		// Check our type and data and (sanity) check the peer we got it
-		// from.
-		require.Equal(t.t, overrideType1, msg.Type, "second message "+
+		// Check our type and data and (sanity) check the peer we got
+		// it from.
+		require.Equal(ht, overrideType1, msg.Type, "second message "+
 			"type")
-		require.Equal(t.t, msgData, msg.Data, "second message data")
-		require.Equal(t.t, net.Bob.PubKey[:], msg.Peer, "second "+
+		require.Equal(ht, msgData, msg.Data, "second message data")
+		require.Equal(ht, bob.PubKey[:], msg.Peer, "second "+
 			"message peer")
 
 	case <-time.After(defaultTimeout):
-		t.t.Fatalf("alice did not receive second custom message")
+		ht.Fatalf("alice did not receive second custom message")
 	}
 }

lntest/itest/lnd_test_list_on_test.go

@@ -8,8 +8,4 @@ var allTestCases = []*testCase{
 		name: "async bidirectional payments",
 		test: testBidirectionalAsyncPayments,
 	},
-	{
-		name: "custom messaging",
-		test: testCustomMessage,
-	},
 }