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

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package drainer
     5  
     6  import (
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/hernad/nomad/ci"
    11  	"github.com/hernad/nomad/nomad/mock"
    12  	"github.com/hernad/nomad/nomad/state"
    13  	"github.com/hernad/nomad/nomad/structs"
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  )
    17  
    18  // testDrainingNode creates a *drainingNode with a 1h deadline but no allocs
    19  func testDrainingNode(t *testing.T) *drainingNode {
    20  	t.Helper()
    21  	state := state.TestStateStore(t)
    22  	node := mock.Node()
    23  	node.DrainStrategy = &structs.DrainStrategy{
    24  		DrainSpec: structs.DrainSpec{
    25  			Deadline: time.Hour,
    26  		},
    27  		ForceDeadline: time.Now().Add(time.Hour),
    28  	}
    29  
    30  	require.Nil(t, state.UpsertNode(structs.MsgTypeTestSetup, 100, node))
    31  	return NewDrainingNode(node, state)
    32  }
    33  
    34  func assertDrainingNode(t *testing.T, dn *drainingNode, isDone bool, remaining, running int) {
    35  	t.Helper()
    36  
    37  	done, err := dn.IsDone()
    38  	require.Nil(t, err)
    39  	assert.Equal(t, isDone, done, "IsDone mismatch")
    40  
    41  	allocs, err := dn.RemainingAllocs()
    42  	require.Nil(t, err)
    43  	assert.Len(t, allocs, remaining, "RemainingAllocs mismatch")
    44  
    45  	jobs, err := dn.DrainingJobs()
    46  	require.Nil(t, err)
    47  	assert.Len(t, jobs, running, "DrainingJobs mismatch")
    48  }
    49  
    50  func TestDrainingNode_Table(t *testing.T) {
    51  	ci.Parallel(t)
    52  
    53  	cases := []struct {
    54  		name      string
    55  		isDone    bool
    56  		remaining int
    57  		running   int
    58  		setup     func(*testing.T, *drainingNode)
    59  	}{
    60  		{
    61  			name:      "Empty",
    62  			isDone:    true,
    63  			remaining: 0,
    64  			running:   0,
    65  			setup:     func(*testing.T, *drainingNode) {},
    66  		},
    67  		{
    68  			name:      "Batch",
    69  			isDone:    false,
    70  			remaining: 1,
    71  			running:   1,
    72  			setup: func(t *testing.T, dn *drainingNode) {
    73  				alloc := mock.BatchAlloc()
    74  				alloc.NodeID = dn.node.ID
    75  				require.Nil(t, dn.state.UpsertJob(structs.MsgTypeTestSetup, 101, nil, alloc.Job))
    76  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 102, []*structs.Allocation{alloc}))
    77  			},
    78  		},
    79  		{
    80  			name:      "Service",
    81  			isDone:    false,
    82  			remaining: 1,
    83  			running:   1,
    84  			setup: func(t *testing.T, dn *drainingNode) {
    85  				alloc := mock.Alloc()
    86  				alloc.NodeID = dn.node.ID
    87  				require.Nil(t, dn.state.UpsertJob(structs.MsgTypeTestSetup, 101, nil, alloc.Job))
    88  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 102, []*structs.Allocation{alloc}))
    89  			},
    90  		},
    91  		{
    92  			name:      "System",
    93  			isDone:    true,
    94  			remaining: 1,
    95  			running:   0,
    96  			setup: func(t *testing.T, dn *drainingNode) {
    97  				alloc := mock.SystemAlloc()
    98  				alloc.NodeID = dn.node.ID
    99  				require.Nil(t, dn.state.UpsertJob(structs.MsgTypeTestSetup, 101, nil, alloc.Job))
   100  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 102, []*structs.Allocation{alloc}))
   101  			},
   102  		},
   103  		{
   104  			name:      "AllTerminal",
   105  			isDone:    true,
   106  			remaining: 0,
   107  			running:   0,
   108  			setup: func(t *testing.T, dn *drainingNode) {
   109  				allocs := []*structs.Allocation{mock.Alloc(), mock.BatchAlloc(), mock.SystemAlloc()}
   110  				for _, a := range allocs {
   111  					a.NodeID = dn.node.ID
   112  					require.Nil(t, dn.state.UpsertJob(structs.MsgTypeTestSetup, 101, nil, a.Job))
   113  				}
   114  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 102, allocs))
   115  
   116  				// StateStore doesn't like inserting new allocs
   117  				// with a terminal status, so set the status in
   118  				// a second pass
   119  				for _, a := range allocs {
   120  					a.ClientStatus = structs.AllocClientStatusComplete
   121  				}
   122  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 103, allocs))
   123  			},
   124  		},
   125  		{
   126  			name:      "ServiceTerminal",
   127  			isDone:    false,
   128  			remaining: 2,
   129  			running:   1,
   130  			setup: func(t *testing.T, dn *drainingNode) {
   131  				allocs := []*structs.Allocation{mock.Alloc(), mock.BatchAlloc(), mock.SystemAlloc()}
   132  				for _, a := range allocs {
   133  					a.NodeID = dn.node.ID
   134  					require.Nil(t, dn.state.UpsertJob(structs.MsgTypeTestSetup, 101, nil, a.Job))
   135  				}
   136  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 102, allocs))
   137  
   138  				// Set only the service job as terminal
   139  				allocs[0].ClientStatus = structs.AllocClientStatusComplete
   140  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 103, allocs))
   141  			},
   142  		},
   143  		{
   144  			name:      "AllTerminalButBatch",
   145  			isDone:    false,
   146  			remaining: 1,
   147  			running:   1,
   148  			setup: func(t *testing.T, dn *drainingNode) {
   149  				allocs := []*structs.Allocation{mock.Alloc(), mock.BatchAlloc(), mock.SystemAlloc()}
   150  				for _, a := range allocs {
   151  					a.NodeID = dn.node.ID
   152  					require.Nil(t, dn.state.UpsertJob(structs.MsgTypeTestSetup, 101, nil, a.Job))
   153  				}
   154  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 102, allocs))
   155  
   156  				// Set only the service and batch jobs as terminal
   157  				allocs[0].ClientStatus = structs.AllocClientStatusComplete
   158  				allocs[2].ClientStatus = structs.AllocClientStatusComplete
   159  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 103, allocs))
   160  			},
   161  		},
   162  		{
   163  			name:      "AllTerminalButSystem",
   164  			isDone:    true,
   165  			remaining: 1,
   166  			running:   0,
   167  			setup: func(t *testing.T, dn *drainingNode) {
   168  				allocs := []*structs.Allocation{mock.Alloc(), mock.BatchAlloc(), mock.SystemAlloc()}
   169  				for _, a := range allocs {
   170  					a.NodeID = dn.node.ID
   171  					require.Nil(t, dn.state.UpsertJob(structs.MsgTypeTestSetup, 101, nil, a.Job))
   172  				}
   173  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 102, allocs))
   174  
   175  				// Set only the service and batch jobs as terminal
   176  				allocs[0].ClientStatus = structs.AllocClientStatusComplete
   177  				allocs[1].ClientStatus = structs.AllocClientStatusComplete
   178  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 103, allocs))
   179  			},
   180  		},
   181  		{
   182  			name:      "HalfTerminal",
   183  			isDone:    false,
   184  			remaining: 3,
   185  			running:   2,
   186  			setup: func(t *testing.T, dn *drainingNode) {
   187  				allocs := []*structs.Allocation{
   188  					mock.Alloc(),
   189  					mock.BatchAlloc(),
   190  					mock.SystemAlloc(),
   191  					mock.Alloc(),
   192  					mock.BatchAlloc(),
   193  					mock.SystemAlloc(),
   194  				}
   195  				for _, a := range allocs {
   196  					a.NodeID = dn.node.ID
   197  					require.Nil(t, dn.state.UpsertJob(structs.MsgTypeTestSetup, 101, nil, a.Job))
   198  				}
   199  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 102, allocs))
   200  
   201  				// Set only the service and batch jobs as terminal
   202  				allocs[0].ClientStatus = structs.AllocClientStatusComplete
   203  				allocs[1].ClientStatus = structs.AllocClientStatusComplete
   204  				allocs[2].ClientStatus = structs.AllocClientStatusComplete
   205  				require.Nil(t, dn.state.UpsertAllocs(structs.MsgTypeTestSetup, 103, allocs))
   206  			},
   207  		},
   208  	}
   209  
   210  	// Default test drainingNode has no allocs, so it should be done and
   211  	// have no remaining allocs.
   212  	for _, tc := range cases {
   213  		tc := tc
   214  		t.Run(tc.name, func(t *testing.T) {
   215  			ci.Parallel(t)
   216  			dn := testDrainingNode(t)
   217  			tc.setup(t, dn)
   218  			assertDrainingNode(t, dn, tc.isDone, tc.remaining, tc.running)
   219  		})
   220  	}
   221  }