github.com/ava-labs/avalanchego@v1.11.11/network/peer/set_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package peer
     5  
     6  import (
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/ava-labs/avalanchego/ids"
    12  	"github.com/ava-labs/avalanchego/utils/constants"
    13  )
    14  
    15  func TestSet(t *testing.T) {
    16  	require := require.New(t)
    17  
    18  	set := NewSet()
    19  
    20  	peer1 := &peer{
    21  		id:              ids.BuildTestNodeID([]byte{0x01}),
    22  		observedUptimes: map[ids.ID]uint32{constants.PrimaryNetworkID: 0},
    23  	}
    24  	updatedPeer1 := &peer{
    25  		id:              ids.BuildTestNodeID([]byte{0x01}),
    26  		observedUptimes: map[ids.ID]uint32{constants.PrimaryNetworkID: 1},
    27  	}
    28  	peer2 := &peer{
    29  		id: ids.BuildTestNodeID([]byte{0x02}),
    30  	}
    31  	unknownPeer := &peer{
    32  		id: ids.BuildTestNodeID([]byte{0xff}),
    33  	}
    34  	peer3 := &peer{
    35  		id: ids.BuildTestNodeID([]byte{0x03}),
    36  	}
    37  	peer4 := &peer{
    38  		id: ids.BuildTestNodeID([]byte{0x04}),
    39  	}
    40  
    41  	// add of first peer is handled
    42  	set.Add(peer1)
    43  	retrievedPeer1, peer1Found := set.GetByID(peer1.id)
    44  	require.True(peer1Found)
    45  	observed1, _ := peer1.ObservedUptime(constants.PrimaryNetworkID)
    46  	observed2, _ := retrievedPeer1.ObservedUptime(constants.PrimaryNetworkID)
    47  	require.Equal(observed1, observed2)
    48  	require.Equal(1, set.Len())
    49  
    50  	// re-addition of peer works as update
    51  	set.Add(updatedPeer1)
    52  	retrievedPeer1, peer1Found = set.GetByID(peer1.id)
    53  	require.True(peer1Found)
    54  	observed1, _ = updatedPeer1.ObservedUptime(constants.PrimaryNetworkID)
    55  	observed2, _ = retrievedPeer1.ObservedUptime(constants.PrimaryNetworkID)
    56  	require.Equal(observed1, observed2)
    57  	require.Equal(1, set.Len())
    58  
    59  	// add of another peer is handled
    60  	set.Add(peer2)
    61  	retrievedPeer2, peer2Found := set.GetByID(peer2.id)
    62  	require.True(peer2Found)
    63  	observed1, _ = peer2.ObservedUptime(constants.PrimaryNetworkID)
    64  	observed2, _ = retrievedPeer2.ObservedUptime(constants.PrimaryNetworkID)
    65  	require.Equal(observed1, observed2)
    66  	require.Equal(2, set.Len())
    67  
    68  	// removal of added peer is handled
    69  	set.Remove(peer1.id)
    70  	_, peer1Found = set.GetByID(peer1.id)
    71  	require.False(peer1Found)
    72  	retrievedPeer2, peer2Found = set.GetByID(peer2.id)
    73  	require.True(peer2Found)
    74  	require.Equal(peer2.id, retrievedPeer2.ID())
    75  	require.Equal(1, set.Len())
    76  
    77  	// query for unknown peer is handled
    78  	_, unknownPeerfound := set.GetByID(unknownPeer.id)
    79  	require.False(unknownPeerfound)
    80  
    81  	// removal of unknown peer is handled
    82  	set.Remove(unknownPeer.id)
    83  	retrievedPeer2, peer2Found = set.GetByID(peer2.id)
    84  	require.True(peer2Found)
    85  	require.Equal(peer2.id, retrievedPeer2.ID())
    86  	require.Equal(1, set.Len())
    87  
    88  	// retrival by inbound index is handled
    89  	set.Add(peer3)
    90  	set.Add(peer4)
    91  	require.Equal(3, set.Len())
    92  
    93  	thirdPeer, ok := set.GetByIndex(1)
    94  	require.True(ok)
    95  	require.Equal(peer3.id, thirdPeer.ID())
    96  
    97  	// retrival by out-of-bounds index is handled
    98  	_, ok = set.GetByIndex(3)
    99  	require.False(ok)
   100  }
   101  
   102  func TestSetSample(t *testing.T) {
   103  	require := require.New(t)
   104  
   105  	set := NewSet()
   106  
   107  	peer1 := &peer{
   108  		id: ids.BuildTestNodeID([]byte{0x01}),
   109  	}
   110  	peer2 := &peer{
   111  		id: ids.BuildTestNodeID([]byte{0x02}),
   112  	}
   113  
   114  	// Case: Empty
   115  	peers := set.Sample(0, NoPrecondition)
   116  	require.Empty(peers)
   117  
   118  	peers = set.Sample(-1, NoPrecondition)
   119  	require.Empty(peers)
   120  
   121  	peers = set.Sample(1, NoPrecondition)
   122  	require.Empty(peers)
   123  
   124  	// Case: 1 peer
   125  	set.Add(peer1)
   126  
   127  	peers = set.Sample(0, NoPrecondition)
   128  	require.Empty(peers)
   129  
   130  	peers = set.Sample(1, NoPrecondition)
   131  	require.Equal([]Peer{peer1}, peers)
   132  
   133  	peers = set.Sample(2, NoPrecondition)
   134  	require.Equal([]Peer{peer1}, peers)
   135  
   136  	// Case: 2 peers
   137  	set.Add(peer2)
   138  
   139  	peers = set.Sample(1, NoPrecondition)
   140  	require.Len(peers, 1)
   141  }