mirror of https://github.com/status-im/go-waku.git
157 lines
3.7 KiB
Go
157 lines
3.7 KiB
Go
package peer_exchange
|
|
|
|
import (
|
|
"crypto/ecdsa"
|
|
"testing"
|
|
|
|
gcrypto "github.com/ethereum/go-ethereum/crypto"
|
|
"github.com/ethereum/go-ethereum/p2p/enode"
|
|
"github.com/ethereum/go-ethereum/p2p/enr"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/waku-org/go-waku/waku/v2/protocol"
|
|
wenr "github.com/waku-org/go-waku/waku/v2/protocol/enr"
|
|
)
|
|
|
|
func getEnode(t *testing.T, key *ecdsa.PrivateKey, cluster uint16, indexes ...uint16) *enode.Node {
|
|
if key == nil {
|
|
var err error
|
|
key, err = gcrypto.GenerateKey()
|
|
require.NoError(t, err)
|
|
}
|
|
myNode, err := wenr.NewLocalnode(key)
|
|
require.NoError(t, err)
|
|
if cluster != 0 {
|
|
shard, err := protocol.NewRelayShards(cluster, indexes...)
|
|
require.NoError(t, err)
|
|
myNode.Set(enr.WithEntry(wenr.ShardingBitVectorEnrField, shard.BitVector()))
|
|
}
|
|
return myNode.Node()
|
|
}
|
|
|
|
func TestLruMoreThanSize(t *testing.T) {
|
|
node1 := getEnode(t, nil, 1, 1)
|
|
node2 := getEnode(t, nil, 1, 1, 2)
|
|
node3 := getEnode(t, nil, 1, 1)
|
|
|
|
lru := newShardLRU(2)
|
|
|
|
err := lru.Add(node1)
|
|
require.NoError(t, err)
|
|
|
|
err = lru.Add(node2)
|
|
require.NoError(t, err)
|
|
|
|
err = lru.Add(node3)
|
|
require.NoError(t, err)
|
|
|
|
nodes := lru.GetRandomNodes(&ShardInfo{1, 1}, 1)
|
|
|
|
require.Equal(t, 1, len(nodes))
|
|
if nodes[0].ID() != node2.ID() && nodes[0].ID() != node3.ID() {
|
|
t.Fatalf("different node found %v", nodes)
|
|
}
|
|
// checks if removed node is deleted from lru or not
|
|
require.Nil(t, lru.Get(node1.ID()))
|
|
|
|
// node 2 is removed from lru for cluster 1/shard 1 but it is still being maintained for cluster1,index2
|
|
{
|
|
err = lru.Add(getEnode(t, nil, 1, 1)) // add two more nodes to 1/1 cluster shard
|
|
require.NoError(t, err)
|
|
|
|
err = lru.Add(getEnode(t, nil, 1, 1))
|
|
require.NoError(t, err)
|
|
|
|
}
|
|
|
|
// node2 still present in lru
|
|
require.Equal(t, node2, lru.Get(node2.ID()))
|
|
|
|
// now node2 is removed from all shards' cache
|
|
{
|
|
err = lru.Add(getEnode(t, nil, 1, 2)) // add two more nodes to 1/2 cluster shard
|
|
require.NoError(t, err)
|
|
|
|
err = lru.Add(getEnode(t, nil, 1, 2))
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// node2 still present in lru
|
|
require.Nil(t, lru.Get(node2.ID()))
|
|
}
|
|
|
|
func TestLruNodeWithNewSeq(t *testing.T) {
|
|
lru := newShardLRU(2)
|
|
//
|
|
key, err := gcrypto.GenerateKey()
|
|
require.NoError(t, err)
|
|
|
|
node1 := getEnode(t, key, 1, 1)
|
|
err = lru.Add(node1)
|
|
require.NoError(t, err)
|
|
|
|
node1 = getEnode(t, key, 1, 2, 3)
|
|
err = lru.Add(node1)
|
|
require.NoError(t, err)
|
|
|
|
//
|
|
nodes := lru.GetRandomNodes(&ShardInfo{1, 1}, 2)
|
|
require.Equal(t, 0, len(nodes))
|
|
//
|
|
nodes = lru.GetRandomNodes(&ShardInfo{1, 2}, 2)
|
|
require.Equal(t, 1, len(nodes))
|
|
//
|
|
nodes = lru.GetRandomNodes(&ShardInfo{1, 3}, 2)
|
|
require.Equal(t, 1, len(nodes))
|
|
}
|
|
|
|
func TestLruNoShard(t *testing.T) {
|
|
lru := newShardLRU(2)
|
|
|
|
node1 := getEnode(t, nil, 0)
|
|
node2 := getEnode(t, nil, 0)
|
|
|
|
err := lru.Add(node1)
|
|
require.NoError(t, err)
|
|
|
|
err = lru.Add(node2)
|
|
require.NoError(t, err)
|
|
|
|
// check returned nodes
|
|
require.Equal(t, 2, lru.len(nil))
|
|
for _, node := range lru.GetRandomNodes(nil, 2) {
|
|
if node.ID() != node1.ID() && node.ID() != node2.ID() {
|
|
t.Fatalf("different node found %v", node)
|
|
}
|
|
}
|
|
}
|
|
|
|
// checks if lru is able to handle nodes with/without shards together
|
|
func TestLruMixedNodes(t *testing.T) {
|
|
lru := newShardLRU(2)
|
|
|
|
node1 := getEnode(t, nil, 0)
|
|
err := lru.Add(node1)
|
|
require.NoError(t, err)
|
|
|
|
node2 := getEnode(t, nil, 1, 1)
|
|
err = lru.Add(node2)
|
|
require.NoError(t, err)
|
|
|
|
node3 := getEnode(t, nil, 1, 2)
|
|
err = lru.Add(node3)
|
|
require.NoError(t, err)
|
|
|
|
// check that default
|
|
require.Equal(t, 3, lru.len(nil))
|
|
|
|
//
|
|
nodes := lru.GetRandomNodes(&ShardInfo{1, 1}, 2)
|
|
require.Equal(t, 1, len(nodes))
|
|
require.Equal(t, node2.ID(), nodes[0].ID())
|
|
|
|
//
|
|
nodes = lru.GetRandomNodes(&ShardInfo{1, 2}, 2)
|
|
require.Equal(t, 1, len(nodes))
|
|
require.Equal(t, node3.ID(), nodes[0].ID())
|
|
}
|