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