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 }