@@ -3,6 +3,7 @@ package prebuilds_test
33import (
44"context"
55"database/sql"
6+ "errors"
67"strings"
78"sync/atomic"
89"testing"
@@ -66,6 +67,32 @@ func (m *storeSpy) ClaimPrebuiltWorkspace(ctx context.Context, arg database.Clai
6667return result ,err
6768}
6869
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+
6996func TestClaimPrebuild (t * testing.T ) {
7097t .Parallel ()
7198
@@ -284,6 +311,171 @@ func TestClaimPrebuild(t *testing.T) {
284311}
285312}
286313
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+
287479func templateWithAgentAndPresetsWithPrebuilds (desiredInstances int32 )* echo.Responses {
288480return & echo.Responses {
289481Parse :echo .ParseComplete ,