From 2154ec130f640b7b12869660ecd1185a2155affc Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Thu, 21 Dec 2017 11:32:54 +0100
Subject: [PATCH] lntest: add new LookUpNodeByPub method to lookup active nodes
 by pubkey

In this commit, we add a new helper function to the NetworkHarness
struct. This helper function serves to allow test authors to look up
pointer to an active node based on its current public key.

Each time a new node is started, its public key will be re-registered
within the global nodesByPub map.
---
 lntest/harness.go | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/lntest/harness.go b/lntest/harness.go
index dbb4d350..c5186955 100644
--- a/lntest/harness.go
+++ b/lntest/harness.go
@@ -32,6 +32,8 @@ type NetworkHarness struct {
 
 	activeNodes map[int]*HarnessNode
 
+	nodesByPub map[string]*HarnessNode
+
 	// Alice and Bob are the initial seeder nodes that are automatically
 	// created to be the initial participants of the test network.
 	Alice *HarnessNode
@@ -56,6 +58,7 @@ type NetworkHarness struct {
 func NewNetworkHarness(r *rpctest.Harness) (*NetworkHarness, error) {
 	n := NetworkHarness{
 		activeNodes:          make(map[int]*HarnessNode),
+		nodesByPub:           make(map[string]*HarnessNode),
 		seenTxns:             make(chan *chainhash.Hash),
 		bitcoinWatchRequests: make(chan *txWatchRequest),
 		lndErrorChan:         make(chan error),
@@ -68,6 +71,21 @@ func NewNetworkHarness(r *rpctest.Harness) (*NetworkHarness, error) {
 	return &n, nil
 }
 
+// LookUpNodeByPub queries the set of active nodes to locate a node according
+// to its public key. The second value will be true if the node was found, and
+// false otherwise.
+func (n *NetworkHarness) LookUpNodeByPub(pubStr string) (*HarnessNode, error) {
+	n.mtx.Lock()
+	defer n.mtx.Unlock()
+
+	node, ok := n.nodesByPub[pubStr]
+	if !ok {
+		return nil, fmt.Errorf("unable to find node")
+	}
+
+	return node, nil
+}
+
 // ProcessErrors returns a channel used for reporting any fatal process errors.
 // If any of the active nodes within the harness' test network incur a fatal
 // error, that error is sent over this channel.
@@ -201,6 +219,7 @@ out:
 
 // TearDownAll tears down all active nodes within the test lightning network.
 func (n *NetworkHarness) TearDownAll() error {
+
 	for _, node := range n.activeNodes {
 		if err := n.ShutdownNode(node); err != nil {
 			return err
@@ -236,6 +255,12 @@ func (n *NetworkHarness) NewNode(extraArgs []string) (*HarnessNode, error) {
 		return nil, err
 	}
 
+	// With the node started, we can now record its public key within the
+	// global mapping.
+	n.mtx.Lock()
+	n.nodesByPub[node.PubKeyStr] = node
+	n.mtx.Unlock()
+
 	return node, nil
 }