github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/nomad/drainer/watch_nodes_test.go (about)

     1  package drainer
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/hashicorp/nomad/ci"
     9  	"github.com/hashicorp/nomad/helper/testlog"
    10  	"github.com/hashicorp/nomad/nomad/mock"
    11  	"github.com/hashicorp/nomad/nomad/state"
    12  	"github.com/hashicorp/nomad/nomad/structs"
    13  	"github.com/hashicorp/nomad/testutil"
    14  	"github.com/stretchr/testify/require"
    15  	"golang.org/x/time/rate"
    16  )
    17  
    18  func testNodeDrainWatcher(t *testing.T) (*nodeDrainWatcher, *state.StateStore, *MockNodeTracker) {
    19  	t.Helper()
    20  	state := state.TestStateStore(t)
    21  	limiter := rate.NewLimiter(100.0, 100)
    22  	logger := testlog.HCLogger(t)
    23  	m := NewMockNodeTracker()
    24  	w := NewNodeDrainWatcher(context.Background(), limiter, state, logger, m)
    25  	return w, state, m
    26  }
    27  
    28  func TestNodeDrainWatcher_Interface(t *testing.T) {
    29  	ci.Parallel(t)
    30  	require := require.New(t)
    31  	w, _, _ := testNodeDrainWatcher(t)
    32  	require.Implements((*DrainingNodeWatcher)(nil), w)
    33  }
    34  
    35  func TestNodeDrainWatcher_AddDraining(t *testing.T) {
    36  	ci.Parallel(t)
    37  	require := require.New(t)
    38  	_, state, m := testNodeDrainWatcher(t)
    39  
    40  	// Create two nodes, one draining and one not draining
    41  	n1, n2 := mock.Node(), mock.Node()
    42  	n2.DrainStrategy = &structs.DrainStrategy{
    43  		DrainSpec: structs.DrainSpec{
    44  			Deadline: time.Hour,
    45  		},
    46  		ForceDeadline: time.Now().Add(time.Hour),
    47  	}
    48  
    49  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 100, n1))
    50  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 101, n2))
    51  
    52  	testutil.WaitForResult(func() (bool, error) {
    53  		return len(m.events()) == 1, nil
    54  	}, func(err error) {
    55  		t.Fatal("No node drain events")
    56  	})
    57  
    58  	tracked := m.TrackedNodes()
    59  	require.NotContains(tracked, n1.ID)
    60  	require.Contains(tracked, n2.ID)
    61  	require.Equal(n2, tracked[n2.ID])
    62  
    63  }
    64  
    65  func TestNodeDrainWatcher_Remove(t *testing.T) {
    66  	ci.Parallel(t)
    67  	require := require.New(t)
    68  	_, state, m := testNodeDrainWatcher(t)
    69  
    70  	// Create a draining node
    71  	n := mock.Node()
    72  	n.DrainStrategy = &structs.DrainStrategy{
    73  		DrainSpec: structs.DrainSpec{
    74  			Deadline: time.Hour,
    75  		},
    76  		ForceDeadline: time.Now().Add(time.Hour),
    77  	}
    78  
    79  	// Wait for it to be tracked
    80  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 100, n))
    81  	testutil.WaitForResult(func() (bool, error) {
    82  		return len(m.events()) == 1, nil
    83  	}, func(err error) {
    84  		t.Fatal("No node drain events")
    85  	})
    86  
    87  	tracked := m.TrackedNodes()
    88  	require.Contains(tracked, n.ID)
    89  	require.Equal(n, tracked[n.ID])
    90  
    91  	// Change the node to be not draining and wait for it to be untracked
    92  	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 101, n.ID, nil, false, 0, nil, nil, ""))
    93  	testutil.WaitForResult(func() (bool, error) {
    94  		return len(m.events()) == 2, nil
    95  	}, func(err error) {
    96  		t.Fatal("No new node drain events")
    97  	})
    98  
    99  	tracked = m.TrackedNodes()
   100  	require.NotContains(tracked, n.ID)
   101  }
   102  
   103  func TestNodeDrainWatcher_Remove_Nonexistent(t *testing.T) {
   104  	ci.Parallel(t)
   105  	require := require.New(t)
   106  	_, state, m := testNodeDrainWatcher(t)
   107  
   108  	// Create a draining node
   109  	n := mock.Node()
   110  	n.DrainStrategy = &structs.DrainStrategy{
   111  		DrainSpec: structs.DrainSpec{
   112  			Deadline: time.Hour,
   113  		},
   114  		ForceDeadline: time.Now().Add(time.Hour),
   115  	}
   116  
   117  	// Wait for it to be tracked
   118  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 100, n))
   119  	testutil.WaitForResult(func() (bool, error) {
   120  		return len(m.events()) == 1, nil
   121  	}, func(err error) {
   122  		t.Fatal("No node drain events")
   123  	})
   124  
   125  	tracked := m.TrackedNodes()
   126  	require.Contains(tracked, n.ID)
   127  	require.Equal(n, tracked[n.ID])
   128  
   129  	// Delete the node
   130  	require.Nil(state.DeleteNode(structs.MsgTypeTestSetup, 101, []string{n.ID}))
   131  	testutil.WaitForResult(func() (bool, error) {
   132  		return len(m.events()) == 2, nil
   133  	}, func(err error) {
   134  		t.Fatal("No new node drain events")
   135  	})
   136  
   137  	tracked = m.TrackedNodes()
   138  	require.NotContains(tracked, n.ID)
   139  }
   140  
   141  func TestNodeDrainWatcher_Update(t *testing.T) {
   142  	ci.Parallel(t)
   143  	require := require.New(t)
   144  	_, state, m := testNodeDrainWatcher(t)
   145  
   146  	// Create a draining node
   147  	n := mock.Node()
   148  	n.DrainStrategy = &structs.DrainStrategy{
   149  		DrainSpec: structs.DrainSpec{
   150  			Deadline: time.Hour,
   151  		},
   152  		ForceDeadline: time.Now().Add(time.Hour),
   153  	}
   154  
   155  	// Wait for it to be tracked
   156  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 100, n))
   157  	testutil.WaitForResult(func() (bool, error) {
   158  		return len(m.events()) == 1, nil
   159  	}, func(err error) {
   160  		t.Fatal("No node drain events")
   161  	})
   162  
   163  	tracked := m.TrackedNodes()
   164  	require.Contains(tracked, n.ID)
   165  	require.Equal(n, tracked[n.ID])
   166  
   167  	// Change the node to have a new spec
   168  	s2 := n.DrainStrategy.Copy()
   169  	s2.Deadline += time.Hour
   170  	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 101, n.ID, s2, false, 0, nil, nil, ""))
   171  
   172  	// Wait for it to be updated
   173  	testutil.WaitForResult(func() (bool, error) {
   174  		return len(m.events()) == 2, nil
   175  	}, func(err error) {
   176  		t.Fatal("No new node drain events")
   177  	})
   178  
   179  	tracked = m.TrackedNodes()
   180  	require.Contains(tracked, n.ID)
   181  	require.Equal(s2, tracked[n.ID].DrainStrategy)
   182  }