network: fix data race in TestHandleMPTData, fix #2241

This commit is contained in:
Roman Khimov 2021-11-15 12:37:01 +03:00
parent cc4a63e454
commit 48de82d902

View file

@ -141,11 +141,7 @@ func TestServerRegisterPeer(t *testing.T) {
ps[i].netaddr.Port = i + 1 ps[i].netaddr.Port = i + 1
} }
ch := startWithChannel(s) startWithCleanup(t, s)
t.Cleanup(func() {
s.Shutdown()
<-ch
})
s.register <- ps[0] s.register <- ps[0]
require.Eventually(t, func() bool { return 1 == s.PeerCount() }, time.Second, time.Millisecond*10) 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 { } else {
s = newTestServer(t, srvCfg) s = newTestServer(t, srvCfg)
} }
startWithCleanup(t, s)
return s
}
func startWithCleanup(t *testing.T, s *Server) {
ch := startWithChannel(s) ch := startWithChannel(s)
t.Cleanup(func() { t.Cleanup(func() {
s.Shutdown() s.Shutdown()
<-ch <-ch
}) })
return s
} }
func TestBlock(t *testing.T) { func TestBlock(t *testing.T) {
@ -826,8 +826,8 @@ func TestHandleMPTData(t *testing.T) {
}) })
t.Run("good", func(t *testing.T) { t.Run("good", func(t *testing.T) {
s := startTestServer(t)
expected := [][]byte{{1, 2, 3}, {2, 3, 4}} expected := [][]byte{{1, 2, 3}, {2, 3, 4}}
s := newTestServer(t, ServerConfig{Port: 0, UserAgent: "/test/"})
s.chain.(*fakechain.FakeChain).P2PStateExchangeExtensions = true s.chain.(*fakechain.FakeChain).P2PStateExchangeExtensions = true
s.stateSync = &fakechain.FakeStateSync{ s.stateSync = &fakechain.FakeStateSync{
AddMPTNodesFunc: func(nodes [][]byte) error { AddMPTNodesFunc: func(nodes [][]byte) error {
@ -835,6 +835,7 @@ func TestHandleMPTData(t *testing.T) {
return nil return nil
}, },
} }
startWithCleanup(t, s)
p := newLocalPeer(t, s) p := newLocalPeer(t, s)
p.handshaked = true p.handshaked = true