@@ -178,3 +178,128 @@ func TestControlRestart(t *testing.T) {
178
178
assert .NoError (t , r .Error )
179
179
}
180
180
}
181
+
182
+ func TestControlSuspend (t * testing.T ) {
183
+ c := NewTestCluster ()
184
+ ctx := context .Background ()
185
+ store := c .store .(* storemocks.Store )
186
+ lock := & lockmocks.DistributedLock {}
187
+ lock .On ("Lock" , mock .Anything ).Return (ctx , nil )
188
+ lock .On ("Unlock" , mock .Anything ).Return (nil )
189
+ store .On ("CreateLock" , mock .Anything , mock .Anything ).Return (lock , nil )
190
+ workload := & types.Workload {
191
+ ID : "id1" ,
192
+ Privileged : true ,
193
+ }
194
+ engine := & enginemocks.API {}
195
+ workload .Engine = engine
196
+ store .On ("GetWorkloads" , mock .Anything , mock .Anything ).Return ([]* types.Workload {workload }, nil )
197
+ // failed, hook true, remove always false
198
+ hook := & types.Hook {
199
+ BeforeSuspend : []string {"cmd1" },
200
+ }
201
+ workload .Hook = hook
202
+ workload .Hook .Force = true
203
+ engine .On ("Execute" , mock .Anything , mock .Anything , mock .Anything ).Return ("" , nil , nil , nil , types .ErrNilEngine )
204
+ ch , err := c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadSuspend , false )
205
+ assert .NoError (t , err )
206
+ for r := range ch {
207
+ assert .Error (t , r .Error )
208
+ }
209
+ // stop failed
210
+ workload .Hook .Force = false
211
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadSuspend , false )
212
+ engine .On ("VirtualizationSuspend" , mock .Anything , mock .Anything , mock .Anything ).Return (types .ErrNilEngine ).Once ()
213
+ assert .NoError (t , err )
214
+ for r := range ch {
215
+ assert .Error (t , r .Error )
216
+ }
217
+ engine .On ("VirtualizationSuspend" , mock .Anything , mock .Anything , mock .Anything ).Return (nil )
218
+ // stop success
219
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadSuspend , false )
220
+ assert .NoError (t , err )
221
+ for r := range ch {
222
+ assert .NoError (t , r .Error )
223
+ }
224
+ }
225
+
226
+ func TestControlResume (t * testing.T ) {
227
+ c := NewTestCluster ()
228
+ ctx := context .Background ()
229
+ store := c .store .(* storemocks.Store )
230
+ lock := & lockmocks.DistributedLock {}
231
+ lock .On ("Lock" , mock .Anything ).Return (ctx , nil )
232
+ lock .On ("Unlock" , mock .Anything ).Return (nil )
233
+ store .On ("CreateLock" , mock .Anything , mock .Anything ).Return (lock , nil )
234
+ // failed by GetWorkloads
235
+ store .On ("GetWorkloads" , mock .Anything , mock .Anything ).Return (nil , types .ErrMockError ).Once ()
236
+ ch , err := c .ControlWorkload (ctx , []string {"id1" }, "" , true )
237
+ assert .NoError (t , err )
238
+ for r := range ch {
239
+ assert .Error (t , r .Error )
240
+ }
241
+ workload := & types.Workload {
242
+ ID : "id1" ,
243
+ Privileged : true ,
244
+ }
245
+ engine := & enginemocks.API {}
246
+ workload .Engine = engine
247
+ store .On ("GetWorkloads" , mock .Anything , mock .Anything ).Return ([]* types.Workload {workload }, nil )
248
+ // failed by type
249
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, "" , true )
250
+ assert .NoError (t , err )
251
+ for r := range ch {
252
+ assert .Error (t , r .Error )
253
+ }
254
+ // failed by start
255
+ engine .On ("VirtualizationResume" , mock .Anything , mock .Anything ).Return (types .ErrNilEngine ).Once ()
256
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadResume , false )
257
+ assert .NoError (t , err )
258
+ for r := range ch {
259
+ assert .Error (t , r .Error )
260
+ }
261
+ engine .On ("VirtualizationResume" , mock .Anything , mock .Anything ).Return (nil )
262
+ // failed by Execute
263
+ hook := & types.Hook {
264
+ AfterResume : []string {"cmd1" , "cmd2" },
265
+ }
266
+ workload .Hook = hook
267
+ workload .Hook .Force = false
268
+ engine .On ("Execute" , mock .Anything , mock .Anything , mock .Anything ).Return ("" , nil , nil , nil , types .ErrNilEngine ).Times (3 )
269
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadResume , false )
270
+ assert .NoError (t , err )
271
+ for r := range ch {
272
+ assert .NoError (t , r .Error )
273
+ }
274
+ // force false, get no error
275
+ workload .Hook .Force = true
276
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadResume , false )
277
+ assert .NoError (t , err )
278
+ for r := range ch {
279
+ assert .Error (t , r .Error )
280
+ assert .Equal (t , r .WorkloadID , "id1" )
281
+ }
282
+ data := io .NopCloser (bytes .NewBufferString ("output" ))
283
+ engine .On ("Execute" , mock .Anything , mock .Anything , mock .Anything ).Return ("eid" , data , nil , nil , nil ).Times (4 )
284
+ // failed by ExecExitCode
285
+ engine .On ("ExecExitCode" , mock .Anything , mock .Anything , mock .Anything ).Return (- 1 , types .ErrNilEngine ).Once ()
286
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadResume , false )
287
+ assert .NoError (t , err )
288
+ for r := range ch {
289
+ assert .Error (t , r .Error )
290
+ }
291
+ // exitCode is not 0
292
+ engine .On ("ExecExitCode" , mock .Anything , mock .Anything , mock .Anything ).Return (- 1 , nil ).Once ()
293
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadResume , false )
294
+ assert .NoError (t , err )
295
+ for r := range ch {
296
+ assert .Error (t , r .Error )
297
+ }
298
+ // exitCode is 0
299
+ engine .On ("ExecExitCode" , mock .Anything , mock .Anything , mock .Anything ).Return (0 , nil )
300
+ ch , err = c .ControlWorkload (ctx , []string {"id1" }, cluster .WorkloadResume , false )
301
+ assert .NoError (t , err )
302
+ for r := range ch {
303
+ assert .NoError (t , r .Error )
304
+ }
305
+ }
0 commit comments