lntest+wait: replace sleeps in assertAmountSent

This commit constructs a helper closure assertAmountSent that can be
reused by other functions. The closure returns an error so that it can
be used with wait.NoError or the new wait.InvariantNoError. The latter
is added since the predicate could otherwise pass immediately for the
sphinx_replay_persistence tests, but change shortly after. It also
rounds out the wait package so that we offer all combinations of
predicate and no-error style waits.

lntest/itest/lnd_test.go

@@ -3704,6 +3704,48 @@ func testChannelForceClosure(net *lntest.NetworkHarness, t *harnessTest) {
 	}
 }
 
+// assertAmountSent generates a closure which queries listchannels for sndr and
+// rcvr, and asserts that sndr sent amt satoshis, and that rcvr received amt
+// satoshis.
+//
+// NOTE: This method assumes that each node only has one channel, and it is the
+// channel used to send the payment.
+func assertAmountSent(amt btcutil.Amount, sndr, rcvr *lntest.HarnessNode) func() error {
+	return func() error {
+		// Both channels should also have properly accounted from the
+		// amount that has been sent/received over the channel.
+		listReq := &lnrpc.ListChannelsRequest{}
+		ctxb := context.Background()
+		ctxt, _ := context.WithTimeout(ctxb, defaultTimeout)
+		sndrListChannels, err := sndr.ListChannels(ctxt, listReq)
+		if err != nil {
+			return fmt.Errorf("unable to query for %s's channel "+
+				"list: %v", sndr.Name(), err)
+		}
+		sndrSatoshisSent := sndrListChannels.Channels[0].TotalSatoshisSent
+		if sndrSatoshisSent != int64(amt) {
+			return fmt.Errorf("%s's satoshis sent is incorrect "+
+				"got %v, expected %v", sndr.Name(),
+				sndrSatoshisSent, amt)
+		}
+
+		ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
+		rcvrListChannels, err := rcvr.ListChannels(ctxt, listReq)
+		if err != nil {
+			return fmt.Errorf("unable to query for %s's channel "+
+				"list: %v", rcvr.Name(), err)
+		}
+		rcvrSatoshisReceived := rcvrListChannels.Channels[0].TotalSatoshisReceived
+		if rcvrSatoshisReceived != int64(amt) {
+			return fmt.Errorf("%s's satoshis received is "+
+				"incorrect got %v, expected %v", rcvr.Name(),
+				rcvrSatoshisReceived, amt)
+		}
+
+		return nil
+	}
+}
+
 // testSphinxReplayPersistence verifies that replayed onion packets are rejected
 // by a remote peer after a restart. We use a combination of unsafe
 // configuration arguments to force Carol to replay the same sphinx packet after
@@ -3753,33 +3795,6 @@ func testSphinxReplayPersistence(net *lntest.NetworkHarness, t *harnessTest) {
 		},
 	)
 
-	assertAmountSent := func(amt btcutil.Amount) {
-		// Both channels should also have properly accounted from the
-		// amount that has been sent/received over the channel.
-		listReq := &lnrpc.ListChannelsRequest{}
-		ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-		carolListChannels, err := carol.ListChannels(ctxt, listReq)
-		if err != nil {
-			t.Fatalf("unable to query for alice's channel list: %v", err)
-		}
-		carolSatoshisSent := carolListChannels.Channels[0].TotalSatoshisSent
-		if carolSatoshisSent != int64(amt) {
-			t.Fatalf("Carol's satoshis sent is incorrect got %v, expected %v",
-				carolSatoshisSent, amt)
-		}
-
-		ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-		daveListChannels, err := dave.ListChannels(ctxt, listReq)
-		if err != nil {
-			t.Fatalf("unable to query for Dave's channel list: %v", err)
-		}
-		daveSatoshisReceived := daveListChannels.Channels[0].TotalSatoshisReceived
-		if daveSatoshisReceived != int64(amt) {
-			t.Fatalf("Dave's satoshis received is incorrect got %v, expected %v",
-				daveSatoshisReceived, amt)
-		}
-	}
-
 	// Now that the channel is open, create an invoice for Dave which
 	// expects a payment of 1000 satoshis from Carol paid via a particular
 	// preimage.
