From fce665ef06ded750335f530c637571a1eaf5a585 Mon Sep 17 00:00:00 2001
From: Sam Kleinman <garen@tychoish.com>
Date: Mon, 19 Apr 2021 18:02:21 -0400
Subject: [PATCH] p2p: fix network update test (#6361)

---
 p2p/p2ptest/network.go | 14 ++++++++++++++
 p2p/p2ptest/require.go |  1 +
 p2p/router_test.go     |  2 +-
 3 files changed, 16 insertions(+), 1 deletion(-)

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"),