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