From 07437f6ec41fcb8a1fcb06e5f7e8873685212d15 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Thu, 13 Apr 2017 15:11:20 -0700
Subject: [PATCH] test: update the ConnectPeer framework method to block until
 connect

This commit modifies the ConnectPeer method on the testing framework to
block (with a timeout) until the target peer is actually detected as
being connected. This was added as the peer connection logic was made
to be more asynchronous in a prior commit.
---
 networktest.go | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/networktest.go b/networktest.go
index db8443c0..ff2d18c6 100644
--- a/networktest.go
+++ b/networktest.go
@@ -786,7 +786,11 @@ func (n *networkHarness) NewNode(extraArgs []string) (*lightningNode, error) {
 }
 
 // ConnectNodes establishes an encrypted+authenticated p2p connection from node
-// a towards node b.
+// a towards node b. The function will return a non-nil error if the connection
+// was unable to be established.
+//
+// NOTE: This function may block for up to 15-seconds as it will not return
+// until the new connection is detected as being known to both nodes.
 func (n *networkHarness) ConnectNodes(ctx context.Context, a, b *lightningNode) error {
 	bobInfo, err := b.GetInfo(ctx, &lnrpc.GetInfoRequest{})
 	if err != nil {
@@ -799,8 +803,32 @@ func (n *networkHarness) ConnectNodes(ctx context.Context, a, b *lightningNode)
 			Host:   b.p2pAddr,
 		},
 	}
-	_, err = a.ConnectPeer(ctx, req)
-	return err
+	if _, err := a.ConnectPeer(ctx, req); err != nil {
+		return err
+	}
+
+	timeout := time.After(time.Second * 15)
+	for {
+
+		select {
+		case <-timeout:
+			return fmt.Errorf("peers not connected within 15 seconds")
+		default:
+		}
+
+		// If node B is seen in the ListPeers response from node A,
+		// then we can exit early as the connection has been fully
+		// established.
+		resp, err := a.ListPeers(ctx, &lnrpc.ListPeersRequest{})
+		if err != nil {
+			return err
+		}
+		for _, peer := range resp.Peers {
+			if peer.PubKey == b.PubKeyStr {
+				return nil
+			}
+		}
+	}
 }
 
 // RestartNode  attempts to restart a lightning node by shutting it down