github.com/MetalBlockchain/metalgo@v1.11.9/utils/sampler/weighted_without_replacement_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 sampler
     5  
     6  import (
     7  	"fmt"
     8  	"math"
     9  	"slices"
    10  	"testing"
    11  
    12  	"github.com/stretchr/testify/require"
    13  
    14  	safemath "github.com/MetalBlockchain/metalgo/utils/math"
    15  )
    16  
    17  var (
    18  	weightedWithoutReplacementSamplers = []struct {
    19  		name    string
    20  		sampler WeightedWithoutReplacement
    21  	}{
    22  		{
    23  			name: "generic with replacer and best",
    24  			sampler: &weightedWithoutReplacementGeneric{
    25  				u: &uniformReplacer{
    26  					rng: globalRNG,
    27  				},
    28  				w: &weightedBest{
    29  					samplers: []Weighted{
    30  						&weightedArray{},
    31  						&weightedHeap{},
    32  						&weightedUniform{
    33  							maxWeight: 1024,
    34  						},
    35  					},
    36  					benchmarkIterations: 30,
    37  				},
    38  			},
    39  		},
    40  	}
    41  	weightedWithoutReplacementTests = []struct {
    42  		name string
    43  		test func(*testing.T, WeightedWithoutReplacement)
    44  	}{
    45  		{
    46  			name: "initialize overflow",
    47  			test: WeightedWithoutReplacementInitializeOverflowTest,
    48  		},
    49  		{
    50  			name: "out of range",
    51  			test: WeightedWithoutReplacementOutOfRangeTest,
    52  		},
    53  		{
    54  			name: "empty without weight",
    55  			test: WeightedWithoutReplacementEmptyWithoutWeightTest,
    56  		},
    57  		{
    58  			name: "empty",
    59  			test: WeightedWithoutReplacementEmptyTest,
    60  		},
    61  		{
    62  			name: "singleton",
    63  			test: WeightedWithoutReplacementSingletonTest,
    64  		},
    65  		{
    66  			name: "with zero",
    67  			test: WeightedWithoutReplacementWithZeroTest,
    68  		},
    69  		{
    70  			name: "distribution",
    71  			test: WeightedWithoutReplacementDistributionTest,
    72  		},
    73  	}
    74  )
    75  
    76  func TestAllWeightedWithoutReplacement(t *testing.T) {
    77  	for _, s := range weightedWithoutReplacementSamplers {
    78  		for _, test := range weightedWithoutReplacementTests {
    79  			t.Run(fmt.Sprintf("sampler %s test %s", s.name, test.name), func(t *testing.T) {
    80  				test.test(t, s.sampler)
    81  			})
    82  		}
    83  	}
    84  }
    85  
    86  func WeightedWithoutReplacementInitializeOverflowTest(
    87  	t *testing.T,
    88  	s WeightedWithoutReplacement,
    89  ) {
    90  	err := s.Initialize([]uint64{1, math.MaxUint64})
    91  	require.ErrorIs(t, err, safemath.ErrOverflow)
    92  }
    93  
    94  func WeightedWithoutReplacementOutOfRangeTest(
    95  	t *testing.T,
    96  	s WeightedWithoutReplacement,
    97  ) {
    98  	require := require.New(t)
    99  
   100  	require.NoError(s.Initialize([]uint64{1}))
   101  
   102  	_, ok := s.Sample(2)
   103  	require.False(ok)
   104  }
   105  
   106  func WeightedWithoutReplacementEmptyWithoutWeightTest(
   107  	t *testing.T,
   108  	s WeightedWithoutReplacement,
   109  ) {
   110  	require := require.New(t)
   111  
   112  	require.NoError(s.Initialize(nil))
   113  
   114  	indices, ok := s.Sample(0)
   115  	require.True(ok)
   116  	require.Empty(indices)
   117  }
   118  
   119  func WeightedWithoutReplacementEmptyTest(
   120  	t *testing.T,
   121  	s WeightedWithoutReplacement,
   122  ) {
   123  	require := require.New(t)
   124  
   125  	require.NoError(s.Initialize([]uint64{1}))
   126  
   127  	indices, ok := s.Sample(0)
   128  	require.True(ok)
   129  	require.Empty(indices)
   130  }
   131  
   132  func WeightedWithoutReplacementSingletonTest(
   133  	t *testing.T,
   134  	s WeightedWithoutReplacement,
   135  ) {
   136  	require := require.New(t)
   137  
   138  	require.NoError(s.Initialize([]uint64{1}))
   139  
   140  	indices, ok := s.Sample(1)
   141  	require.True(ok)
   142  	require.Equal([]int{0}, indices)
   143  }
   144  
   145  func WeightedWithoutReplacementWithZeroTest(
   146  	t *testing.T,
   147  	s WeightedWithoutReplacement,
   148  ) {
   149  	require := require.New(t)
   150  
   151  	require.NoError(s.Initialize([]uint64{0, 1}))
   152  
   153  	indices, ok := s.Sample(1)
   154  	require.True(ok)
   155  	require.Equal([]int{1}, indices)
   156  }
   157  
   158  func WeightedWithoutReplacementDistributionTest(
   159  	t *testing.T,
   160  	s WeightedWithoutReplacement,
   161  ) {
   162  	require := require.New(t)
   163  
   164  	require.NoError(s.Initialize([]uint64{1, 1, 2}))
   165  
   166  	indices, ok := s.Sample(4)
   167  	require.True(ok)
   168  
   169  	slices.Sort(indices)
   170  	require.Equal([]int{0, 1, 2, 2}, indices)
   171  }