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 }