refactor pex_reactor tests

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.

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())
 

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)
+	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())
 
 	// 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
+}