autopilot/agent_test: pass in agent's quit chan to heuristic

In this commit, we alter our mock heuristic to also take in a quit chan.
It's possible that at the end of a test the agent is blocked on a
NeedMoreChans/Select call as their mock implementations use channels. To
prevent this, we use the agent's quit chan so that the heuristic can
safely exit once the agent does.

autopilot/agent_test.go

@@ -2,13 +2,12 @@ package autopilot
 
 import (
 	"bytes"
+	"errors"
 	"net"
 	"sync"
 	"testing"
 	"time"
 
-	"errors"
-
 	"github.com/btcsuite/btcd/btcec"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcutil"
@@ -31,21 +30,32 @@ type mockHeuristic struct {
 
 	directiveResps chan []AttachmentDirective
 	directiveArgs  chan directiveArg
+
+	quit chan struct{}
 }
 
 func (m *mockHeuristic) NeedMoreChans(chans []Channel,
 	balance btcutil.Amount) (btcutil.Amount, uint32, bool) {
 
 	if m.moreChanArgs != nil {
-		m.moreChanArgs <- moreChanArg{
+		moreChan := moreChanArg{
 			chans:   chans,
 			balance: balance,
 		}
 
+		select {
+		case m.moreChanArgs <- moreChan:
+		case <-m.quit:
+			return 0, 0, false
+		}
 	}
 
-	resp := <-m.moreChansResps
+	select {
-	return resp.amt, resp.numMore, resp.needMore
+	case resp := <-m.moreChansResps:
+		return resp.amt, resp.numMore, resp.needMore
+	case <-m.quit:
+		return 0, 0, false
+	}
 }
 
 type directiveArg struct {
@@ -60,16 +70,26 @@ func (m *mockHeuristic) Select(self *btcec.PublicKey, graph ChannelGraph,
 	skipChans map[NodeID]struct{}) ([]AttachmentDirective, error) {
 
 	if m.directiveArgs != nil {
-		m.directiveArgs <- directiveArg{
+		directive := directiveArg{
 			self:  self,
 			graph: graph,
 			amt:   amtToUse,
 			skip:  skipChans,
 		}
+
+		select {
+		case m.directiveArgs <- directive:
+		case <-m.quit:
+			return nil, errors.New("exiting")
+		}
 	}
 
-	resp := <-m.directiveResps
+	select {
-	return resp, nil
+	case resp := <-m.directiveResps:
+		return resp, nil
+	case <-m.quit:
+		return nil, errors.New("exiting")
+	}
 }
 
 var _ AttachmentHeuristic = (*mockHeuristic)(nil)
@@ -155,6 +175,10 @@ func TestAgentChannelOpenSignal(t *testing.T) {
 		t.Fatalf("unable to create agent: %v", err)
 	}
 
+	// To ensure the heuristic doesn't block on quitting the agent, we'll
+	// use the agent's quit chan to signal when it should also stop.
+	heuristic.quit = agent.quit
+
 	// With the autopilot agent and all its dependencies we'll star the
 	// primary controller goroutine.
 	if err := agent.Start(); err != nil {
@@ -296,6 +320,10 @@ func TestAgentChannelFailureSignal(t *testing.T) {
 		t.Fatalf("unable to create agent: %v", err)
 	}
 
+	// To ensure the heuristic doesn't block on quitting the agent, we'll
+	// use the agent's quit chan to signal when it should also stop.
+	heuristic.quit = agent.quit
+
 	// With the autopilot agent and all its dependencies we'll start the
 	// primary controller goroutine.
 	if err := agent.Start(); err != nil {
@@ -402,6 +430,10 @@ func TestAgentChannelCloseSignal(t *testing.T) {
 		t.Fatalf("unable to create agent: %v", err)
 	}
 
+	// To ensure the heuristic doesn't block on quitting the agent, we'll
+	// use the agent's quit chan to signal when it should also stop.
+	heuristic.quit = agent.quit
+
 	// With the autopilot agent and all its dependencies we'll star the
 	// primary controller goroutine.
 	if err := agent.Start(); err != nil {
@@ -521,6 +553,10 @@ func TestAgentBalanceUpdate(t *testing.T) {
 		t.Fatalf("unable to create agent: %v", err)
 	}
 
+	// To ensure the heuristic doesn't block on quitting the agent, we'll
+	// use the agent's quit chan to signal when it should also stop.
+	heuristic.quit = agent.quit
+
 	// With the autopilot agent and all its dependencies we'll star the
 	// primary controller goroutine.
 	if err := agent.Start(); err != nil {
@@ -643,6 +679,10 @@ func TestAgentImmediateAttach(t *testing.T) {
 		t.Fatalf("unable to create agent: %v", err)
 	}
 
+	// To ensure the heuristic doesn't block on quitting the agent, we'll
+	// use the agent's quit chan to signal when it should also stop.
+	heuristic.quit = agent.quit
+
 	// With the autopilot agent and all its dependencies we'll star the
 	// primary controller goroutine.
 	if err := agent.Start(); err != nil {
@@ -781,6 +821,10 @@ func TestAgentPrivateChannels(t *testing.T) {
 		t.Fatalf("unable to create agent: %v", err)
 	}
 
+	// To ensure the heuristic doesn't block on quitting the agent, we'll
+	// use the agent's quit chan to signal when it should also stop.
+	heuristic.quit = agent.quit
+
 	// With the autopilot agent and all its dependencies we'll star the
 	// primary controller goroutine.
 	if err := agent.Start(); err != nil {
@@ -914,6 +958,10 @@ func TestAgentPendingChannelState(t *testing.T) {
 		t.Fatalf("unable to create agent: %v", err)
 	}
 
+	// To ensure the heuristic doesn't block on quitting the agent, we'll
+	// use the agent's quit chan to signal when it should also stop.
+	heuristic.quit = agent.quit
+
 	// With the autopilot agent and all its dependencies we'll start the
 	// primary controller goroutine.
 	if err := agent.Start(); err != nil {
@@ -966,7 +1014,6 @@ func TestAgentPendingChannelState(t *testing.T) {
 	}
 	select {
 	case heuristic.directiveResps <- []AttachmentDirective{nodeDirective}:
-		return
 	case <-time.After(time.Second * 10):
 		t.Fatalf("heuristic wasn't queried in time")
 	}
@@ -994,8 +1041,6 @@ func TestAgentPendingChannelState(t *testing.T) {
 	// heuristic.
 	agent.OnBalanceChange(0.4 * btcutil.SatoshiPerBitcoin)
 
-	wg = sync.WaitGroup{}
-
 	// The heuristic should be queried, and the argument for the set of
 	// channels passed in should include the pending channels that
 	// should've been created above.