From 7b43c9a750f61afbcbc8a784b3ba3d9eebfc209f Mon Sep 17 00:00:00 2001
From: Alex <alex@akselrod.org>
Date: Tue, 7 Nov 2017 12:05:53 -0700
Subject: [PATCH] lnwallet: add timeout and more accurate check to reorg
 interface test

---
 lnwallet/interface_test.go | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/lnwallet/interface_test.go b/lnwallet/interface_test.go
index 928d5bb6..86eafb3a 100644
--- a/lnwallet/interface_test.go
+++ b/lnwallet/interface_test.go
@@ -1177,17 +1177,29 @@ func testReorgWalletBalance(r *rpctest.Harness, w *lnwallet.LightningWallet,
 	// one block on the passed miner and two on the created miner,
 	// connecting them, and waiting for them to sync.
 	for i := 0; i < 5; i++ {
+		peers, err := r2.Node.GetPeerInfo()
+		if err != nil {
+			t.Fatalf("unable to get peer info: %v", err)
+		}
+		numPeers := len(peers)
 		err = r2.Node.AddNode(r.P2PAddress(), rpcclient.ANRemove)
 		if err != nil {
 			t.Fatalf("unable to disconnect mining nodes: %v", err)
 		}
 		// Wait for disconnection
+		timeout := time.After(30 * time.Second)
 		for true {
-			peers, err := r2.Node.GetPeerInfo()
+			// Allow for timeout
+			select {
+			case <-timeout:
+				t.Fatalf("timeout waiting for miner disconnect")
+			default:
+			}
+			peers, err = r2.Node.GetPeerInfo()
 			if err != nil {
 				t.Fatalf("unable to get peer info: %v", err)
 			}
-			if len(peers) == 0 {
+			if len(peers) < numPeers {
 				break
 			}
 			time.Sleep(100 * time.Millisecond)