github.com/MetalBlockchain/metalgo@v1.11.9/snow/consensus/snowman/bootstrapper/sampler_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 bootstrapper 5 6 import ( 7 "math" 8 "testing" 9 10 "github.com/stretchr/testify/require" 11 12 "github.com/MetalBlockchain/metalgo/ids" 13 "github.com/MetalBlockchain/metalgo/utils/set" 14 15 safemath "github.com/MetalBlockchain/metalgo/utils/math" 16 ) 17 18 func TestSample(t *testing.T) { 19 tests := []struct { 20 name string 21 elements map[ids.NodeID]uint64 22 maxSize int 23 expectedSampled set.Set[ids.NodeID] 24 expectedErr error 25 }{ 26 { 27 name: "sample everything", 28 elements: map[ids.NodeID]uint64{ 29 nodeID0: 1, 30 nodeID1: 1, 31 }, 32 maxSize: 2, 33 expectedSampled: set.Of(nodeID0, nodeID1), 34 expectedErr: nil, 35 }, 36 { 37 name: "limit sample due to too few elements", 38 elements: map[ids.NodeID]uint64{ 39 nodeID0: 1, 40 }, 41 maxSize: 2, 42 expectedSampled: set.Of(nodeID0), 43 expectedErr: nil, 44 }, 45 { 46 name: "limit sample", 47 elements: map[ids.NodeID]uint64{ 48 nodeID0: math.MaxUint64 - 1, 49 nodeID1: 1, 50 }, 51 maxSize: 1, 52 expectedSampled: set.Of(nodeID0), 53 expectedErr: nil, 54 }, 55 { 56 name: "overflow", 57 elements: map[ids.NodeID]uint64{ 58 nodeID0: math.MaxUint64, 59 nodeID1: 1, 60 }, 61 maxSize: 1, 62 expectedSampled: nil, 63 expectedErr: safemath.ErrOverflow, 64 }, 65 } 66 for _, test := range tests { 67 t.Run(test.name, func(t *testing.T) { 68 require := require.New(t) 69 70 sampled, err := Sample(test.elements, test.maxSize) 71 require.ErrorIs(err, test.expectedErr) 72 require.Equal(test.expectedSampled, sampled) 73 }) 74 } 75 }