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.
This commit is contained in:
Wilmer Paulino 2018-08-06 18:17:38 -07:00
parent dd5b6394d9
commit 956acdfa14
No known key found for this signature in database
GPG Key ID: 6DF57B9F9514972F

@ -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
return resp.amt, resp.numMore, resp.needMore
select {
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
return resp, nil
select {
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.