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  }