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  }