From 86e6d230f2cb22d817da5258f265e5044900f1b9 Mon Sep 17 00:00:00 2001
From: "Johan T. Halseth" <johanth@gmail.com>
Date: Thu, 6 Dec 2018 14:24:15 +0100
Subject: [PATCH] autopilot/agent: add attachment directive goroutine to wait
 group

---
 autopilot/agent.go      |  3 +++
 autopilot/agent_test.go | 22 +++++++++++++++++++---
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/autopilot/agent.go b/autopilot/agent.go
index 12761db3..a59602c0 100644
--- a/autopilot/agent.go
+++ b/autopilot/agent.go
@@ -545,7 +545,10 @@ func (a *Agent) controller() {
 			}
 			pendingConns[nodeID] = struct{}{}
 
+			a.wg.Add(1)
 			go func(directive AttachmentDirective) {
+				defer a.wg.Done()
+
 				// We'll start out by attempting to connect to
 				// the peer in order to begin the funding
 				// workflow.
diff --git a/autopilot/agent_test.go b/autopilot/agent_test.go
index 3330ffd3..98b03a57 100644
--- a/autopilot/agent_test.go
+++ b/autopilot/agent_test.go
@@ -1231,6 +1231,7 @@ func TestAgentSkipPendingConns(t *testing.T) {
 	const walletBalance = btcutil.SatoshiPerBitcoin * 6
 
 	connect := make(chan chan error)
+	quit := make(chan struct{})
 
 	// With the dependencies we created, we can now create the initial
 	// agent itself.
@@ -1243,9 +1244,19 @@ func TestAgentSkipPendingConns(t *testing.T) {
 		},
 		ConnectToPeer: func(*btcec.PublicKey, []net.Addr) (bool, error) {
 			errChan := make(chan error)
-			connect <- errChan
-			err := <-errChan
-			return false, err
+
+			select {
+			case connect <- errChan:
+			case <-quit:
+				return false, errors.New("quit")
+			}
+
+			select {
+			case err := <-errChan:
+				return false, err
+			case <-quit:
+				return false, errors.New("quit")
+			}
 		},
 		DisconnectPeer: func(*btcec.PublicKey) error {
 			return nil
@@ -1272,6 +1283,11 @@ func TestAgentSkipPendingConns(t *testing.T) {
 	}
 	defer agent.Stop()
 
+	// We must defer the closing of quit after the defer agent.Stop(), to
+	// make sure ConnectToPeer won't block preventing the agent from
+	// exiting.
+	defer close(quit)
+
 	// We'll send an initial "yes" response to advance the agent past its
 	// initial check. This will cause it to try to get directives from the
 	// graph.