From d39303f2469305360281b500e1b3e3206e763d49 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Thu, 10 Nov 2022 14:47:48 +0800
Subject: [PATCH] itest: wrap `assertTxLabel` inside `wait`

---
 lntest/itest/lnd_misc_test.go | 36 ++++++++++++++++++++++++++++-------
 lntest/wait/wait.go           |  3 +++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/lntest/itest/lnd_misc_test.go b/lntest/itest/lnd_misc_test.go
index a548a77c6..d9af97f3c 100644
--- a/lntest/itest/lnd_misc_test.go
+++ b/lntest/itest/lnd_misc_test.go
@@ -1262,23 +1262,45 @@ func testSweepAllCoins(ht *lntemp.HarnessTest) {
 
 	// assertTxLabel is a helper function which finds a target tx in our
 	// set of transactions and checks that it has the desired label.
-	assertTxLabel := func(targetTx, label string) {
+	assertTxLabel := func(targetTx, label string) error {
 		// List all transactions relevant to our wallet, and find the
 		// tx so that we can check the correct label has been set.
 		txResp := ainz.RPC.GetTransactions()
 
-		// Find our transaction in the set of transactions returned and
-		// check its label.
+		var target *lnrpc.Transaction
+
+		// First we need to find the target tx.
 		for _, txn := range txResp.Transactions {
 			if txn.TxHash == targetTx {
-				require.Equal(ht, label, txn.Label,
-					"labels not match")
+				target = txn
 			}
 		}
+
+		// If we cannot find it, return an error.
+		if target == nil {
+			return fmt.Errorf("target tx %v not found", targetTx)
+		}
+
+		// Otherwise, check the labels are matched.
+		if target.Label == label {
+			return nil
+		}
+
+		return fmt.Errorf("labels not match, want: "+
+			"%v, got %v", label, target.Label)
+	}
+
+	// waitTxLabel waits until the desired tx label is found or timeout.
+	waitTxLabel := func(targetTx, label string) {
+		err := wait.NoError(func() error {
+			return assertTxLabel(targetTx, label)
+		}, defaultTimeout)
+
+		require.NoError(ht, err, "timeout assertTxLabel")
 	}
 
 	sweepTxStr := sweepTx.TxHash().String()
-	assertTxLabel(sweepTxStr, sendCoinsLabel)
+	waitTxLabel(sweepTxStr, sendCoinsLabel)
 
 	// While we are looking at labels, we test our label transaction
 	// command to make sure it is behaving as expected. First, we try to
@@ -1322,7 +1344,7 @@ func testSweepAllCoins(ht *lntemp.HarnessTest) {
 	}
 	ainz.RPC.LabelTransaction(req)
 
-	assertTxLabel(sweepTxStr, newLabel)
+	waitTxLabel(sweepTxStr, newLabel)
 
 	// Finally, Ainz should now have no coins at all within his wallet.
 	resp := ainz.RPC.WalletBalance()
diff --git a/lntest/wait/wait.go b/lntest/wait/wait.go
index 16a5f2ab2..b6274086c 100644
--- a/lntest/wait/wait.go
+++ b/lntest/wait/wait.go
@@ -13,6 +13,9 @@ const PollInterval = 200 * time.Millisecond
 // timing doesn't always line up well when running integration tests with
 // several running lnd nodes. This function gives callers a way to assert that
 // some property is upheld within a particular time frame.
+//
+// TODO(yy): build a counter here so we know how many times we've tried the
+// `pred`.
 func Predicate(pred func() bool, timeout time.Duration) error {
 	exitTimer := time.After(timeout)
 	result := make(chan bool, 1)