package p2p import ( "fmt" "io/ioutil" "math/rand" "testing" "github.com/stretchr/testify/assert" ) func createTempFileName(prefix string) string { f, err := ioutil.TempFile("", prefix) if err != nil { panic(err) } fname := f.Name() err = f.Close() if err != nil { panic(err) } return fname } func TestAddrBookSaveLoad(t *testing.T) { fname := createTempFileName("addrbook_test") // 0 addresses book := NewAddrBook(fname, true) book.saveToFile(fname) book = NewAddrBook(fname, true) book.loadFromFile(fname) assert.Zero(t, book.Size()) // 100 addresses randAddrs := randNetAddressPairs(t, 100) for _, addrSrc := range randAddrs { book.AddAddress(addrSrc.addr, addrSrc.src) } assert.Equal(t, 100, book.Size()) book.saveToFile(fname) book = NewAddrBook(fname, true) book.loadFromFile(fname) assert.Equal(t, 100, book.Size()) } func TestAddrBookLookup(t *testing.T) { fname := createTempFileName("addrbook_test") randAddrs := randNetAddressPairs(t, 100) book := NewAddrBook(fname, true) for _, addrSrc := range randAddrs { addr := addrSrc.addr src := addrSrc.src book.AddAddress(addr, src) ka := book.addrLookup[addr.String()] assert.NotNil(t, ka, "Expected to find KnownAddress %v but wasn't there.", addr) if !(ka.Addr.Equals(addr) && ka.Src.Equals(src)) { t.Fatalf("KnownAddress doesn't match addr & src") } } } func TestAddrBookPromoteToOld(t *testing.T) { fname := createTempFileName("addrbook_test") randAddrs := randNetAddressPairs(t, 100) book := NewAddrBook(fname, true) for _, addrSrc := range randAddrs { book.AddAddress(addrSrc.addr, addrSrc.src) } // Attempt all addresses. for _, addrSrc := range randAddrs { book.MarkAttempt(addrSrc.addr) } // Promote half of them for i, addrSrc := range randAddrs { if i%2 == 0 { book.MarkGood(addrSrc.addr) } } // TODO: do more testing :) selection := book.GetSelection() t.Logf("selection: %v", selection) if len(selection) > book.Size() { t.Errorf("selection could not be bigger than the book") } } func TestAddrBookHandlesDuplicates(t *testing.T) { fname := createTempFileName("addrbook_test") book := NewAddrBook(fname, true) randAddrs := randNetAddressPairs(t, 100) differentSrc := randIPv4Address(t) for _, addrSrc := range randAddrs { book.AddAddress(addrSrc.addr, addrSrc.src) book.AddAddress(addrSrc.addr, addrSrc.src) // duplicate book.AddAddress(addrSrc.addr, differentSrc) // different src } assert.Equal(t, 100, book.Size()) } type netAddressPair struct { addr *NetAddress src *NetAddress } func randNetAddressPairs(t *testing.T, n int) []netAddressPair { randAddrs := make([]netAddressPair, n) for i := 0; i < n; i++ { randAddrs[i] = netAddressPair{addr: randIPv4Address(t), src: randIPv4Address(t)} } return randAddrs } func randIPv4Address(t *testing.T) *NetAddress { for { ip := fmt.Sprintf("%v.%v.%v.%v", rand.Intn(254)+1, rand.Intn(255), rand.Intn(255), rand.Intn(255), ) port := rand.Intn(65535-1) + 1 addr, err := NewNetAddressString(fmt.Sprintf("%v:%v", ip, port)) assert.Nil(t, err, "error generating rand network address") if addr.Routable() { return addr } } } func TestAddrBookRemoveAddress(t *testing.T) { fname := createTempFileName("addrbook_test") book := NewAddrBook(fname, true) addr := randIPv4Address(t) book.AddAddress(addr, addr) assert.Equal(t, 1, book.Size()) book.RemoveAddress(addr) assert.Equal(t, 0, book.Size()) nonExistingAddr := randIPv4Address(t) book.RemoveAddress(nonExistingAddr) assert.Equal(t, 0, book.Size()) }