@@ -136,6 +136,14 @@ func (s *scaletestStrategyFlags) toStrategy() harness.ExecutionStrategy {
136
136
return strategy
137
137
}
138
138
139
+ func (s * scaletestStrategyFlags ) toContext (ctx context.Context ) (context.Context , context.CancelFunc ) {
140
+ if s .timeout > 0 {
141
+ return context .WithTimeout (ctx , s .timeout )
142
+ }
143
+
144
+ return context .WithCancel (ctx )
145
+ }
146
+
139
147
type scaleTestOutputFormat string
140
148
141
149
const (
@@ -272,11 +280,39 @@ func requireAdmin(ctx context.Context, client *codersdk.Client) (codersdk.User,
272
280
return me , nil
273
281
}
274
282
283
+ // userCleanupRunner is a runner that deletes a user in the Run phase.
284
+ type userCleanupRunner struct {
285
+ client * codersdk.Client
286
+ userID uuid.UUID
287
+ }
288
+
289
+ var _ harness.Runnable = & userCleanupRunner {}
290
+
291
+ // Run implements Runnable.
292
+ func (r * userCleanupRunner ) Run (ctx context.Context , _ string , _ io.Writer ) error {
293
+ if r .userID == uuid .Nil {
294
+ return nil
295
+ }
296
+ ctx , span := tracing .StartSpan (ctx )
297
+ defer span .End ()
298
+
299
+ err := r .client .DeleteUser (ctx , r .userID )
300
+ if err != nil {
301
+ return xerrors .Errorf ("delete user %q: %w" , r .userID , err )
302
+ }
303
+
304
+ return nil
305
+ }
306
+
275
307
func scaletestCleanup () * cobra.Command {
308
+ var (
309
+ cleanupStrategy = & scaletestStrategyFlags {cleanup : true }
310
+ )
311
+
276
312
cmd := & cobra.Command {
277
313
Use : "cleanup" ,
278
314
Short : "Cleanup any orphaned scaletest resources" ,
279
- Long : "Cleanup scaletest workspaces." ,
315
+ Long : "Cleanup scaletest workspaces, then cleanup scaletest users. The strategy flags will apply to each stage of the cleanup process ." ,
280
316
RunE : func (cmd * cobra.Command , args []string ) error {
281
317
ctx := cmd .Context ()
282
318
client , err := CreateClient (cmd )
@@ -291,16 +327,126 @@ func scaletestCleanup() *cobra.Command {
291
327
292
328
client .BypassRatelimits = true
293
329
294
- // TODO: cleanup workspaces
295
- _ = isScaleTestWorkspace
330
+ cmd .PrintErrln ("Fetching scaletest workspaces..." )
331
+ var (
332
+ pageNumber = 0
333
+ limit = 100
334
+ workspaces []codersdk.Workspace
335
+ )
336
+ for {
337
+ page , err := client .Workspaces (ctx , codersdk.WorkspaceFilter {
338
+ Name : "scaletest-" ,
339
+ Offset : pageNumber * limit ,
340
+ Limit : limit ,
341
+ })
342
+ if err != nil {
343
+ return xerrors .Errorf ("fetch scaletest workspaces page %d: %w" , pageNumber , err )
344
+ }
345
+
346
+ if len (page .Workspaces ) == 0 {
347
+ break
348
+ }
349
+
350
+ pageWorkspaces := make ([]codersdk.Workspace , 0 , len (page .Workspaces ))
351
+ for _ , w := range page .Workspaces {
352
+ if isScaleTestWorkspace (w ) {
353
+ pageWorkspaces = append (pageWorkspaces , w )
354
+ }
355
+ }
356
+ workspaces = append (workspaces , pageWorkspaces ... )
357
+ }
358
+
359
+ cmd .PrintErrf ("Found %d scaletest workspaces\n " , len (workspaces ))
360
+ if len (workspaces ) != 0 {
361
+ cmd .Println ("Deleting scaletest workspaces..." )
362
+ harness := harness .NewTestHarness (cleanupStrategy .toStrategy (), harness.ConcurrentExecutionStrategy {})
363
+
364
+ for i , w := range workspaces {
365
+ const testName = "cleanup-workspace"
366
+ r := workspacebuild .NewCleanupRunner (client , w .ID )
367
+ harness .AddRun (testName , strconv .Itoa (i ), r )
368
+ }
369
+
370
+ ctx , cancel := cleanupStrategy .toContext (ctx )
371
+ defer cancel ()
372
+ err := harness .Run (ctx )
373
+ if err != nil {
374
+ return xerrors .Errorf ("run test harness to delete workspaces (harness failure, not a test failure): %w" , err )
375
+ }
296
376
297
- // TODO: cleanup users
298
- _ = isScaleTestUser
377
+ cmd .Println ("Done deleting scaletest workspaces:" )
378
+ res := harness .Results ()
379
+ res .PrintText (cmd .ErrOrStderr ())
380
+
381
+ if res .TotalFail > 0 {
382
+ return xerrors .Errorf ("failed to delete scaletest workspaces" )
383
+ }
384
+ }
385
+
386
+ cmd .PrintErrln ("Fetching scaletest users..." )
387
+ pageNumber = 0
388
+ limit = 100
389
+ var users []codersdk.User
390
+ for {
391
+ page , err := client .Users (ctx , codersdk.UsersRequest {
392
+ Search : "scaletest-" ,
393
+ Pagination : codersdk.Pagination {
394
+ Offset : pageNumber * limit ,
395
+ Limit : limit ,
396
+ },
397
+ })
398
+ if err != nil {
399
+ return xerrors .Errorf ("fetch scaletest users page %d: %w" , pageNumber , err )
400
+ }
401
+
402
+ if len (page .Users ) == 0 {
403
+ break
404
+ }
405
+
406
+ pageUsers := make ([]codersdk.User , 0 , len (page .Users ))
407
+ for _ , u := range page .Users {
408
+ if isScaleTestUser (u ) {
409
+ pageUsers = append (pageUsers , u )
410
+ }
411
+ }
412
+ users = append (users , pageUsers ... )
413
+ }
414
+
415
+ cmd .PrintErrf ("Found %d scaletest users\n " , len (users ))
416
+ if len (workspaces ) != 0 {
417
+ cmd .Println ("Deleting scaletest users..." )
418
+ harness := harness .NewTestHarness (cleanupStrategy .toStrategy (), harness.ConcurrentExecutionStrategy {})
419
+
420
+ for i , u := range users {
421
+ const testName = "cleanup-users"
422
+ r := & userCleanupRunner {
423
+ client : client ,
424
+ userID : u .ID ,
425
+ }
426
+ harness .AddRun (testName , strconv .Itoa (i ), r )
427
+ }
428
+
429
+ ctx , cancel := cleanupStrategy .toContext (ctx )
430
+ defer cancel ()
431
+ err := harness .Run (ctx )
432
+ if err != nil {
433
+ return xerrors .Errorf ("run test harness to delete users (harness failure, not a test failure): %w" , err )
434
+ }
435
+
436
+ cmd .Println ("Done deleting scaletest users:" )
437
+ res := harness .Results ()
438
+ res .PrintText (cmd .ErrOrStderr ())
439
+
440
+ if res .TotalFail > 0 {
441
+ return xerrors .Errorf ("failed to delete scaletest users" )
442
+ }
443
+ }
299
444
300
445
return nil
301
446
},
302
447
}
303
448
449
+ cleanupStrategy .attach (cmd )
304
450
return cmd
305
451
}
306
452
@@ -494,7 +640,7 @@ func scaletestCreateWorkspaces() *cobra.Command {
494
640
495
641
config := createworkspaces.Config {
496
642
User : createworkspaces.UserConfig {
497
- // TODO: configuration org
643
+ // TODO: configurable org
498
644
OrganizationID : me .OrganizationIDs [0 ],
499
645
Username : username ,
500
646
Email : email ,
@@ -563,12 +709,8 @@ func scaletestCreateWorkspaces() *cobra.Command {
563
709
564
710
// TODO: live progress output
565
711
_ , _ = fmt .Fprintln (cmd .ErrOrStderr (), "Running load test..." )
566
- testCtx := ctx
567
- if strategy .timeout > 0 {
568
- var testCancel func ()
569
- testCtx , testCancel = context .WithTimeout (testCtx , strategy .timeout )
570
- defer testCancel ()
571
- }
712
+ testCtx , testCancel := strategy .toContext (ctx )
713
+ defer testCancel ()
572
714
err = th .Run (testCtx )
573
715
if err != nil {
574
716
return xerrors .Errorf ("run test harness (harness failure, not a test failure): %w" , err )
@@ -583,12 +725,8 @@ func scaletestCreateWorkspaces() *cobra.Command {
583
725
}
584
726
585
727
_ , _ = fmt .Fprintln (cmd .ErrOrStderr (), "\n Cleaning up..." )
586
- cleanupCtx := ctx
587
- if cleanupStrategy .timeout > 0 {
588
- var cleanupCancel context.CancelFunc
589
- cleanupCtx , cleanupCancel = context .WithTimeout (ctx , cleanupStrategy .timeout )
590
- defer cleanupCancel ()
591
- }
728
+ cleanupCtx , cleanupCancel := cleanupStrategy .toContext (ctx )
729
+ defer cleanupCancel ()
592
730
err = th .Cleanup (cleanupCtx )
593
731
if err != nil {
594
732
return xerrors .Errorf ("cleanup tests: %w" , err )
0 commit comments