github.com/status-im/status-go@v1.1.0/peers/cache_test.go (about)

     1  package peers
     2  
     3  import (
     4  	"net"
     5  	"testing"
     6  
     7  	"github.com/ethereum/go-ethereum/p2p/enode"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  	"github.com/stretchr/testify/require"
    11  	"github.com/syndtr/goleveldb/leveldb"
    12  	"github.com/syndtr/goleveldb/leveldb/storage"
    13  
    14  	"github.com/ethereum/go-ethereum/crypto"
    15  	"github.com/ethereum/go-ethereum/p2p/discv5"
    16  )
    17  
    18  func TestPeersRange(t *testing.T) {
    19  	peersDB, err := newInMemoryCache()
    20  	require.NoError(t, err)
    21  	topic := discv5.Topic("test")
    22  
    23  	peerCount := 3
    24  	peers, err := createDiscv5Peers(peerCount)
    25  	require.NoError(t, err)
    26  
    27  	addedPeerIDs := make(map[string]struct{})
    28  	for _, peer := range peers {
    29  		assert.NoError(t, peersDB.AddPeer(peer, topic))
    30  		addedPeerIDs[peer.String()] = struct{}{}
    31  	}
    32  
    33  	// check if all peers are cached
    34  	cachedNodes := peersDB.GetPeersRange(topic, peerCount)
    35  	require.Len(t, cachedNodes, peerCount)
    36  	// and returned peers are these we just created
    37  	for _, node := range cachedNodes {
    38  		_, ok := addedPeerIDs[node.String()]
    39  		assert.True(t, ok)
    40  	}
    41  
    42  	// test removing peers
    43  	pk, err := peers[0].ID.Pubkey()
    44  	require.NoError(t, err)
    45  	require.NoError(t, peersDB.RemovePeer(enode.PubkeyToIDV4(pk), topic))
    46  	cachedNodes = peersDB.GetPeersRange(topic, peerCount)
    47  	require.Len(t, cachedNodes, peerCount-1)
    48  }
    49  
    50  func TestMultipleTopics(t *testing.T) {
    51  	peersDB, err := newInMemoryCache()
    52  	require.NoError(t, err)
    53  	topics := []discv5.Topic{discv5.Topic("first"), discv5.Topic("second")}
    54  
    55  	peerCount := 3
    56  	peers, err := createDiscv5Peers(peerCount)
    57  	require.NoError(t, err)
    58  
    59  	for i := range topics {
    60  		for _, peer := range peers {
    61  			assert.NoError(t, peersDB.AddPeer(peer, topics[i]))
    62  		}
    63  	}
    64  
    65  	for i := range topics {
    66  		nodes := peersDB.GetPeersRange(topics[i], peerCount+1)
    67  		assert.Len(t, nodes, peerCount)
    68  	}
    69  }
    70  
    71  // newInMemoryCache creates a cache for tests
    72  func newInMemoryCache() (*Cache, error) {
    73  	memdb, err := leveldb.Open(storage.NewMemStorage(), nil)
    74  	if err != nil {
    75  		return nil, err
    76  	}
    77  	return NewCache(memdb), nil
    78  }
    79  
    80  func createDiscv5Peers(count int) ([]*discv5.Node, error) {
    81  	nodes := make([]*discv5.Node, count)
    82  
    83  	for i := 0; i < count; i++ {
    84  		id, err := crypto.GenerateKey()
    85  		if err != nil {
    86  			return nil, err
    87  		}
    88  
    89  		nodes[i] = discv5.NewNode(
    90  			discv5.PubkeyID(&id.PublicKey),
    91  			net.IPv4(10, 10, 0, byte(i)),
    92  			32311,
    93  			32311,
    94  		)
    95  	}
    96  
    97  	return nodes, nil
    98  }