github.com/MetalBlockchain/metalgo@v1.11.9/utils/sampler/weighted_heap_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  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  )
    11  
    12  func TestWeightedHeapInitialize(t *testing.T) {
    13  	require := require.New(t)
    14  
    15  	h := weightedHeap{}
    16  
    17  	require.NoError(h.Initialize([]uint64{2, 2, 1, 3, 3, 1, 3}))
    18  
    19  	expectedOrdering := []int{3, 4, 6, 0, 1, 2, 5}
    20  	for i, elem := range h.heap {
    21  		expected := expectedOrdering[i]
    22  		require.Equal(expected, elem.index)
    23  	}
    24  }
    25  
    26  func TestWeightedHeapElementCompare(t *testing.T) {
    27  	type test struct {
    28  		name     string
    29  		elt1     weightedHeapElement
    30  		elt2     weightedHeapElement
    31  		expected int
    32  	}
    33  	tests := []test{
    34  		{
    35  			name:     "all same",
    36  			elt1:     weightedHeapElement{},
    37  			elt2:     weightedHeapElement{},
    38  			expected: 0,
    39  		},
    40  		{
    41  			name: "lower weight",
    42  			elt1: weightedHeapElement{},
    43  			elt2: weightedHeapElement{
    44  				weight: 1,
    45  			},
    46  			expected: 1,
    47  		},
    48  		{
    49  			name: "higher index",
    50  			elt1: weightedHeapElement{
    51  				index: 1,
    52  			},
    53  			elt2:     weightedHeapElement{},
    54  			expected: 1,
    55  		},
    56  	}
    57  
    58  	for _, tt := range tests {
    59  		t.Run(tt.name, func(t *testing.T) {
    60  			require := require.New(t)
    61  
    62  			require.Equal(tt.expected, tt.elt1.Compare(tt.elt2))
    63  			require.Equal(-tt.expected, tt.elt2.Compare(tt.elt1))
    64  		})
    65  	}
    66  }