diff --git a/p2p/p2ptest/network.go b/p2p/p2ptest/network.go index 74faadb36..47332a240 100644 --- a/p2p/p2ptest/network.go +++ b/p2p/p2ptest/network.go @@ -294,11 +294,25 @@ func (n *Node) MakeChannelNoCleanup( // MakePeerUpdates opens a peer update subscription, with automatic cleanup. // It checks that all updates have been consumed during cleanup. func (n *Node) MakePeerUpdates(t *testing.T) *p2p.PeerUpdates { + t.Helper() sub := n.PeerManager.Subscribe() t.Cleanup(func() { + t.Helper() RequireNoUpdates(t, sub) sub.Close() }) return sub } + +// MakePeerUpdatesNoRequireEmpty opens a peer update subscription, with automatic cleanup. +// It does *not* check that all updates have been consumed, but will +// close the update channel. +func (n *Node) MakePeerUpdatesNoRequireEmpty(t *testing.T) *p2p.PeerUpdates { + sub := n.PeerManager.Subscribe() + t.Cleanup(func() { + sub.Close() + }) + + return sub +} diff --git a/p2p/p2ptest/require.go b/p2p/p2ptest/require.go index 47deadf7f..cce88a42a 100644 --- a/p2p/p2ptest/require.go +++ b/p2p/p2ptest/require.go @@ -93,6 +93,7 @@ func RequireSendReceive( // RequireNoUpdates requires that a PeerUpdates subscription is empty. func RequireNoUpdates(t *testing.T, peerUpdates *p2p.PeerUpdates) { + t.Helper() select { case update := <-peerUpdates.Updates(): require.Fail(t, "unexpected peer updates", "got %v", update) diff --git a/p2p/router_test.go b/p2p/router_test.go index acc8fca05..6e6acc828 100644 --- a/p2p/router_test.go +++ b/p2p/router_test.go @@ -80,7 +80,7 @@ func TestRouter_Network(t *testing.T) { // We then submit an error for a peer, and watch it get disconnected and // then reconnected as the router retries it. - peerUpdates := local.MakePeerUpdates(t) + peerUpdates := local.MakePeerUpdatesNoRequireEmpty(t) channel.Error <- p2p.PeerError{ NodeID: peers[0].NodeID, Err: errors.New("boom"),