From 48de82d902f9c758ead8364c60a87bd8b625722c Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Mon, 15 Nov 2021 12:37:01 +0300
Subject: [PATCH] network: fix data race in TestHandleMPTData, fix #2241

---
 pkg/network/server_test.go | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/pkg/network/server_test.go b/pkg/network/server_test.go
index 51e9fe633..6cc83d3fc 100644
--- a/pkg/network/server_test.go
+++ b/pkg/network/server_test.go
@@ -141,11 +141,7 @@ func TestServerRegisterPeer(t *testing.T) {
 		ps[i].netaddr.Port = i + 1
 	}
 
-	ch := startWithChannel(s)
-	t.Cleanup(func() {
-		s.Shutdown()
-		<-ch
-	})
+	startWithCleanup(t, s)
 
 	s.register <- ps[0]
 	require.Eventually(t, func() bool { return 1 == s.PeerCount() }, time.Second, time.Millisecond*10)
@@ -395,12 +391,16 @@ func startTestServer(t *testing.T, protocolCfg ...func(*config.ProtocolConfigura
 	} else {
 		s = newTestServer(t, srvCfg)
 	}
+	startWithCleanup(t, s)
+	return s
+}
+
+func startWithCleanup(t *testing.T, s *Server) {
 	ch := startWithChannel(s)
 	t.Cleanup(func() {
 		s.Shutdown()
 		<-ch
 	})
-	return s
 }
 
 func TestBlock(t *testing.T) {
@@ -826,8 +826,8 @@ func TestHandleMPTData(t *testing.T) {
 	})
 
 	t.Run("good", func(t *testing.T) {
-		s := startTestServer(t)
 		expected := [][]byte{{1, 2, 3}, {2, 3, 4}}
+		s := newTestServer(t, ServerConfig{Port: 0, UserAgent: "/test/"})
 		s.chain.(*fakechain.FakeChain).P2PStateExchangeExtensions = true
 		s.stateSync = &fakechain.FakeStateSync{
 			AddMPTNodesFunc: func(nodes [][]byte) error {
@@ -835,6 +835,7 @@ func TestHandleMPTData(t *testing.T) {
 				return nil
 			},
 		}
+		startWithCleanup(t, s)
 
 		p := newLocalPeer(t, s)
 		p.handshaked = true