From 74f8fe482d05ccf9dd81f4cc0755b63afad09b78 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 19 Aug 2021 21:00:21 +0800
Subject: [PATCH] itest: manage context inside EnsureConnected

---
 lntest/harness.go                            |  6 +++--
 lntest/itest/assertions.go                   |  2 +-
 lntest/itest/lnd_channel_backup_test.go      |  6 ++---
 lntest/itest/lnd_channel_balance_test.go     |  5 ++--
 lntest/itest/lnd_channel_graph_test.go       | 14 ++++-------
 lntest/itest/lnd_forward_interceptor_test.go |  5 +---
 lntest/itest/lnd_funding_test.go             | 13 ++++------
 lntest/itest/lnd_max_channel_size_test.go    |  6 ++---
 lntest/itest/lnd_misc_test.go                | 11 ++++-----
 lntest/itest/lnd_mpp_test.go                 |  5 +---
 lntest/itest/lnd_multi-hop_test.go           |  5 ++--
 lntest/itest/lnd_payment_test.go             |  3 +--
 lntest/itest/lnd_psbt_test.go                |  8 +++----
 lntest/itest/lnd_rest_api_test.go            |  2 +-
 lntest/itest/lnd_switch_test.go              | 25 +++++++-------------
 lntest/itest/lnd_test.go                     |  4 +---
 lntest/itest/lnd_wallet_import_test.go       |  4 +---
 lntest/itest/lnd_wumbo_channels_test.go      |  4 ++--
 18 files changed, 46 insertions(+), 82 deletions(-)

diff --git a/lntest/harness.go b/lntest/harness.go
index 103219d15..f0698c19c 100644
--- a/lntest/harness.go
+++ b/lntest/harness.go
@@ -548,8 +548,10 @@ tryconnect:
 // behave the same as ConnectNodes. If a pending connection request has already
 // been made, the method will block until the two nodes appear in each other's
 // peers list, or until the 15s timeout expires.