@@ -3844,8 +3859,12 @@ func testSphinxReplayPersistence(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// With the payment sent but hedl, all balance related stats should not
 	// have changed.
-	time.Sleep(time.Millisecond * 200)
-	assertAmountSent(0)
+	err = wait.InvariantNoError(
+		assertAmountSent(0, carol, dave), 3*time.Second,
+	)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
 
 	// With the first payment sent, restart dave to make sure he is
 	// persisting the information required to detect replayed sphinx
@@ -3874,7 +3893,12 @@ func testSphinxReplayPersistence(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// Since the payment failed, the balance should still be left
 	// unaltered.
-	assertAmountSent(0)
+	err = wait.InvariantNoError(
+		assertAmountSent(0, carol, dave), 3*time.Second,
+	)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
 
 	ctxt, _ = context.WithTimeout(ctxb, channelCloseTimeout)
 	closeChannelAndAssert(ctxt, t, net, carol, chanPoint, true)
@@ -3897,33 +3921,6 @@ func testSingleHopInvoice(net *lntest.NetworkHarness, t *harnessTest) {
 		},
 	)
 
-	assertAmountSent := func(amt btcutil.Amount) {
-		// Both channels should also have properly accounted from the
-		// amount that has been sent/received over the channel.
-		listReq := &lnrpc.ListChannelsRequest{}
-		ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-		aliceListChannels, err := net.Alice.ListChannels(ctxt, listReq)
-		if err != nil {
-			t.Fatalf("unable to query for alice's channel list: %v", err)
-		}
-		aliceSatoshisSent := aliceListChannels.Channels[0].TotalSatoshisSent
-		if aliceSatoshisSent != int64(amt) {
-			t.Fatalf("Alice's satoshis sent is incorrect got %v, expected %v",
-				aliceSatoshisSent, amt)
-		}
-
-		ctxt, _ = context.WithTimeout(ctxb, defaultTimeout)
-		bobListChannels, err := net.Bob.ListChannels(ctxt, listReq)
-		if err != nil {
-			t.Fatalf("unable to query for bob's channel list: %v", err)
-		}
-		bobSatoshisReceived := bobListChannels.Channels[0].TotalSatoshisReceived
-		if bobSatoshisReceived != int64(amt) {
-			t.Fatalf("Bob's satoshis received is incorrect got %v, expected %v",
-				bobSatoshisReceived, amt)
-		}
-	}
-
 	// Now that the channel is open, create an invoice for Bob which
 	// expects a payment of 1000 satoshis from Alice paid via a particular
 	// preimage.
@@ -3989,8 +3986,13 @@ func testSingleHopInvoice(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// With the payment completed all balance related stats should be
 	// properly updated.
-	time.Sleep(time.Millisecond * 200)
-	assertAmountSent(paymentAmt)
+	err = wait.NoError(
+		assertAmountSent(paymentAmt, net.Alice, net.Bob),
+		3*time.Second,
+	)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
 
 	// Create another invoice for Bob, this time leaving off the preimage
 	// to one will be randomly generated. We'll test the proper
@@ -4021,8 +4023,13 @@ func testSingleHopInvoice(net *lntest.NetworkHarness, t *harnessTest) {
 
 	// The second payment should also have succeeded, with the balances
 	// being update accordingly.
-	time.Sleep(time.Millisecond * 200)
-	assertAmountSent(paymentAmt * 2)
+	err = wait.NoError(
+		assertAmountSent(2*paymentAmt, net.Alice, net.Bob),
+		3*time.Second,
+	)
+	if err != nil {
+		t.Fatalf(err.Error())
+	}
 
 	ctxt, _ = context.WithTimeout(ctxb, channelCloseTimeout)
 	closeChannelAndAssert(ctxt, t, net, net.Alice, chanPoint, false)

lntest/wait/wait.go

@@ -76,3 +76,23 @@ func Invariant(statement func() bool, timeout time.Duration) error {
 		}
 	}
 }
+
+// InvariantNoError is a wrapper around Invariant that waits out the duration
+// specified by timeout. It fails if the predicate ever returns an error during
+// that time.
+func InvariantNoError(f func() error, timeout time.Duration) error {
+	var predErr error
+	pred := func() bool {
+		if err := f(); err != nil {
+			predErr = err
+			return false
+		}
+		return true
+	}
+
+	if err := Invariant(pred, timeout); err != nil {
+		return predErr
+	}
+
+	return nil
+}