diff --git a/autopilot/agent_test.go b/autopilot/agent_test.go index 9baf96eb..eb1009b3 100644 --- a/autopilot/agent_test.go +++ b/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 - 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.