github.com/hernad/nomad@v1.6.112/nomad/drainer/drain_heap_test.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package drainer
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/hernad/nomad/ci"
    12  	"github.com/hernad/nomad/testutil"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func TestDeadlineHeap_Interface(t *testing.T) {
    17  	ci.Parallel(t)
    18  	require := require.New(t)
    19  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
    20  	require.Implements((*DrainDeadlineNotifier)(nil), h)
    21  }
    22  
    23  func TestDeadlineHeap_WatchAndGet(t *testing.T) {
    24  	ci.Parallel(t)
    25  	require := require.New(t)
    26  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
    27  
    28  	now := time.Now()
    29  	nodeID := "1"
    30  	wait := 10 * time.Millisecond
    31  	deadline := now.Add(wait)
    32  	h.Watch(nodeID, deadline)
    33  
    34  	var batch []string
    35  	select {
    36  	case batch = <-h.NextBatch():
    37  	case <-time.After(testutil.Timeout(3 * wait)):
    38  		t.Fatal("timeout")
    39  	}
    40  
    41  	require.Len(batch, 1)
    42  	require.Equal(nodeID, batch[0])
    43  }
    44  
    45  func TestDeadlineHeap_WatchThenUpdateAndGet(t *testing.T) {
    46  	ci.Parallel(t)
    47  	require := require.New(t)
    48  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
    49  
    50  	now := time.Now()
    51  	nodeID := "1"
    52  	wait := 10 * time.Millisecond
    53  	deadline := now.Add(wait)
    54  
    55  	// Initially watch way in the future
    56  	h.Watch(nodeID, now.Add(24*time.Hour))
    57  
    58  	// Rewatch
    59  	h.Watch(nodeID, deadline)
    60  
    61  	var batch []string
    62  	select {
    63  	case batch = <-h.NextBatch():
    64  	case <-time.After(testutil.Timeout(2 * wait)):
    65  		t.Fatal("timeout")
    66  	}
    67  
    68  	require.Len(batch, 1)
    69  	require.Equal(nodeID, batch[0])
    70  }
    71  
    72  func TestDeadlineHeap_MultiwatchAndDelete(t *testing.T) {
    73  	ci.Parallel(t)
    74  	require := require.New(t)
    75  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
    76  
    77  	now := time.Now()
    78  	wait := 50 * time.Millisecond
    79  	deadline := now.Add(wait)
    80  
    81  	nodeID1 := "1"
    82  	nodeID2 := "2"
    83  	h.Watch(nodeID1, deadline)
    84  	h.Watch(nodeID2, deadline)
    85  
    86  	time.Sleep(1 * time.Millisecond)
    87  	h.Remove(nodeID2)
    88  
    89  	var batch []string
    90  	select {
    91  	case batch = <-h.NextBatch():
    92  	case <-time.After(testutil.Timeout(2 * wait)):
    93  		t.Fatal("timeout")
    94  	}
    95  
    96  	require.Len(batch, 1)
    97  	require.Equal(nodeID1, batch[0])
    98  }
    99  
   100  func TestDeadlineHeap_WatchCoalesce(t *testing.T) {
   101  	ci.Parallel(t)
   102  	require := require.New(t)
   103  	h := NewDeadlineHeap(context.Background(), 100*time.Millisecond)
   104  
   105  	now := time.Now()
   106  
   107  	group1 := map[string]time.Time{
   108  		"1": now.Add(5 * time.Millisecond),
   109  		"2": now.Add(10 * time.Millisecond),
   110  		"3": now.Add(20 * time.Millisecond),
   111  		"4": now.Add(100 * time.Millisecond),
   112  	}
   113  
   114  	group2 := map[string]time.Time{
   115  		"10": now.Add(350 * time.Millisecond),
   116  		"11": now.Add(360 * time.Millisecond),
   117  	}
   118  
   119  	for _, g := range []map[string]time.Time{group1, group2} {
   120  		for n, d := range g {
   121  			h.Watch(n, d)
   122  		}
   123  	}
   124  
   125  	var batch []string
   126  	select {
   127  	case batch = <-h.NextBatch():
   128  	case <-time.After(testutil.Timeout(time.Second)):
   129  		t.Fatal("timeout")
   130  	}
   131  
   132  	require.Len(batch, len(group1))
   133  	for nodeID := range group1 {
   134  		require.Contains(batch, nodeID)
   135  	}
   136  	batch = nil
   137  
   138  	select {
   139  	case batch = <-h.NextBatch():
   140  	case <-time.After(testutil.Timeout(2 * time.Second)):
   141  		t.Fatal("timeout")
   142  	}
   143  
   144  	require.Len(batch, len(group2))
   145  	for nodeID := range group2 {
   146  		require.Contains(batch, nodeID)
   147  	}
   148  
   149  	select {
   150  	case <-h.NextBatch():
   151  		t.Fatal("unexpected batch")
   152  	case <-time.After(testutil.Timeout(100 * time.Millisecond)):
   153  	}
   154  }
   155  
   156  func TestDeadlineHeap_MultipleForce(t *testing.T) {
   157  	ci.Parallel(t)
   158  	require := require.New(t)
   159  	h := NewDeadlineHeap(context.Background(), 1*time.Second)
   160  
   161  	nodeID := "1"
   162  	deadline := time.Time{}
   163  	h.Watch(nodeID, deadline)
   164  
   165  	var batch []string
   166  	select {
   167  	case batch = <-h.NextBatch():
   168  	case <-time.After(testutil.Timeout(10 * time.Millisecond)):
   169  		t.Fatal("timeout")
   170  	}
   171  
   172  	require.Len(batch, 1)
   173  	require.Equal(nodeID, batch[0])
   174  
   175  	nodeID = "2"
   176  	h.Watch(nodeID, deadline)
   177  	select {
   178  	case batch = <-h.NextBatch():
   179  	case <-time.After(testutil.Timeout(10 * time.Millisecond)):
   180  		t.Fatal("timeout")
   181  	}
   182  
   183  	require.Len(batch, 1)
   184  	require.Equal(nodeID, batch[0])
   185  }