-func (n *NetworkHarness) EnsureConnected(ctx context.Context,
-	t *testing.T, a, b *HarnessNode) {
+func (n *NetworkHarness) EnsureConnected(t *testing.T, a, b *HarnessNode) {
+	ctxb := context.Background()
+	ctx, cancel := context.WithTimeout(ctxb, DefaultTimeout)
+	defer cancel()
 
 	// errConnectionRequested is used to signal that a connection was
 	// requested successfully, which is distinct from already being
diff --git a/lntest/itest/assertions.go b/lntest/itest/assertions.go
index 357d82a6c..6819a23a6 100644
--- a/lntest/itest/assertions.go
+++ b/lntest/itest/assertions.go
@@ -1279,7 +1279,7 @@ func assertDLPExecuted(net *lntest.NetworkHarness, t *harnessTest,
 	// To make sure the nodes are initiating DLP now, we have to manually
 	// re-connect them.
 	ctxb := context.Background()
-	net.EnsureConnected(ctxb, t.t, carol, dave)
+	net.EnsureConnected(t.t, carol, dave)
 
 	// Upon reconnection, the nodes should detect that Dave is out of sync.
 	// Carol should force close the channel using her latest commitment.
diff --git a/lntest/itest/lnd_channel_backup_test.go b/lntest/itest/lnd_channel_backup_test.go
index 65a6fe64f..95e805df2 100644
--- a/lntest/itest/lnd_channel_backup_test.go
+++ b/lntest/itest/lnd_channel_backup_test.go
@@ -1079,8 +1079,7 @@ func testChanRestoreScenario(t *harnessTest, net *lntest.NetworkHarness,
 		// Now that we have our new node up, we expect that it'll
 		// re-connect to Carol automatically based on the restored
 		// backup.
-		ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-		net.EnsureConnected(ctxt, t.t, dave, carol)
+		net.EnsureConnected(t.t, dave, carol)
 
 		assertTimeLockSwept(
 			net, t, carol, carolStartingBalance, dave,
@@ -1172,8 +1171,7 @@ func testChanRestoreScenario(t *harnessTest, net *lntest.NetworkHarness,
 
 	// Now that we have our new node up, we expect that it'll re-connect to
 	// Carol automatically based on the restored backup.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, dave, carol)
+	net.EnsureConnected(t.t, dave, carol)
 
 	// TODO(roasbeef): move dave restarts?
 
diff --git a/lntest/itest/lnd_channel_balance_test.go b/lntest/itest/lnd_channel_balance_test.go
index db531179a..d847e6ba6 100644
--- a/lntest/itest/lnd_channel_balance_test.go
+++ b/lntest/itest/lnd_channel_balance_test.go
@@ -51,8 +51,7 @@ func testChannelBalance(net *lntest.NetworkHarness, t *harnessTest) {
 	}
 
 	// Before beginning, make sure alice and bob are connected.
-	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, net.Alice, net.Bob)
+	net.EnsureConnected(t.t, net.Alice, net.Bob)
 
 	chanPoint := openChannelAndAssert(
 		t, net, net.Alice, net.Bob,
@@ -62,7 +61,7 @@ func testChannelBalance(net *lntest.NetworkHarness, t *harnessTest) {
 	)
 
 	// Wait for both Alice and Bob to recognize this new channel.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
 	err := net.Alice.WaitForNetworkChannelOpen(ctxt, chanPoint)
 	if err != nil {
 		t.Fatalf("alice didn't advertise channel before "+
diff --git a/lntest/itest/lnd_channel_graph_test.go b/lntest/itest/lnd_channel_graph_test.go
index ee9ca705f..c0f9c2020 100644
--- a/lntest/itest/lnd_channel_graph_test.go
+++ b/lntest/itest/lnd_channel_graph_test.go
@@ -204,8 +204,7 @@ func testUpdateChanStatus(net *lntest.NetworkHarness, t *harnessTest) {
 	)
 
 	// Reconnecting the nodes should propagate a "Disabled = false" update.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, alice, bob)
+	net.EnsureConnected(t.t, alice, bob)
 	expectedPolicy.Disabled = false
 	waitForChannelUpdate(
 		t, carolSub,
@@ -248,8 +247,7 @@ func testUpdateChanStatus(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// Bob sends a "Disabled = false" update upon detecting the
 	// reconnect.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, alice, bob)
+	net.EnsureConnected(t.t, alice, bob)
 	expectedPolicy.Disabled = false
 	waitForChannelUpdate(
 		t, carolSub,
@@ -282,7 +280,7 @@ func testUpdateChanStatus(net *lntest.NetworkHarness, t *harnessTest) {
 	// BOTH Alice and Bob should set the channel state back to "enabled"
 	// on reconnect.
 	sendReq(alice, chanPoint, routerrpc.ChanStatusAction_AUTO)
-	net.EnsureConnected(ctxt, t.t, alice, bob)
+	net.EnsureConnected(t.t, alice, bob)
 	expectedPolicy.Disabled = false
 	waitForChannelUpdate(
 		t, carolSub,
@@ -440,8 +438,7 @@ func testGraphTopologyNtfns(net *lntest.NetworkHarness, t *harnessTest, pinned b
 	defer shutdownAndAssert(net, t, alice)
 
 	// Connect Alice and Bob.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, alice, bob)
+	net.EnsureConnected(t.t, alice, bob)
 
 	// Alice stimmy.
 	net.SendCoins(t.t, btcutil.SatoshiPerBitcoin, alice)
@@ -609,8 +606,7 @@ out:
 	// and Carol. Note that we will also receive a node announcement from
 	// Bob, since a node will update its node announcement after a new
 	// channel is opened.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, alice, bob)
+	net.EnsureConnected(t.t, alice, bob)
 
 	// We should receive an update advertising the newly connected node,
 	// Bob's new node announcement, and the channel between Bob and Carol.
diff --git a/lntest/itest/lnd_forward_interceptor_test.go b/lntest/itest/lnd_forward_interceptor_test.go
index 587a4c478..2698cf4ca 100644
--- a/lntest/itest/lnd_forward_interceptor_test.go
+++ b/lntest/itest/lnd_forward_interceptor_test.go
@@ -286,14 +286,11 @@ func newInterceptorTestContext(t *harnessTest,
 	net *lntest.NetworkHarness,
 	alice, bob, carol *lntest.HarnessNode) *interceptorTestContext {
 
-	ctxb := context.Background()
-
 	// Connect nodes
 	nodes := []*lntest.HarnessNode{alice, bob, carol}
 	for i := 0; i < len(nodes); i++ {
 		for j := i + 1; j < len(nodes); j++ {
-			ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
-			net.EnsureConnected(ctxt, t.t, nodes[i], nodes[j])
+			net.EnsureConnected(t.t, nodes[i], nodes[j])
 		}
 	}
 
diff --git a/lntest/itest/lnd_funding_test.go b/lntest/itest/lnd_funding_test.go
index fe457d8a9..683be7586 100644
--- a/lntest/itest/lnd_funding_test.go
+++ b/lntest/itest/lnd_funding_test.go
@@ -27,8 +27,6 @@ import (
 // transaction was mined.
 func testBasicChannelFunding(net *lntest.NetworkHarness, t *harnessTest) {
 
-	ctxb := context.Background()
-
 	// Run through the test with combinations of all the different
 	// commitment types.
 	allTypes := []commitType{
@@ -57,8 +55,7 @@ func testBasicChannelFunding(net *lntest.NetworkHarness, t *harnessTest) {
 
 		// Before we start the test, we'll ensure both sides are
 		// connected to the funding flow can properly be executed.
-		ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
-		net.EnsureConnected(ctxt, t.t, carol, dave)
+		net.EnsureConnected(t.t, carol, dave)
 
 		carolChan, daveChan, closeChan, err := basicChannelFundingTest(
 			t, net, carol, dave, nil,
@@ -390,8 +387,7 @@ func testExternalFundingChanPoint(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// Before we start the test, we'll ensure both sides are connected to
 	// the funding flow can properly be executed.
-	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, carol, dave)
+	net.EnsureConnected(t.t, carol, dave)
 
 	// At this point, we're ready to simulate our external channel funding
 	// flow. To start with, we'll create a pending channel with a shim for
@@ -407,7 +403,7 @@ func testExternalFundingChanPoint(net *lntest.NetworkHarness, t *harnessTest) {
 			FundingShim: fundingShim1,
 		},
 	)
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
 	assertNumOpenChannelsPending(ctxt, t, carol, dave, 1)
 
 	// That channel is now pending forever and normally would saturate the
@@ -570,8 +566,7 @@ func testChannelFundingPersistence(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// The following block ensures that after both nodes have restarted,
 	// they have reconnected before the execution of the next test.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, net.Alice, carol)
+	net.EnsureConnected(t.t, net.Alice, carol)
 
 	// Next, mine enough blocks s.t the channel will open with a single
 	// additional block mined.
diff --git a/lntest/itest/lnd_max_channel_size_test.go b/lntest/itest/lnd_max_channel_size_test.go
index 335c8f6da..4491e61c6 100644
--- a/lntest/itest/lnd_max_channel_size_test.go
+++ b/lntest/itest/lnd_max_channel_size_test.go
@@ -35,7 +35,7 @@ func testMaxChannelSize(net *lntest.NetworkHarness, t *harnessTest) {
 	// Next we'll connect both nodes, then attempt to make a wumbo channel
 	// funding request, which should fail as it exceeds the default wumbo
 	// soft limit of 10 BTC.
-	net.EnsureConnected(ctxb, t.t, wumboNode, wumboNode2)
+	net.EnsureConnected(t.t, wumboNode, wumboNode2)
 
 	chanAmt := funding.MaxBtcFundingAmountWumbo + 1
 	_, err := net.OpenChannel(
@@ -58,7 +58,7 @@ func testMaxChannelSize(net *lntest.NetworkHarness, t *harnessTest) {
 	miniNode := net.NewNode(t.t, "mini", nil)
 	defer shutdownAndAssert(net, t, miniNode)
 
-	net.EnsureConnected(ctxb, t.t, wumboNode, miniNode)
+	net.EnsureConnected(t.t, wumboNode, miniNode)
 
 	_, err = net.OpenChannel(
 		ctxb, wumboNode, miniNode, lntest.OpenChannelParams{
@@ -89,7 +89,7 @@ func testMaxChannelSize(net *lntest.NetworkHarness, t *harnessTest) {
 	defer shutdownAndAssert(net, t, wumboNode3)
 
 	// Creating a wumbo channel between these two nodes should succeed.
-	net.EnsureConnected(ctxb, t.t, wumboNode, wumboNode3)
+	net.EnsureConnected(t.t, wumboNode, wumboNode3)
 	chanPoint := openChannelAndAssert(
 		t, net, wumboNode, wumboNode3,
 		lntest.OpenChannelParams{
diff --git a/lntest/itest/lnd_misc_test.go b/lntest/itest/lnd_misc_test.go
index c95c7e55d..e93e57063 100644
--- a/lntest/itest/lnd_misc_test.go
+++ b/lntest/itest/lnd_misc_test.go
@@ -862,8 +862,7 @@ func testDataLossProtection(net *lntest.NetworkHarness, t *harnessTest) {
 
 		// We must let the node communicate with Carol before they are
 		// able to open channel, so we connect them.
-		ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
-		net.EnsureConnected(ctxt, t.t, carol, node)
+		net.EnsureConnected(t.t, carol, node)
 
 		// We'll first open up a channel between them with a 0.5 BTC
 		// value.
@@ -888,7 +887,7 @@ func testDataLossProtection(net *lntest.NetworkHarness, t *harnessTest) {
 
 		// Wait for Carol to receive the channel edge from the funding
 		// manager.
-		ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+		ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
 		err = carol.WaitForNetworkChannelOpen(ctxt, chanPoint)
 		if err != nil {
 			t.Fatalf("carol didn't see the carol->%s channel "+
@@ -1479,8 +1478,7 @@ func testSendUpdateDisableChannel(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// Reconnect Carol and Eve, this should cause them to reenable the
 	// channel from both ends after a short delay.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, carol, eve)
+	net.EnsureConnected(t.t, carol, eve)
 
 	expectedPolicy.Disabled = false
 	waitForChannelUpdate(
@@ -1502,8 +1500,7 @@ func testSendUpdateDisableChannel(net *lntest.NetworkHarness, t *harnessTest) {
 		t.Fatalf("unable to disconnect Carol from Eve: %v", err)
 	}
 	time.Sleep(time.Second)
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, eve, carol)
+	net.EnsureConnected(t.t, eve, carol)
 
 	// Since the disable should have been canceled by both Carol and Eve, we
 	// expect no channel updates to appear on the network.
diff --git a/lntest/itest/lnd_mpp_test.go b/lntest/itest/lnd_mpp_test.go
index 9b12bff15..7c51e66f2 100644
--- a/lntest/itest/lnd_mpp_test.go
+++ b/lntest/itest/lnd_mpp_test.go
@@ -255,8 +255,6 @@ type mppTestContext struct {
 func newMppTestContext(t *harnessTest,
 	net *lntest.NetworkHarness) *mppTestContext {
 
-	ctxb := context.Background()
-
 	alice := net.NewNode(t.t, "alice", nil)
 	bob := net.NewNode(t.t, "bob", []string{"--accept-amp"})
 
@@ -270,8 +268,7 @@ func newMppTestContext(t *harnessTest,
 	nodes := []*lntest.HarnessNode{alice, bob, carol, dave, eve}
 	for i := 0; i < len(nodes); i++ {
 		for j := i + 1; j < len(nodes); j++ {
-			ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
-			net.EnsureConnected(ctxt, t.t, nodes[i], nodes[j])
+			net.EnsureConnected(t.t, nodes[i], nodes[j])
 		}
 	}
 
diff --git a/lntest/itest/lnd_multi-hop_test.go b/lntest/itest/lnd_multi-hop_test.go
index ea0ecfd00..bf245d7b1 100644
--- a/lntest/itest/lnd_multi-hop_test.go
+++ b/lntest/itest/lnd_multi-hop_test.go
@@ -206,8 +206,7 @@ func createThreeHopNetwork(t *harnessTest, net *lntest.NetworkHarness,
 
 	ctxb := context.Background()
 
-	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, alice, bob)
+	net.EnsureConnected(t.t, alice, bob)
 
 	// Make sure there are enough utxos for anchoring.
 	for i := 0; i < 2; i++ {
@@ -225,7 +224,7 @@ func createThreeHopNetwork(t *harnessTest, net *lntest.NetworkHarness,
 		},
 	)
 
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+	ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
 	err := alice.WaitForNetworkChannelOpen(ctxt, aliceChanPoint)
 	if err != nil {
 		t.Fatalf("alice didn't report channel: %v", err)
diff --git a/lntest/itest/lnd_payment_test.go b/lntest/itest/lnd_payment_test.go
index 1980abffd..98c254392 100644
--- a/lntest/itest/lnd_payment_test.go
+++ b/lntest/itest/lnd_payment_test.go
@@ -196,8 +196,7 @@ func testPaymentFollowingChannelOpen(net *lntest.NetworkHarness, t *harnessTest)
 	}
 
 	// We ensure that Bob reconnects to Alice.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, net.Bob, net.Alice)
+	net.EnsureConnected(t.t, net.Bob, net.Alice)
 
 	// We mine one block for the channel to be confirmed.
 	_ = mineBlocks(t, net, 6, 1)[0]
diff --git a/lntest/itest/lnd_psbt_test.go b/lntest/itest/lnd_psbt_test.go
index 445134f4d..1ab8315d3 100644
--- a/lntest/itest/lnd_psbt_test.go
+++ b/lntest/itest/lnd_psbt_test.go
@@ -35,10 +35,8 @@ func testPsbtChanFunding(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// Before we start the test, we'll ensure both sides are connected so
 	// the funding flow can be properly executed.
-	ctxt, cancel := context.WithTimeout(ctxb, defaultTimeout)
-	defer cancel()
-	net.EnsureConnected(ctxt, t.t, carol, dave)
-	net.EnsureConnected(ctxt, t.t, carol, net.Alice)
+	net.EnsureConnected(t.t, carol, dave)
+	net.EnsureConnected(t.t, carol, net.Alice)
 
 	// At this point, we can begin our PSBT channel funding workflow. We'll
 	// start by generating a pending channel ID externally that will be used
@@ -55,7 +53,7 @@ func testPsbtChanFunding(net *lntest.NetworkHarness, t *harnessTest) {
 	// Now that we have the pending channel ID, Carol will open the channel
 	// by specifying a PSBT shim. We use the NoPublish flag here to avoid
 	// publishing the whole batch TX too early.
-	ctxt, cancel = context.WithTimeout(ctxb, defaultTimeout)
+	ctxt, cancel := context.WithTimeout(ctxb, defaultTimeout)
 	defer cancel()
 	chanUpdates, tempPsbt, err := openChannelPsbt(
 		ctxt, carol, dave, lntest.OpenChannelParams{
diff --git a/lntest/itest/lnd_rest_api_test.go b/lntest/itest/lnd_rest_api_test.go
index ad009c941..96b6dcd89 100644
--- a/lntest/itest/lnd_rest_api_test.go
+++ b/lntest/itest/lnd_rest_api_test.go
@@ -499,7 +499,7 @@ func wsTestCaseBiDirectionalSubscription(ht *harnessTest,
 
 	// Before we start opening channels, make sure the two nodes are
 	// connected.
-	net.EnsureConnected(context.Background(), ht.t, net.Alice, net.Bob)
+	net.EnsureConnected(ht.t, net.Alice, net.Bob)
 
 	// Open 3 channels to make sure multiple requests and responses can be
 	// sent over the web socket.
diff --git a/lntest/itest/lnd_switch_test.go b/lntest/itest/lnd_switch_test.go
index 168632fb5..a3b411a60 100644
--- a/lntest/itest/lnd_switch_test.go
+++ b/lntest/itest/lnd_switch_test.go
@@ -189,11 +189,8 @@ func testSwitchCircuitPersistence(net *lntest.NetworkHarness, t *harnessTest) {
 	}
 
 	// Ensure all of the intermediate links are reconnected.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, net.Alice, dave)
-
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, net.Bob, net.Alice)
+	net.EnsureConnected(t.t, net.Alice, dave)
+	net.EnsureConnected(t.t, net.Bob, net.Alice)
 
 	// Ensure all nodes in the network still have 5 outstanding htlcs.
 	err = wait.Predicate(func() bool {
@@ -211,8 +208,7 @@ func testSwitchCircuitPersistence(net *lntest.NetworkHarness, t *harnessTest) {
 		t.Fatalf("Node restart failed: %v", err)
 	}
 
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, dave, carol)
+	net.EnsureConnected(t.t, dave, carol)
 
 	// After the payments settle, there should be no active htlcs on any of
 	// the nodes in the network.
@@ -523,8 +519,7 @@ func testSwitchOfflineDelivery(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// Now that the settles have reached Dave, reconnect him with Alice,
 	// allowing the settles to return to the sender.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, dave, net.Alice)
+	net.EnsureConnected(t.t, dave, net.Alice)
 
 	// Wait until all outstanding htlcs in the network have been settled.
 	err = wait.Predicate(func() bool {
@@ -783,8 +778,7 @@ func testSwitchOfflineDeliveryPersistence(net *lntest.NetworkHarness, t *harness
 
 	// Make Carol and Dave are reconnected before waiting for the htlcs to
 	// clear.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, dave, carol)
+	net.EnsureConnected(t.t, dave, carol)
 
 	// Wait for Carol to report no outstanding htlcs, and also for Dav to
 	// receive all the settles from Carol.
@@ -813,8 +807,7 @@ func testSwitchOfflineDeliveryPersistence(net *lntest.NetworkHarness, t *harness
 
 	// Force Dave and Alice to reconnect before waiting for the htlcs to
 	// clear.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, dave, net.Alice)
+	net.EnsureConnected(t.t, dave, net.Alice)
 
 	// After reconnection succeeds, the settles should be propagated all
 	// the way back to the sender. All nodes should report no active htlcs.
@@ -866,8 +859,7 @@ func testSwitchOfflineDeliveryPersistence(net *lntest.NetworkHarness, t *harness
 
 	// Before completing the final payment request, ensure that the
 	// connection between Dave and Carol has been healed.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, dave, carol)
+	net.EnsureConnected(t.t, dave, carol)
 
 	// Using Carol as the source, pay to the 5 invoices from Bob created
 	// above.
@@ -1114,8 +1106,7 @@ func testSwitchOfflineDeliveryOutgoingOffline(
 
 	// Ensure that Dave is reconnected to Alice before waiting for the
 	// htlcs to clear.
-	ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-	net.EnsureConnected(ctxt, t.t, dave, net.Alice)
+	net.EnsureConnected(t.t, dave, net.Alice)
 
 	// Since Carol has been shutdown permanently, we will wait until all
 	// other nodes in the network report no active htlcs.
diff --git a/lntest/itest/lnd_test.go b/lntest/itest/lnd_test.go
index 20acfdea9..5b152ddcf 100644
--- a/lntest/itest/lnd_test.go
+++ b/lntest/itest/lnd_test.go
@@ -1,7 +1,6 @@
 package itest
 
 import (
-	"context"
 	"flag"
 	"fmt"
 	"os"
@@ -225,8 +224,7 @@ func TestLightningNetworkDaemon(t *testing.T) {
 			}()
 
 			lndHarness.EnsureConnected(
-				context.Background(), t1,
-				lndHarness.Alice, lndHarness.Bob,
+				t1, lndHarness.Alice, lndHarness.Bob,
 			)
 
 			logLine := fmt.Sprintf(
diff --git a/lntest/itest/lnd_wallet_import_test.go b/lntest/itest/lnd_wallet_import_test.go
index 3ddb32384..854d9674b 100644
--- a/lntest/itest/lnd_wallet_import_test.go
+++ b/lntest/itest/lnd_wallet_import_test.go
@@ -332,9 +332,7 @@ func fundChanAndCloseFromImportedAccount(t *harnessTest, srcNode, destNode,
 
 	// Now, start the channel funding process. We'll need to connect both
 	// nodes first.
-	ctxt, cancel = context.WithTimeout(ctxb, defaultTimeout)
-	defer cancel()
-	t.lndHarness.EnsureConnected(ctxt, t.t, srcNode, destNode)
+	t.lndHarness.EnsureConnected(t.t, srcNode, destNode)
 
 	// The source node will then fund the channel through a PSBT shim.
 	var pendingChanID [32]byte
diff --git a/lntest/itest/lnd_wumbo_channels_test.go b/lntest/itest/lnd_wumbo_channels_test.go
index 3fd8fe3e3..5284db562 100644
--- a/lntest/itest/lnd_wumbo_channels_test.go
+++ b/lntest/itest/lnd_wumbo_channels_test.go
@@ -34,7 +34,7 @@ func testWumboChannels(net *lntest.NetworkHarness, t *harnessTest) {
 	// Next we'll connect both nodes, then attempt to make a wumbo channel
 	// funding request to the mini node we created above. The wumbo request
 	// should fail as the node isn't advertising wumbo channels.
-	net.EnsureConnected(ctxb, t.t, wumboNode, miniNode)
+	net.EnsureConnected(t.t, wumboNode, miniNode)
 
 	chanAmt := funding.MaxBtcFundingAmount + 1
 	_, err := net.OpenChannel(
@@ -61,7 +61,7 @@ func testWumboChannels(net *lntest.NetworkHarness, t *harnessTest) {
 	defer shutdownAndAssert(net, t, wumboNode2)
 
 	// Creating a wumbo channel between these two nodes should succeed.
-	net.EnsureConnected(ctxb, t.t, wumboNode, wumboNode2)
+	net.EnsureConnected(t.t, wumboNode, wumboNode2)
 	chanPoint := openChannelAndAssert(
 		t, net, wumboNode, wumboNode2,
 		lntest.OpenChannelParams{