github.com/manicqin/nomad@v0.9.5/lib/delayheap/delay_heap_test.go (about)

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