github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/lib/delayheap/delay_heap_test.go (about)

     1  package delayheap
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/hashicorp/nomad/ci"
     8  	"github.com/stretchr/testify/require"
     9  )
    10  
    11  // HeapNodeImpl satisfies the HeapNode interface
    12  type heapNodeImpl struct {
    13  	dataObject interface{}
    14  	id         string
    15  	namespace  string
    16  }
    17  
    18  func (d *heapNodeImpl) Data() interface{} {
    19  	return d.dataObject
    20  }
    21  
    22  func (d *heapNodeImpl) ID() string {
    23  	return d.id
    24  }
    25  
    26  func (d *heapNodeImpl) Namespace() string {
    27  	return d.namespace
    28  }
    29  
    30  func TestDelayHeap_PushPop(t *testing.T) {
    31  	ci.Parallel(t)
    32  
    33  	delayHeap := NewDelayHeap()
    34  	now := time.Now()
    35  	require := require.New(t)
    36  	// a dummy type to use as the inner object in the heap
    37  	type myObj struct {
    38  		a int
    39  		b string
    40  	}
    41  	dataNode1 := &heapNodeImpl{
    42  		dataObject: &myObj{a: 0, b: "hey"},
    43  		id:         "101",
    44  		namespace:  "default",
    45  	}
    46  	delayHeap.Push(dataNode1, now.Add(-10*time.Minute))
    47  
    48  	dataNode2 := &heapNodeImpl{
    49  		dataObject: &myObj{a: 0, b: "hey"},
    50  		id:         "102",
    51  		namespace:  "default",
    52  	}
    53  	delayHeap.Push(dataNode2, now.Add(10*time.Minute))
    54  
    55  	dataNode3 := &heapNodeImpl{
    56  		dataObject: &myObj{a: 0, b: "hey"},
    57  		id:         "103",
    58  		namespace:  "default",
    59  	}
    60  	delayHeap.Push(dataNode3, now.Add(-15*time.Second))
    61  
    62  	dataNode4 := &heapNodeImpl{
    63  		dataObject: &myObj{a: 0, b: "hey"},
    64  		id:         "101",
    65  		namespace:  "test-namespace",
    66  	}
    67  	delayHeap.Push(dataNode4, now.Add(2*time.Hour))
    68  
    69  	expectedWaitTimes := []time.Time{now.Add(-10 * time.Minute), now.Add(-15 * time.Second), now.Add(10 * time.Minute), now.Add(2 * time.Hour)}
    70  	entries := getHeapEntries(delayHeap, now)
    71  	for i, entry := range entries {
    72  		require.Equal(expectedWaitTimes[i], entry.WaitUntil)
    73  	}
    74  
    75  }
    76  
    77  func TestDelayHeap_Update(t *testing.T) {
    78  	ci.Parallel(t)
    79  
    80  	delayHeap := NewDelayHeap()
    81  	now := time.Now()
    82  	require := require.New(t)
    83  	// a dummy type to use as the inner object in the heap
    84  	type myObj struct {
    85  		a int
    86  		b string
    87  	}
    88  	dataNode1 := &heapNodeImpl{
    89  		dataObject: &myObj{a: 0, b: "hey"},
    90  		id:         "101",
    91  		namespace:  "default",
    92  	}
    93  	delayHeap.Push(dataNode1, now.Add(-10*time.Minute))
    94  
    95  	dataNode2 := &heapNodeImpl{
    96  		dataObject: &myObj{a: 0, b: "hey"},
    97  		id:         "102",
    98  		namespace:  "default",
    99  	}
   100  	delayHeap.Push(dataNode2, now.Add(10*time.Minute))
   101  	delayHeap.Update(dataNode1, now.Add(20*time.Minute))
   102  
   103  	expectedWaitTimes := []time.Time{now.Add(10 * time.Minute), now.Add(20 * time.Minute)}
   104  	expectedIdOrder := []string{"102", "101"}
   105  	entries := getHeapEntries(delayHeap, now)
   106  	for i, entry := range entries {
   107  		require.Equal(expectedWaitTimes[i], entry.WaitUntil)
   108  		require.Equal(expectedIdOrder[i], entry.Node.ID())
   109  	}
   110  
   111  }
   112  
   113  func getHeapEntries(delayHeap *DelayHeap, now time.Time) []*delayHeapNode {
   114  	var entries []*delayHeapNode
   115  	for node := delayHeap.Pop(); node != nil; {
   116  		entries = append(entries, node)
   117  		node = delayHeap.Pop()
   118  	}
   119  	return entries
   120  }