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 }