github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/lib/kheap/score_heap_test.go (about) 1 package kheap 2 3 import ( 4 "container/heap" 5 "testing" 6 7 "github.com/hashicorp/nomad/ci" 8 "github.com/stretchr/testify/require" 9 ) 10 11 type heapItem struct { 12 Value string 13 ScoreVal float64 14 } 15 16 func (h *heapItem) Data() interface{} { 17 return h.Value 18 } 19 20 func (h *heapItem) Score() float64 { 21 return h.ScoreVal 22 } 23 24 func TestScoreHeap(t *testing.T) { 25 ci.Parallel(t) 26 27 type testCase struct { 28 desc string 29 items map[string]float64 30 expected []*heapItem 31 } 32 33 cases := []testCase{ 34 { 35 desc: "More than K elements", 36 items: map[string]float64{ 37 "banana": 3.0, 38 "apple": 2.25, 39 "pear": 2.32, 40 "watermelon": 5.45, 41 "orange": 0.20, 42 "strawberry": 9.03, 43 "blueberry": 0.44, 44 "lemon": 3.9, 45 "cherry": 0.03, 46 }, 47 expected: []*heapItem{ 48 {Value: "pear", ScoreVal: 2.32}, 49 {Value: "banana", ScoreVal: 3.0}, 50 {Value: "lemon", ScoreVal: 3.9}, 51 {Value: "watermelon", ScoreVal: 5.45}, 52 {Value: "strawberry", ScoreVal: 9.03}, 53 }, 54 }, 55 { 56 desc: "Less than K elements", 57 items: map[string]float64{ 58 "eggplant": 9.0, 59 "okra": -1.0, 60 "corn": 0.25, 61 }, 62 expected: []*heapItem{ 63 {Value: "okra", ScoreVal: -1.0}, 64 {Value: "corn", ScoreVal: 0.25}, 65 {Value: "eggplant", ScoreVal: 9.0}, 66 }, 67 }, 68 } 69 70 for _, tc := range cases { 71 t.Run("", func(t *testing.T) { 72 // Create Score heap, push elements into it 73 pq := NewScoreHeap(5) 74 for value, score := range tc.items { 75 heapItem := &heapItem{ 76 Value: value, 77 ScoreVal: score, 78 } 79 heap.Push(pq, heapItem) 80 } 81 82 // Take the items out; they arrive in increasing Score order 83 require := require.New(t) 84 require.Equal(len(tc.expected), pq.Len()) 85 86 i := 0 87 for pq.Len() > 0 { 88 item := heap.Pop(pq).(*heapItem) 89 require.Equal(tc.expected[i], item) 90 i++ 91 } 92 }) 93 } 94 95 }