itest: manage context inside WaitForChannelOpen and WaitForChannelClose

lntest/harness.go

@@ -1123,9 +1123,13 @@ func (n *NetworkHarness) OpenPendingChannel(ctx context.Context,
 // consuming a message from the past open channel stream. If the passed context
 // has a timeout, then if the timeout is reached before the channel has been
 // opened, then an error is returned.
-func (n *NetworkHarness) WaitForChannelOpen(ctx context.Context,
+func (n *NetworkHarness) WaitForChannelOpen(
 	openChanStream lnrpc.Lightning_OpenChannelClient) (*lnrpc.ChannelPoint, error) {
 
+	ctxb := context.Background()
+	ctx, cancel := context.WithTimeout(ctxb, ChannelOpenTimeout)
+	defer cancel()
+
 	errChan := make(chan error)
 	respChan := make(chan *lnrpc.ChannelPoint)
 	go func() {
@@ -1294,9 +1298,13 @@ func (n *NetworkHarness) CloseChannel(ctx context.Context,
 // stream that the node has deemed the channel has been fully closed. If the
 // passed context has a timeout, then if the timeout is reached before the
 // notification is received then an error is returned.
-func (n *NetworkHarness) WaitForChannelClose(ctx context.Context,
+func (n *NetworkHarness) WaitForChannelClose(
 	closeChanStream lnrpc.Lightning_CloseChannelClient) (*chainhash.Hash, error) {
 
+	ctxb := context.Background()
+	ctx, cancel := context.WithTimeout(ctxb, ChannelCloseTimeout)
+	defer cancel()
+
 	errChan := make(chan error)
 	updateChan := make(chan *lnrpc.CloseStatusUpdate_ChanClose)
 	go func() {

lntest/itest/assertions.go

@@ -84,7 +84,7 @@ func openChannelAndAssert(t *harnessTest, net *lntest.NetworkHarness,
 	// case that the channel is public, it is announced to the network.
 	block := mineBlocks(t, net, 6, 1)[0]
 
-	fundingChanPoint, err := net.WaitForChannelOpen(ctx, chanOpenUpdate)
+	fundingChanPoint, err := net.WaitForChannelOpen(chanOpenUpdate)
 	require.NoError(t.t, err, "error while waiting for channel open")
 
 	fundingTxID, err := lnrpc.GetChanPointFundingTxid(fundingChanPoint)
@@ -343,7 +343,7 @@ func assertChannelClosed(ctx context.Context, t *harnessTest,
 
 	block := mineBlocks(t, net, 1, expectedTxes)[0]
 
-	closingTxid, err := net.WaitForChannelClose(ctx, closeUpdates)
+	closingTxid, err := net.WaitForChannelClose(closeUpdates)
 	require.NoError(t.t, err, "error while waiting for channel close")
 
 	assertTxInBlock(t, block, closingTxid)

lntest/itest/lnd_channel_graph_test.go

@@ -314,8 +314,7 @@ func testUnannouncedChannels(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// One block is enough to make the channel ready for use, since the
 	// nodes have defaultNumConfs=1 set.
-	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
+	fundingChanPoint, err := net.WaitForChannelOpen(chanOpenUpdate)
-	fundingChanPoint, err := net.WaitForChannelOpen(ctxt, chanOpenUpdate)
 	if err != nil {
 		t.Fatalf("error while waiting for channel open: %v", err)
 	}
@@ -324,7 +323,7 @@ func testUnannouncedChannels(net *lntest.NetworkHarness, t *harnessTest) {
 	req := &lnrpc.ChannelGraphRequest{
 		IncludeUnannounced: true,
 	}
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
 	chanGraph, err := net.Alice.DescribeGraph(ctxt, req)
 	if err != nil {
 		t.Fatalf("unable to query alice's graph: %v", err)

lntest/itest/lnd_funding_test.go

@@ -354,8 +354,7 @@ func testUnconfirmedChannelFunding(net *lntest.NetworkHarness, t *harnessTest) {
 	// parties. Two transactions should be mined, the unconfirmed spend and
 	// the funding tx.
 	mineBlocks(t, net, 6, 2)
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+	chanPoint, err := net.WaitForChannelOpen(chanOpenUpdate)
-	chanPoint, err := net.WaitForChannelOpen(ctxt, chanOpenUpdate)
 	require.NoError(t.t, err, "error while waitinng for channel open")
 
 	// With the channel open, we'll check the balances on each side of the

lntest/itest/lnd_misc_test.go

@@ -547,8 +547,7 @@ func testMaxPendingChannels(net *lntest.NetworkHarness, t *harnessTest) {
 
 	chanPoints := make([]*lnrpc.ChannelPoint, maxPendingChannels)
 	for i, stream := range openStreams {
-		ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+		fundingChanPoint, err := net.WaitForChannelOpen(stream)
-		fundingChanPoint, err := net.WaitForChannelOpen(ctxt, stream)
 		if err != nil {
 			t.Fatalf("error while waiting for channel open: %v", err)
 		}

lntest/itest/lnd_revocation_test.go

@@ -203,8 +203,7 @@ func testRevokedCloseRetribution(net *lntest.NetworkHarness, t *harnessTest) {
 	// block.
 	block := mineBlocks(t, net, 1, 1)[0]
 
-	ctxt, _ = context.WithTimeout(ctxb, channelCloseTimeout)
+	breachTXID, err := net.WaitForChannelClose(closeUpdates)
-	breachTXID, err := net.WaitForChannelClose(ctxt, closeUpdates)
 	if err != nil {
 		t.Fatalf("error while waiting for channel close: %v", err)
 	}
@@ -439,8 +438,7 @@ func testRevokedCloseRetributionZeroValueRemoteOutput(net *lntest.NetworkHarness
 		t.Fatalf("unable to stop Dave's node: %v", err)
 	}
 
-	ctxt, _ = context.WithTimeout(ctxb, channelCloseTimeout)
+	breachTXID, err := net.WaitForChannelClose(closeUpdates)
-	breachTXID, err := net.WaitForChannelClose(ctxt, closeUpdates)
 	if err != nil {
 		t.Fatalf("error while waiting for channel close: %v", err)
 	}
@@ -746,8 +744,7 @@ func testRevokedCloseRetributionRemoteHodl(net *lntest.NetworkHarness,
 
 	// Finally, wait for the final close status update, then ensure that
 	// the closing transaction was included in the block.
-	ctxt, _ = context.WithTimeout(ctxb, channelCloseTimeout)
+	breachTXID, err := net.WaitForChannelClose(closeUpdates)
-	breachTXID, err := net.WaitForChannelClose(ctxt, closeUpdates)
 	if err != nil {
 		t.Fatalf("error while waiting for channel close: %v", err)
 	}
@@ -1181,8 +1178,7 @@ func testRevokedCloseRetributionAltruistWatchtowerCase(
 	// block.
 	block := mineBlocks(t, net, 1, 1)[0]
 
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+	breachTXID, err := net.WaitForChannelClose(closeUpdates)
-	breachTXID, err := net.WaitForChannelClose(ctxt, closeUpdates)
 	if err != nil {
 		t.Fatalf("error while waiting for channel close: %v", err)
 	}

lntest/itest/lnd_routing_test.go

@@ -840,8 +840,7 @@ func testPrivateChannels(net *lntest.NetworkHarness, t *harnessTest) {
 	// nodes have defaultNumConfs=1 set.
 	block := mineBlocks(t, net, 1, 1)[0]
 
-	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
+	chanPointPrivate, err := net.WaitForChannelOpen(chanOpenUpdate)
-	chanPointPrivate, err := net.WaitForChannelOpen(ctxt, chanOpenUpdate)
 	if err != nil {
 		t.Fatalf("error while waiting for channel open: %v", err)
 	}
@@ -857,7 +856,7 @@ func testPrivateChannels(net *lntest.NetworkHarness, t *harnessTest) {
 		Hash:  *fundingTxID,
 		Index: chanPointPrivate.OutputIndex,
 	}
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
 	err = net.AssertChannelExists(ctxt, carol, &privateFundPoint)
 	if err != nil {
 		t.Fatalf("unable to assert channel existence: %v", err)