From 6a432c4be27f11c15fa66519423b96772bccf0e7 Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@mit.edu>
Date: Wed, 14 Mar 2018 02:18:32 -0700
Subject: [PATCH] lntest: adds EnsureConnected helper method

---
 lntest/harness.go | 62 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/lntest/harness.go b/lntest/harness.go
index 858a4cd6..971148d7 100644
--- a/lntest/harness.go
+++ b/lntest/harness.go
@@ -3,6 +3,7 @@ package lntest
 import (
 	"fmt"
 	"io/ioutil"
+	"strings"
 	"sync"
 	"time"
 
@@ -264,6 +265,67 @@ func (n *NetworkHarness) NewNode(extraArgs []string) (*HarnessNode, error) {
 	return node, nil
 }
 
+// EnsureConnected will try to connect to two nodes, returning no error if they
+// are already connected. If the nodes were not connected previously, this will
+// behave the same as ConnectNodes. If a pending connection request has already
+// been made, the method will block until the two nodes appear in each other's
+// peers list, or until the 15s timeout expires.
+func (n *NetworkHarness) EnsureConnected(ctx context.Context, a, b *HarnessNode) error {
+	bobInfo, err := b.GetInfo(ctx, &lnrpc.GetInfoRequest{})
+	if err != nil {
+		return err
+	}
+
+	req := &lnrpc.ConnectPeerRequest{
+		Addr: &lnrpc.LightningAddress{
+			Pubkey: bobInfo.IdentityPubkey,
+			Host:   b.cfg.P2PAddr(),
+		},
+	}
+
+	_, err = a.ConnectPeer(ctx, req)
+	switch {
+
+	// Request was successful, wait for both to display the connection.
+	case err == nil:
+
+	// If we already have pending connection, we will wait until bob appears
+	// in alice's peer list.
+	case strings.Contains(err.Error(), "connection attempt to ") &&
+		strings.Contains(err.Error(), " is pending"):
+
+	// If the two are already connected, we return early with no error.
+	case strings.Contains(err.Error(), "already connected to peer"):
+		return nil
+
+	default:
+		return err
+	}
+
+	err = WaitPredicate(func() bool {
+		// 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 false
+		}
+
+		for _, peer := range resp.Peers {
+			if peer.PubKey == b.PubKeyStr {
+				return true
+			}
+		}
+
+		return false
+	}, time.Second*15)
+	if err != nil {
+		return fmt.Errorf("peers not connected within 15 seconds")
+	}
+
+	return nil
+}
+
 // ConnectNodes establishes an encrypted+authenticated p2p connection from node
 // a towards node b. The function will return a non-nil error if the connection
 // was unable to be established.