From 0a335bbea2ead6d103fe28deb3dca2957522b04b Mon Sep 17 00:00:00 2001
From: Wilmer Paulino <wilmer.paulino@gmail.com>
Date: Thu, 9 Aug 2018 21:15:40 -0700
Subject: [PATCH] lntest: extract WaitForBalance into a method

---
 lntest/harness.go | 22 ++--------------------
 lntest/node.go    | 27 +++++++++++++++++++++++++++
 2 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/lntest/harness.go b/lntest/harness.go
index 60c0f4ef..d31c5592 100644
--- a/lntest/harness.go
+++ b/lntest/harness.go
@@ -1171,24 +1171,6 @@ func (n *NetworkHarness) sendCoins(ctx context.Context, amt btcutil.Amount,
 		return err
 	}
 
-	// Pause until the nodes current wallet balances reflects the amount
-	// sent to it above.
-	// TODO(roasbeef): factor out into helper func
-	balanceTicker := time.Tick(time.Millisecond * 50)
-	balanceTimeout := time.After(time.Second * 30)
-	for {
-		select {
-		case <-balanceTicker:
-			currentBal, err := target.WalletBalance(ctx, balReq)
-			if err != nil {
-				return err
-			}
-
-			if currentBal.ConfirmedBalance == initialBalance.ConfirmedBalance+int64(amt) {
-				return nil
-			}
-		case <-balanceTimeout:
-			return fmt.Errorf("balances not synced after deadline")
-		}
-	}
+	expectedBalance := initialBalance.ConfirmedBalance + int64(amt)
+	return target.WaitForBalance(expectedBalance, true)
 }
diff --git a/lntest/node.go b/lntest/node.go
index c3c44590..aa965568 100644
--- a/lntest/node.go
+++ b/lntest/node.go
@@ -895,6 +895,33 @@ func (hn *HarnessNode) WaitForBlockchainSync(ctx context.Context) error {
 	}
 }
 
+// WaitForBalance waits until the node sees the expected confirmed/unconfirmed
+// balance within their wallet.
+func (hn *HarnessNode) WaitForBalance(expectedBalance int64, confirmed bool) error {
+	ctx := context.Background()
+	req := &lnrpc.WalletBalanceRequest{}
+
+	doesBalanceMatch := func() bool {
+		balance, err := hn.WalletBalance(ctx, req)
+		if err != nil {
+			return false
+		}
+
+		if confirmed {
+			return balance.ConfirmedBalance == expectedBalance
+		}
+
+		return balance.UnconfirmedBalance == expectedBalance
+	}
+
+	err := WaitPredicate(doesBalanceMatch, 30*time.Second)
+	if err != nil {
+		return errors.New("balances not synced after deadline")
+	}
+
+	return nil
+}
+
 // fileExists reports whether the named file or directory exists.
 // This function is taken from https://github.com/btcsuite/btcd
 func fileExists(name string) bool {