diff --git a/p2p/pex/addrbook.go b/p2p/pex/addrbook.go index 95ad70fe7..3a3920486 100644 --- a/p2p/pex/addrbook.go +++ b/p2p/pex/addrbook.go @@ -139,6 +139,10 @@ func (a *addrBook) Wait() { a.wg.Wait() } +func (a *addrBook) FilePath() string { + return a.filePath +} + //------------------------------------------------------- // AddOurAddress one of our addresses. diff --git a/p2p/pex/addrbook_test.go b/p2p/pex/addrbook_test.go index 166d31847..4a8df7166 100644 --- a/p2p/pex/addrbook_test.go +++ b/p2p/pex/addrbook_test.go @@ -5,6 +5,7 @@ import ( "fmt" "io/ioutil" "math/rand" + "os" "testing" "github.com/stretchr/testify/assert" @@ -26,17 +27,24 @@ func createTempFileName(prefix string) string { return fname } +func deleteTempFile(fname string) { + err := os.Remove(fname) + if err != nil { + panic(err) + } +} + func TestAddrBookPickAddress(t *testing.T) { - assert := assert.New(t) fname := createTempFileName("addrbook_test") + defer deleteTempFile(fname) // 0 addresses book := NewAddrBook(fname, true) book.SetLogger(log.TestingLogger()) - assert.Zero(book.Size()) + assert.Zero(t, book.Size()) addr := book.PickAddress(50) - assert.Nil(addr, "expected no address") + assert.Nil(t, addr, "expected no address") randAddrs := randNetAddressPairs(t, 1) addrSrc := randAddrs[0] @@ -44,26 +52,27 @@ func TestAddrBookPickAddress(t *testing.T) { // pick an address when we only have new address addr = book.PickAddress(0) - assert.NotNil(addr, "expected an address") + assert.NotNil(t, addr, "expected an address") addr = book.PickAddress(50) - assert.NotNil(addr, "expected an address") + assert.NotNil(t, addr, "expected an address") addr = book.PickAddress(100) - assert.NotNil(addr, "expected an address") + assert.NotNil(t, addr, "expected an address") // pick an address when we only have old address book.MarkGood(addrSrc.addr) addr = book.PickAddress(0) - assert.NotNil(addr, "expected an address") + assert.NotNil(t, addr, "expected an address") addr = book.PickAddress(50) - assert.NotNil(addr, "expected an address") + assert.NotNil(t, addr, "expected an address") // in this case, nNew==0 but we biased 100% to new, so we return nil addr = book.PickAddress(100) - assert.Nil(addr, "did not expected an address") + assert.Nil(t, addr, "did not expected an address") } func TestAddrBookSaveLoad(t *testing.T) { fname := createTempFileName("addrbook_test") + defer deleteTempFile(fname) // 0 addresses book := NewAddrBook(fname, true) @@ -95,6 +104,7 @@ func TestAddrBookSaveLoad(t *testing.T) { func TestAddrBookLookup(t *testing.T) { fname := createTempFileName("addrbook_test") + defer deleteTempFile(fname) randAddrs := randNetAddressPairs(t, 100) @@ -115,8 +125,8 @@ func TestAddrBookLookup(t *testing.T) { } func TestAddrBookPromoteToOld(t *testing.T) { - assert := assert.New(t) fname := createTempFileName("addrbook_test") + defer deleteTempFile(fname) randAddrs := randNetAddressPairs(t, 100) @@ -147,11 +157,12 @@ func TestAddrBookPromoteToOld(t *testing.T) { t.Errorf("selection could not be bigger than the book") } - assert.Equal(book.Size(), 100, "expecting book size to be 100") + assert.Equal(t, book.Size(), 100, "expecting book size to be 100") } func TestAddrBookHandlesDuplicates(t *testing.T) { fname := createTempFileName("addrbook_test") + defer deleteTempFile(fname) book := NewAddrBook(fname, true) book.SetLogger(log.TestingLogger()) @@ -202,6 +213,8 @@ func randIPv4Address(t *testing.T) *p2p.NetAddress { func TestAddrBookRemoveAddress(t *testing.T) { fname := createTempFileName("addrbook_test") + defer deleteTempFile(fname) + book := NewAddrBook(fname, true) book.SetLogger(log.TestingLogger()) diff --git a/p2p/pex/pex_reactor_test.go b/p2p/pex/pex_reactor_test.go index 41da867ad..6b610f009 100644 --- a/p2p/pex/pex_reactor_test.go +++ b/p2p/pex/pex_reactor_test.go @@ -4,6 +4,7 @@ import ( "fmt" "io/ioutil" "os" + "path/filepath" "testing" "time" @@ -30,49 +31,33 @@ func init() { } func TestPEXReactorBasic(t *testing.T) { - assert, require := assert.New(t), require.New(t) + r, book := createReactor(&PEXReactorConfig{}) + defer teardownReactor(book) - dir, err := ioutil.TempDir("", "pex_reactor") - require.Nil(err) - defer os.RemoveAll(dir) // nolint: errcheck - book := NewAddrBook(dir+"addrbook.json", true) - book.SetLogger(log.TestingLogger()) - - r := NewPEXReactor(book, &PEXReactorConfig{}) - r.SetLogger(log.TestingLogger()) - - assert.NotNil(r) - assert.NotEmpty(r.GetChannels()) + assert.NotNil(t, r) + assert.NotEmpty(t, r.GetChannels()) } func TestPEXReactorAddRemovePeer(t *testing.T) { - assert, require := assert.New(t), require.New(t) - - dir, err := ioutil.TempDir("", "pex_reactor") - require.Nil(err) - defer os.RemoveAll(dir) // nolint: errcheck - book := NewAddrBook(dir+"addrbook.json", true) - book.SetLogger(log.TestingLogger()) - - r := NewPEXReactor(book, &PEXReactorConfig{}) - r.SetLogger(log.TestingLogger()) + r, book := createReactor(&PEXReactorConfig{}) + defer teardownReactor(book) size := book.Size() peer := p2p.CreateRandomPeer(false) r.AddPeer(peer) - assert.Equal(size+1, book.Size()) + assert.Equal(t, size+1, book.Size()) r.RemovePeer(peer, "peer not available") - assert.Equal(size+1, book.Size()) + assert.Equal(t, size+1, book.Size()) outboundPeer := p2p.CreateRandomPeer(true) r.AddPeer(outboundPeer) - assert.Equal(size+1, book.Size(), "outbound peers should not be added to the address book") + assert.Equal(t, size+1, book.Size(), "outbound peers should not be added to the address book") r.RemovePeer(outboundPeer, "peer not available") - assert.Equal(size+1, book.Size()) + assert.Equal(t, size+1, book.Size()) } func TestPEXReactorRunning(t *testing.T) { @@ -82,7 +67,7 @@ func TestPEXReactorRunning(t *testing.T) { dir, err := ioutil.TempDir("", "pex_reactor") require.Nil(t, err) defer os.RemoveAll(dir) // nolint: errcheck - book := NewAddrBook(dir+"addrbook.json", false) + book := NewAddrBook(filepath.Join(dir, "addrbook.json"), false) book.SetLogger(log.TestingLogger()) // create switches @@ -120,16 +105,8 @@ func TestPEXReactorRunning(t *testing.T) { } func TestPEXReactorReceive(t *testing.T) { - assert, require := assert.New(t), require.New(t) - - dir, err := ioutil.TempDir("", "pex_reactor") - require.Nil(err) - defer os.RemoveAll(dir) // nolint: errcheck - book := NewAddrBook(dir+"addrbook.json", false) - book.SetLogger(log.TestingLogger()) - - r := NewPEXReactor(book, &PEXReactorConfig{}) - r.SetLogger(log.TestingLogger()) + r, book := createReactor(&PEXReactorConfig{}) + defer teardownReactor(book) peer := p2p.CreateRandomPeer(false) @@ -140,98 +117,77 @@ func TestPEXReactorReceive(t *testing.T) { addrs := []*p2p.NetAddress{peer.NodeInfo().NetAddress()} msg := wire.BinaryBytes(struct{ PexMessage }{&pexAddrsMessage{Addrs: addrs}}) r.Receive(PexChannel, peer, msg) - assert.Equal(size+1, book.Size()) + assert.Equal(t, size+1, book.Size()) msg = wire.BinaryBytes(struct{ PexMessage }{&pexRequestMessage{}}) r.Receive(PexChannel, peer, msg) } func TestPEXReactorRequestMessageAbuse(t *testing.T) { - assert, require := assert.New(t), require.New(t) + r, book := createReactor(&PEXReactorConfig{}) + defer teardownReactor(book) - dir, err := ioutil.TempDir("", "pex_reactor") - require.Nil(err) - defer os.RemoveAll(dir) // nolint: errcheck - book := NewAddrBook(dir+"addrbook.json", true) - book.SetLogger(log.TestingLogger()) - - r := NewPEXReactor(book, &PEXReactorConfig{}) - sw := p2p.MakeSwitch(config, 0, "127.0.0.1", "123.123.123", func(i int, sw *p2p.Switch) *p2p.Switch { return sw }) - sw.SetLogger(log.TestingLogger()) - sw.AddReactor("PEX", r) - r.SetSwitch(sw) - r.SetLogger(log.TestingLogger()) + sw := createSwitchAndAddReactors(r) peer := newMockPeer() p2p.AddPeerToSwitch(sw, peer) - assert.True(sw.Peers().Has(peer.ID())) + assert.True(t, sw.Peers().Has(peer.ID())) id := string(peer.ID()) msg := wire.BinaryBytes(struct{ PexMessage }{&pexRequestMessage{}}) // first time creates the entry r.Receive(PexChannel, peer, msg) - assert.True(r.lastReceivedRequests.Has(id)) - assert.True(sw.Peers().Has(peer.ID())) + assert.True(t, r.lastReceivedRequests.Has(id)) + assert.True(t, sw.Peers().Has(peer.ID())) // next time sets the last time value r.Receive(PexChannel, peer, msg) - assert.True(r.lastReceivedRequests.Has(id)) - assert.True(sw.Peers().Has(peer.ID())) + assert.True(t, r.lastReceivedRequests.Has(id)) + assert.True(t, sw.Peers().Has(peer.ID())) // third time is too many too soon - peer is removed r.Receive(PexChannel, peer, msg) - assert.False(r.lastReceivedRequests.Has(id)) - assert.False(sw.Peers().Has(peer.ID())) + assert.False(t, r.lastReceivedRequests.Has(id)) + assert.False(t, sw.Peers().Has(peer.ID())) } func TestPEXReactorAddrsMessageAbuse(t *testing.T) { - assert, require := assert.New(t), require.New(t) - - dir, err := ioutil.TempDir("", "pex_reactor") - require.Nil(err) - defer os.RemoveAll(dir) // nolint: errcheck - book := NewAddrBook(dir+"addrbook.json", true) - book.SetLogger(log.TestingLogger()) + r, book := createReactor(&PEXReactorConfig{}) + defer teardownReactor(book) - r := NewPEXReactor(book, &PEXReactorConfig{}) - sw := p2p.MakeSwitch(config, 0, "127.0.0.1", "123.123.123", func(i int, sw *p2p.Switch) *p2p.Switch { return sw }) - sw.SetLogger(log.TestingLogger()) - sw.AddReactor("PEX", r) - r.SetSwitch(sw) - r.SetLogger(log.TestingLogger()) + sw := createSwitchAndAddReactors(r) peer := newMockPeer() p2p.AddPeerToSwitch(sw, peer) - assert.True(sw.Peers().Has(peer.ID())) + assert.True(t, sw.Peers().Has(peer.ID())) id := string(peer.ID()) // request addrs from the peer r.RequestAddrs(peer) - assert.True(r.requestsSent.Has(id)) - assert.True(sw.Peers().Has(peer.ID())) + assert.True(t, r.requestsSent.Has(id)) + assert.True(t, sw.Peers().Has(peer.ID())) addrs := []*p2p.NetAddress{peer.NodeInfo().NetAddress()} msg := wire.BinaryBytes(struct{ PexMessage }{&pexAddrsMessage{Addrs: addrs}}) // receive some addrs. should clear the request r.Receive(PexChannel, peer, msg) - assert.False(r.requestsSent.Has(id)) - assert.True(sw.Peers().Has(peer.ID())) + assert.False(t, r.requestsSent.Has(id)) + assert.True(t, sw.Peers().Has(peer.ID())) // receiving more addrs causes a disconnect r.Receive(PexChannel, peer, msg) - assert.False(sw.Peers().Has(peer.ID())) + assert.False(t, sw.Peers().Has(peer.ID())) } func TestPEXReactorUsesSeedsIfNeeded(t *testing.T) { - dir, err := ioutil.TempDir("", "pex_reactor") require.Nil(t, err) defer os.RemoveAll(dir) // nolint: errcheck - book := NewAddrBook(dir+"addrbook.json", false) + book := NewAddrBook(filepath.Join(dir, "addrbook.json"), false) book.SetLogger(log.TestingLogger()) // 1. create seed @@ -288,22 +244,11 @@ func TestPEXReactorUsesSeedsIfNeeded(t *testing.T) { } func TestPEXReactorCrawlStatus(t *testing.T) { - assert, require := assert.New(t), require.New(t) + pexR, book := createReactor(&PEXReactorConfig{SeedMode: true}) + defer teardownReactor(book) - dir, err := ioutil.TempDir("", "pex_reactor") - require.Nil(err) - defer os.RemoveAll(dir) // nolint: errcheck - book := NewAddrBook(dir+"addrbook.json", false) - book.SetLogger(log.TestingLogger()) - - pexR := NewPEXReactor(book, &PEXReactorConfig{SeedMode: true}) // Seed/Crawler mode uses data from the Switch - p2p.MakeSwitch(config, 0, "127.0.0.1", "123.123.123", func(i int, sw *p2p.Switch) *p2p.Switch { - pexR.SetLogger(log.TestingLogger()) - sw.SetLogger(log.TestingLogger().With("switch", i)) - sw.AddReactor("pex", pexR) - return sw - }) + _ = createSwitchAndAddReactors(pexR) // Create a peer, add it to the peer set and the addrbook. peer := p2p.CreateRandomPeer(false) @@ -319,7 +264,7 @@ func TestPEXReactorCrawlStatus(t *testing.T) { peerInfos := pexR.getPeersToCrawl() // Make sure it has the proper number of elements - assert.Equal(2, len(peerInfos)) + assert.Equal(t, 2, len(peerInfos)) // TODO: test } @@ -400,3 +345,33 @@ func assertPeersWithTimeout( } } } + +func createReactor(config *PEXReactorConfig) (r *PEXReactor, book *addrBook) { + dir, err := ioutil.TempDir("", "pex_reactor") + if err != nil { + panic(err) + } + book = NewAddrBook(filepath.Join(dir, "addrbook.json"), true) + book.SetLogger(log.TestingLogger()) + + r = NewPEXReactor(book, &PEXReactorConfig{}) + r.SetLogger(log.TestingLogger()) + return +} + +func teardownReactor(book *addrBook) { + err := os.RemoveAll(filepath.Dir(book.FilePath())) + if err != nil { + panic(err) + } +} + +func createSwitchAndAddReactors(reactors ...p2p.Reactor) *p2p.Switch { + sw := p2p.MakeSwitch(config, 0, "127.0.0.1", "123.123.123", func(i int, sw *p2p.Switch) *p2p.Switch { return sw }) + sw.SetLogger(log.TestingLogger()) + for _, r := range reactors { + sw.AddReactor(r.String(), r) + r.SetSwitch(sw) + } + return sw +}