status-go/peers/cache_test.go

99 lines
2.3 KiB
Go

package peers
import (
"net"
"testing"
"github.com/ethereum/go-ethereum/p2p/enode"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/syndtr/goleveldb/leveldb"
"github.com/syndtr/goleveldb/leveldb/storage"
"github.com/ethereum/go-ethereum/crypto"
"github.com/ethereum/go-ethereum/p2p/discv5"
)
func TestPeersRange(t *testing.T) {
peersDB, err := newInMemoryCache()
require.NoError(t, err)
topic := discv5.Topic("test")
peerCount := 3
peers, err := createDiscv5Peers(peerCount)
require.NoError(t, err)
addedPeerIDs := make(map[string]struct{})
for _, peer := range peers {
assert.NoError(t, peersDB.AddPeer(peer, topic))
addedPeerIDs[peer.String()] = struct{}{}
}
// check if all peers are cached
cachedNodes := peersDB.GetPeersRange(topic, peerCount)
require.Len(t, cachedNodes, peerCount)
// and returned peers are these we just created
for _, node := range cachedNodes {
_, ok := addedPeerIDs[node.String()]
assert.True(t, ok)
}
// test removing peers
pk, err := peers[0].ID.Pubkey()
require.NoError(t, err)
require.NoError(t, peersDB.RemovePeer(enode.PubkeyToIDV4(pk), topic))
cachedNodes = peersDB.GetPeersRange(topic, peerCount)
require.Len(t, cachedNodes, peerCount-1)
}
func TestMultipleTopics(t *testing.T) {
peersDB, err := newInMemoryCache()
require.NoError(t, err)
topics := []discv5.Topic{discv5.Topic("first"), discv5.Topic("second")}
peerCount := 3
peers, err := createDiscv5Peers(peerCount)
require.NoError(t, err)
for i := range topics {
for _, peer := range peers {
assert.NoError(t, peersDB.AddPeer(peer, topics[i]))
}
}
for i := range topics {
nodes := peersDB.GetPeersRange(topics[i], peerCount+1)
assert.Len(t, nodes, peerCount)
}
}
// newInMemoryCache creates a cache for tests
func newInMemoryCache() (*Cache, error) {
memdb, err := leveldb.Open(storage.NewMemStorage(), nil)
if err != nil {
return nil, err
}
return NewCache(memdb), nil
}
func createDiscv5Peers(count int) ([]*discv5.Node, error) {
nodes := make([]*discv5.Node, count)
for i := 0; i < count; i++ {
id, err := crypto.GenerateKey()
if err != nil {
return nil, err
}
nodes[i] = discv5.NewNode(
discv5.PubkeyID(&id.PublicKey),
net.IPv4(10, 10, 0, byte(i)),
32311,
32311,
)
}
return nodes, nil
}