github.com/prysmaticlabs/prysm@v1.4.4/beacon-chain/p2p/peers/peerdata/store_test.go (about)

     1  package peerdata_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  
     7  	"github.com/libp2p/go-libp2p-core/peer"
     8  	"github.com/prysmaticlabs/prysm/beacon-chain/p2p/peers/peerdata"
     9  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    10  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    11  )
    12  
    13  func TestStore_GetSetDelete(t *testing.T) {
    14  	ctx, cancel := context.WithCancel(context.Background())
    15  	defer cancel()
    16  
    17  	store := peerdata.NewStore(ctx, &peerdata.StoreConfig{
    18  		MaxPeers: 12,
    19  	})
    20  	assert.NotNil(t, store)
    21  	assert.Equal(t, 12, store.Config().MaxPeers)
    22  
    23  	// Check non-existent data.
    24  	pid := peer.ID("00001")
    25  	peerData, ok := store.PeerData(pid)
    26  	assert.Equal(t, false, ok)
    27  	assert.Equal(t, (*peerdata.PeerData)(nil), peerData)
    28  	assert.Equal(t, 0, len(store.Peers()))
    29  
    30  	// Associate data.
    31  	store.SetPeerData(pid, &peerdata.PeerData{
    32  		BadResponses:    3,
    33  		ProcessedBlocks: 42,
    34  	})
    35  	peerData, ok = store.PeerData(pid)
    36  	assert.Equal(t, true, ok)
    37  	assert.Equal(t, 3, peerData.BadResponses)
    38  	assert.Equal(t, uint64(42), peerData.ProcessedBlocks)
    39  	require.Equal(t, 1, len(store.Peers()))
    40  	peers := store.Peers()
    41  	_, ok = peers[pid]
    42  	require.Equal(t, true, ok)
    43  	assert.Equal(t, 3, peers[pid].BadResponses)
    44  	assert.Equal(t, uint64(42), peers[pid].ProcessedBlocks)
    45  
    46  	// Remove data from peer.
    47  	store.DeletePeerData(pid)
    48  	peerData, ok = store.PeerData(pid)
    49  	assert.Equal(t, false, ok)
    50  	assert.Equal(t, (*peerdata.PeerData)(nil), peerData)
    51  	assert.Equal(t, 0, len(store.Peers()))
    52  }
    53  
    54  func TestStore_PeerDataGetOrCreate(t *testing.T) {
    55  	ctx, cancel := context.WithCancel(context.Background())
    56  	defer cancel()
    57  
    58  	store := peerdata.NewStore(ctx, &peerdata.StoreConfig{
    59  		MaxPeers: 12,
    60  	})
    61  	assert.NotNil(t, store)
    62  	assert.Equal(t, 12, store.Config().MaxPeers)
    63  
    64  	pid := peer.ID("00001")
    65  	peerData, ok := store.PeerData(pid)
    66  	assert.Equal(t, false, ok)
    67  	assert.Equal(t, (*peerdata.PeerData)(nil), peerData)
    68  	assert.Equal(t, 0, len(store.Peers()))
    69  
    70  	peerData = store.PeerDataGetOrCreate(pid)
    71  	assert.NotNil(t, peerData)
    72  	assert.Equal(t, 0, peerData.BadResponses)
    73  	assert.Equal(t, uint64(0), peerData.ProcessedBlocks)
    74  	require.Equal(t, 1, len(store.Peers()))
    75  
    76  	// Method must be idempotent, check.
    77  	peerData = store.PeerDataGetOrCreate(pid)
    78  	assert.NotNil(t, peerData)
    79  	assert.Equal(t, 0, peerData.BadResponses)
    80  	assert.Equal(t, uint64(0), peerData.ProcessedBlocks)
    81  	require.Equal(t, 1, len(store.Peers()))
    82  }