@@ -3,6 +3,7 @@ package prebuilds_test
3
3
import (
4
4
"context"
5
5
"database/sql"
6
+ "errors"
6
7
"strings"
7
8
"sync/atomic"
8
9
"testing"
@@ -66,6 +67,32 @@ func (m *storeSpy) ClaimPrebuiltWorkspace(ctx context.Context, arg database.Clai
66
67
return result ,err
67
68
}
68
69
70
+ type errorStore struct {
71
+ claimingErr error
72
+
73
+ database.Store
74
+ }
75
+
76
+ func newErrorStore (db database.Store ,claimingErr error )* errorStore {
77
+ return & errorStore {
78
+ Store :db ,
79
+ claimingErr :claimingErr ,
80
+ }
81
+ }
82
+
83
+ func (es * errorStore )InTx (fn func (store database.Store )error ,opts * database.TxOptions )error {
84
+ // Pass failure store down into transaction store.
85
+ return es .Store .InTx (func (store database.Store )error {
86
+ newES := newErrorStore (store ,es .claimingErr )
87
+
88
+ return fn (newES )
89
+ },opts )
90
+ }
91
+
92
+ func (es * errorStore )ClaimPrebuiltWorkspace (ctx context.Context ,arg database.ClaimPrebuiltWorkspaceParams ) (database.ClaimPrebuiltWorkspaceRow ,error ) {
93
+ return database.ClaimPrebuiltWorkspaceRow {},es .claimingErr
94
+ }
95
+
69
96
func TestClaimPrebuild (t * testing.T ) {
70
97
t .Parallel ()
71
98
@@ -284,6 +311,171 @@ func TestClaimPrebuild(t *testing.T) {
284
311
}
285
312
}
286
313
314
+ func TestClaimPrebuild_CheckDifferentErrors (t * testing.T ) {
315
+ t .Parallel ()
316
+
317
+ if ! dbtestutil .WillUsePostgres () {
318
+ t .Skip ("This test requires postgres" )
319
+ }
320
+
321
+ const (
322
+ desiredInstances = 1
323
+ presetCount = 2
324
+
325
+ expectedPrebuildsCount = desiredInstances * presetCount
326
+ )
327
+
328
+ cases := map [string ]struct {
329
+ claimingErr error
330
+ checkFn func (
331
+ t * testing.T ,
332
+ ctx context.Context ,
333
+ store database.Store ,
334
+ userClient * codersdk.Client ,
335
+ user codersdk.User ,
336
+ templateVersionID uuid.UUID ,
337
+ presetID uuid.UUID ,
338
+ )
339
+ }{
340
+ "ErrNoClaimablePrebuiltWorkspaces is returned" : {
341
+ claimingErr :agplprebuilds .ErrNoClaimablePrebuiltWorkspaces ,
342
+ checkFn :func (
343
+ t * testing.T ,
344
+ ctx context.Context ,
345
+ store database.Store ,
346
+ userClient * codersdk.Client ,
347
+ user codersdk.User ,
348
+ templateVersionID uuid.UUID ,
349
+ presetID uuid.UUID ,
350
+ ) {
351
+ // When: a user creates a new workspace with a preset for which prebuilds are configured.
352
+ workspaceName := strings .ReplaceAll (testutil .GetRandomName (t ),"_" ,"-" )
353
+ userWorkspace ,err := userClient .CreateUserWorkspace (ctx ,user .Username , codersdk.CreateWorkspaceRequest {
354
+ TemplateVersionID :templateVersionID ,
355
+ Name :workspaceName ,
356
+ TemplateVersionPresetID :presetID ,
357
+ })
358
+
359
+ require .NoError (t ,err )
360
+ coderdtest .AwaitWorkspaceBuildJobCompleted (t ,userClient ,userWorkspace .LatestBuild .ID )
361
+
362
+ // Then: the number of running prebuilds hasn't changed because claiming prebuild is failed and we fallback to creating new workspace.
363
+ currentPrebuilds ,err := store .GetRunningPrebuiltWorkspaces (ctx )
364
+ require .NoError (t ,err )
365
+ require .Equal (t ,expectedPrebuildsCount ,len (currentPrebuilds ))
366
+ },
367
+ },
368
+ "unexpected error during claim is returned" : {
369
+ claimingErr :errors .New ("unexpected error during claim" ),
370
+ checkFn :func (
371
+ t * testing.T ,
372
+ ctx context.Context ,
373
+ store database.Store ,
374
+ userClient * codersdk.Client ,
375
+ user codersdk.User ,
376
+ templateVersionID uuid.UUID ,
377
+ presetID uuid.UUID ,
378
+ ) {
379
+ // When: a user creates a new workspace with a preset for which prebuilds are configured.
380
+ workspaceName := strings .ReplaceAll (testutil .GetRandomName (t ),"_" ,"-" )
381
+ _ ,err := userClient .CreateUserWorkspace (ctx ,user .Username , codersdk.CreateWorkspaceRequest {
382
+ TemplateVersionID :templateVersionID ,
383
+ Name :workspaceName ,
384
+ TemplateVersionPresetID :presetID ,
385
+ })
386
+
387
+ // Then: unexpected error happened and was propagated all the way to the caller
388
+ require .Error (t ,err )
389
+ require .ErrorContains (t ,err ,"unexpected error during claim" )
390
+ },
391
+ },
392
+ }
393
+
394
+ for name ,tc := range cases {
395
+ t .Run (name ,func (t * testing.T ) {
396
+ t .Parallel ()
397
+
398
+ // Setup.
399
+ ctx := testutil .Context (t ,testutil .WaitMedium )
400
+ db ,pubsub := dbtestutil .NewDB (t )
401
+ failureStore := newErrorStore (db ,tc .claimingErr )
402
+
403
+ logger := testutil .Logger (t )
404
+ client ,_ ,api ,owner := coderdenttest .NewWithAPI (t ,& coderdenttest.Options {
405
+ Options :& coderdtest.Options {
406
+ IncludeProvisionerDaemon :true ,
407
+ Database :failureStore ,
408
+ Pubsub :pubsub ,
409
+ },
410
+
411
+ EntitlementsUpdateInterval :time .Second ,
412
+ })
413
+
414
+ reconciler := prebuilds .NewStoreReconciler (failureStore ,pubsub , codersdk.PrebuildsConfig {},logger ,quartz .NewMock (t ))
415
+ var claimer agplprebuilds.Claimer = & prebuilds.EnterpriseClaimer {}
416
+ api .AGPL .PrebuildsClaimer .Store (& claimer )
417
+
418
+ version := coderdtest .CreateTemplateVersion (t ,client ,owner .OrganizationID ,templateWithAgentAndPresetsWithPrebuilds (desiredInstances ))
419
+ _ = coderdtest .AwaitTemplateVersionJobCompleted (t ,client ,version .ID )
420
+ coderdtest .CreateTemplate (t ,client ,owner .OrganizationID ,version .ID )
421
+ presets ,err := client .TemplateVersionPresets (ctx ,version .ID )
422
+ require .NoError (t ,err )
423
+ require .Len (t ,presets ,presetCount )
424
+
425
+ userClient ,user := coderdtest .CreateAnotherUser (t ,client ,owner .OrganizationID ,rbac .RoleMember ())
426
+
427
+ ctx = dbauthz .AsPrebuildsOrchestrator (ctx )
428
+
429
+ // Given: the reconciliation state is snapshot.
430
+ state ,err := reconciler .SnapshotState (ctx ,failureStore )
431
+ require .NoError (t ,err )
432
+ require .Len (t ,state .Presets ,presetCount )
433
+
434
+ // When: a reconciliation is setup for each preset.
435
+ for _ ,preset := range presets {
436
+ ps ,err := state .FilterByPreset (preset .ID )
437
+ require .NoError (t ,err )
438
+ require .NotNil (t ,ps )
439
+ actions ,err := reconciler .CalculateActions (ctx ,* ps )
440
+ require .NoError (t ,err )
441
+ require .NotNil (t ,actions )
442
+
443
+ require .NoError (t ,reconciler .ReconcilePreset (ctx ,* ps ))
444
+ }
445
+
446
+ // Given: a set of running, eligible prebuilds eventually starts up.
447
+ runningPrebuilds := make (map [uuid.UUID ]database.GetRunningPrebuiltWorkspacesRow ,desiredInstances * presetCount )
448
+ require .Eventually (t ,func ()bool {
449
+ rows ,err := failureStore .GetRunningPrebuiltWorkspaces (ctx )
450
+ require .NoError (t ,err )
451
+
452
+ for _ ,row := range rows {
453
+ runningPrebuilds [row .CurrentPresetID .UUID ]= row
454
+
455
+ agents ,err := db .GetWorkspaceAgentsInLatestBuildByWorkspaceID (ctx ,row .ID )
456
+ require .NoError (t ,err )
457
+
458
+ // Workspaces are eligible once its agent is marked "ready".
459
+ for _ ,agent := range agents {
460
+ require .NoError (t ,db .UpdateWorkspaceAgentLifecycleStateByID (ctx , database.UpdateWorkspaceAgentLifecycleStateByIDParams {
461
+ ID :agent .ID ,
462
+ LifecycleState :database .WorkspaceAgentLifecycleStateReady ,
463
+ StartedAt : sql.NullTime {Time :time .Now ().Add (time .Hour ),Valid :true },
464
+ ReadyAt : sql.NullTime {Time :time .Now ().Add (- 1 * time .Hour ),Valid :true },
465
+ }))
466
+ }
467
+ }
468
+
469
+ t .Logf ("found %d running prebuilds so far, want %d" ,len (runningPrebuilds ),expectedPrebuildsCount )
470
+
471
+ return len (runningPrebuilds )== expectedPrebuildsCount
472
+ },testutil .WaitSuperLong ,testutil .IntervalSlow )
473
+
474
+ tc .checkFn (t ,ctx ,failureStore ,userClient ,user ,version .ID ,presets [0 ].ID )
475
+ })
476
+ }
477
+ }
478
+
287
479
func templateWithAgentAndPresetsWithPrebuilds (desiredInstances int32 )* echo.Responses {
288
480
return & echo.Responses {
289
481
Parse :echo .ParseComplete ,