Skip to content

Commit 9167cbf

Browse files
evgeniy-scherbinaevgeniy-scherbina
authored
refactor: claim prebuilt workspace tests (#17567)
Follow-up to: #17458 Specifically it addresses these discussions: - #17458 (comment)
1 parent 3ab3ef8 commit 9167cbf

File tree

1 file changed

+57
-206
lines changed

1 file changed

+57
-206
lines changed

‎enterprise/coderd/prebuilds/claim_test.go

+57-206
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package prebuilds_test
33
import (
44
"context"
55
"database/sql"
6+
"errors"
67
"slices"
78
"strings"
89
"sync/atomic"
@@ -35,21 +36,25 @@ type storeSpy struct {
3536
claims *atomic.Int32
3637
claimParams *atomic.Pointer[database.ClaimPrebuiltWorkspaceParams]
3738
claimedWorkspace *atomic.Pointer[database.ClaimPrebuiltWorkspaceRow]
39+
40+
// if claimingErr is not nil - error will be returned when ClaimPrebuiltWorkspace is called
41+
claimingErr error
3842
}
3943

40-
func newStoreSpy(db database.Store) *storeSpy {
44+
func newStoreSpy(db database.Store, claimingErr error) *storeSpy {
4145
return &storeSpy{
4246
Store: db,
4347
claims: &atomic.Int32{},
4448
claimParams: &atomic.Pointer[database.ClaimPrebuiltWorkspaceParams]{},
4549
claimedWorkspace: &atomic.Pointer[database.ClaimPrebuiltWorkspaceRow]{},
50+
claimingErr: claimingErr,
4651
}
4752
}
4853

4954
func (m *storeSpy) InTx(fn func(store database.Store) error, opts *database.TxOptions) error {
5055
// Pass spy down into transaction store.
5156
return m.Store.InTx(func(store database.Store) error {
52-
spy := newStoreSpy(store)
57+
spy := newStoreSpy(store, m.claimingErr)
5358
spy.claims = m.claims
5459
spy.claimParams = m.claimParams
5560
spy.claimedWorkspace = m.claimedWorkspace
@@ -59,6 +64,10 @@ func (m *storeSpy) InTx(fn func(store database.Store) error, opts *database.TxOp
5964
}
6065

6166
func (m *storeSpy) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
67+
if m.claimingErr != nil {
68+
return database.ClaimPrebuiltWorkspaceRow{}, m.claimingErr
69+
}
70+
6271
m.claims.Add(1)
6372
m.claimParams.Store(&arg)
6473
result, err := m.Store.ClaimPrebuiltWorkspace(ctx, arg)
@@ -68,32 +77,6 @@ func (m *storeSpy) ClaimPrebuiltWorkspace(ctx context.Context, arg database.Clai
6877
return result, err
6978
}
7079

71-
type errorStore struct {
72-
claimingErr error
73-
74-
database.Store
75-
}
76-
77-
func newErrorStore(db database.Store, claimingErr error) *errorStore {
78-
return &errorStore{
79-
Store: db,
80-
claimingErr: claimingErr,
81-
}
82-
}
83-
84-
func (es *errorStore) InTx(fn func(store database.Store) error, opts *database.TxOptions) error {
85-
// Pass failure store down into transaction store.
86-
return es.Store.InTx(func(store database.Store) error {
87-
newES := newErrorStore(store, es.claimingErr)
88-
89-
return fn(newES)
90-
}, opts)
91-
}
92-
93-
func (es *errorStore) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
94-
return database.ClaimPrebuiltWorkspaceRow{}, es.claimingErr
95-
}
96-
9780
func TestClaimPrebuild(t *testing.T) {
9881
t.Parallel()
9982

@@ -106,9 +89,13 @@ func TestClaimPrebuild(t *testing.T) {
10689
presetCount = 2
10790
)
10891

92+
unexpectedClaimingError := xerrors.New("unexpected claiming error")
93+
10994
cases := map[string]struct {
11095
expectPrebuildClaimed bool
11196
markPrebuildsClaimable bool
97+
// if claimingErr is not nil - error will be returned when ClaimPrebuiltWorkspace is called
98+
claimingErr error
11299
}{
113100
"no eligible prebuilds to claim": {
114101
expectPrebuildClaimed: false,
@@ -118,6 +105,17 @@ func TestClaimPrebuild(t *testing.T) {
118105
expectPrebuildClaimed: true,
119106
markPrebuildsClaimable: true,
120107
},
108+
109+
"no claimable prebuilt workspaces error is returned": {
110+
expectPrebuildClaimed: false,
111+
markPrebuildsClaimable: true,
112+
claimingErr: agplprebuilds.ErrNoClaimablePrebuiltWorkspaces,
113+
},
114+
"unexpected claiming error is returned": {
115+
expectPrebuildClaimed: false,
116+
markPrebuildsClaimable: true,
117+
claimingErr: unexpectedClaimingError,
118+
},
121119
}
122120

123121
for name, tc := range cases {
@@ -129,7 +127,8 @@ func TestClaimPrebuild(t *testing.T) {
129127
// Setup.
130128
ctx := testutil.Context(t, testutil.WaitSuperLong)
131129
db, pubsub := dbtestutil.NewDB(t)
132-
spy := newStoreSpy(db)
130+
131+
spy := newStoreSpy(db, tc.claimingErr)
133132
expectedPrebuildsCount := desiredInstances * presetCount
134133

135134
logger := testutil.Logger(t)
@@ -225,8 +224,35 @@ func TestClaimPrebuild(t *testing.T) {
225224
TemplateVersionPresetID: presets[0].ID,
226225
})
227226

228-
require.NoError(t, err)
229-
coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID)
227+
switch {
228+
case tc.claimingErr != nil && errors.Is(tc.claimingErr, agplprebuilds.ErrNoClaimablePrebuiltWorkspaces):
229+
require.NoError(t, err)
230+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID)
231+
232+
// Then: the number of running prebuilds hasn't changed because claiming prebuild is failed and we fallback to creating new workspace.
233+
currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx)
234+
require.NoError(t, err)
235+
require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds))
236+
return
237+
238+
case tc.claimingErr != nil && errors.Is(tc.claimingErr, unexpectedClaimingError):
239+
// Then: unexpected error happened and was propagated all the way to the caller
240+
require.Error(t, err)
241+
require.ErrorContains(t, err, unexpectedClaimingError.Error())
242+
243+
// Then: the number of running prebuilds hasn't changed because claiming prebuild is failed.
244+
currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx)
245+
require.NoError(t, err)
246+
require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds))
247+
return
248+
249+
default:
250+
// tc.claimingErr is nil scenario
251+
require.NoError(t, err)
252+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID)
253+
}
254+
255+
// at this point we know that tc.claimingErr is nil
230256

231257
// Then: a prebuild should have been claimed.
232258
require.EqualValues(t, spy.claims.Load(), 1)
@@ -315,181 +341,6 @@ func TestClaimPrebuild(t *testing.T) {
315341
}
316342
}
317343

318-
func TestClaimPrebuild_CheckDifferentErrors(t *testing.T) {
319-
t.Parallel()
320-
321-
if !dbtestutil.WillUsePostgres() {
322-
t.Skip("This test requires postgres")
323-
}
324-
325-
const (
326-
desiredInstances = 1
327-
presetCount = 2
328-
329-
expectedPrebuildsCount = desiredInstances * presetCount
330-
)
331-
332-
cases := map[string]struct {
333-
claimingErr error
334-
checkFn func(
335-
t *testing.T,
336-
ctx context.Context,
337-
store database.Store,
338-
userClient *codersdk.Client,
339-
user codersdk.User,
340-
templateVersionID uuid.UUID,
341-
presetID uuid.UUID,
342-
)
343-
}{
344-
"ErrNoClaimablePrebuiltWorkspaces is returned": {
345-
claimingErr: agplprebuilds.ErrNoClaimablePrebuiltWorkspaces,
346-
checkFn: func(
347-
t *testing.T,
348-
ctx context.Context,
349-
store database.Store,
350-
userClient *codersdk.Client,
351-
user codersdk.User,
352-
templateVersionID uuid.UUID,
353-
presetID uuid.UUID,
354-
) {
355-
// When: a user creates a new workspace with a preset for which prebuilds are configured.
356-
workspaceName := strings.ReplaceAll(testutil.GetRandomName(t), "_", "-")
357-
userWorkspace, err := userClient.CreateUserWorkspace(ctx, user.Username, codersdk.CreateWorkspaceRequest{
358-
TemplateVersionID: templateVersionID,
359-
Name: workspaceName,
360-
TemplateVersionPresetID: presetID,
361-
})
362-
363-
require.NoError(t, err)
364-
coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID)
365-
366-
// Then: the number of running prebuilds hasn't changed because claiming prebuild is failed and we fallback to creating new workspace.
367-
currentPrebuilds, err := store.GetRunningPrebuiltWorkspaces(ctx)
368-
require.NoError(t, err)
369-
require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds))
370-
},
371-
},
372-
"unexpected error during claim is returned": {
373-
claimingErr: xerrors.New("unexpected error during claim"),
374-
checkFn: func(
375-
t *testing.T,
376-
ctx context.Context,
377-
store database.Store,
378-
userClient *codersdk.Client,
379-
user codersdk.User,
380-
templateVersionID uuid.UUID,
381-
presetID uuid.UUID,
382-
) {
383-
// When: a user creates a new workspace with a preset for which prebuilds are configured.
384-
workspaceName := strings.ReplaceAll(testutil.GetRandomName(t), "_", "-")
385-
_, err := userClient.CreateUserWorkspace(ctx, user.Username, codersdk.CreateWorkspaceRequest{
386-
TemplateVersionID: templateVersionID,
387-
Name: workspaceName,
388-
TemplateVersionPresetID: presetID,
389-
})
390-
391-
// Then: unexpected error happened and was propagated all the way to the caller
392-
require.Error(t, err)
393-
require.ErrorContains(t, err, "unexpected error during claim")
394-
395-
// Then: the number of running prebuilds hasn't changed because claiming prebuild is failed.
396-
currentPrebuilds, err := store.GetRunningPrebuiltWorkspaces(ctx)
397-
require.NoError(t, err)
398-
require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds))
399-
},
400-
},
401-
}
402-
403-
for name, tc := range cases {
404-
t.Run(name, func(t *testing.T) {
405-
t.Parallel()
406-
407-
// Setup.
408-
ctx := testutil.Context(t, testutil.WaitSuperLong)
409-
db, pubsub := dbtestutil.NewDB(t)
410-
errorStore := newErrorStore(db, tc.claimingErr)
411-
412-
logger := testutil.Logger(t)
413-
client, _, api, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
414-
Options: &coderdtest.Options{
415-
IncludeProvisionerDaemon: true,
416-
Database: errorStore,
417-
Pubsub: pubsub,
418-
},
419-
420-
EntitlementsUpdateInterval: time.Second,
421-
})
422-
423-
reconciler := prebuilds.NewStoreReconciler(errorStore, pubsub, codersdk.PrebuildsConfig{}, logger, quartz.NewMock(t), api.PrometheusRegistry)
424-
var claimer agplprebuilds.Claimer = prebuilds.NewEnterpriseClaimer(errorStore)
425-
api.AGPL.PrebuildsClaimer.Store(&claimer)
426-
427-
version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, templateWithAgentAndPresetsWithPrebuilds(desiredInstances))
428-
_ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
429-
coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID)
430-
presets, err := client.TemplateVersionPresets(ctx, version.ID)
431-
require.NoError(t, err)
432-
require.Len(t, presets, presetCount)
433-
434-
userClient, user := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID, rbac.RoleMember())
435-
436-
// Given: the reconciliation state is snapshot.
437-
state, err := reconciler.SnapshotState(ctx, errorStore)
438-
require.NoError(t, err)
439-
require.Len(t, state.Presets, presetCount)
440-
441-
// When: a reconciliation is setup for each preset.
442-
for _, preset := range presets {
443-
ps, err := state.FilterByPreset(preset.ID)
444-
require.NoError(t, err)
445-
require.NotNil(t, ps)
446-
actions, err := reconciler.CalculateActions(ctx, *ps)
447-
require.NoError(t, err)
448-
require.NotNil(t, actions)
449-
450-
require.NoError(t, reconciler.ReconcilePreset(ctx, *ps))
451-
}
452-
453-
// Given: a set of running, eligible prebuilds eventually starts up.
454-
runningPrebuilds := make(map[uuid.UUID]database.GetRunningPrebuiltWorkspacesRow, desiredInstances*presetCount)
455-
require.Eventually(t, func() bool {
456-
rows, err := errorStore.GetRunningPrebuiltWorkspaces(ctx)
457-
if err != nil {
458-
return false
459-
}
460-
461-
for _, row := range rows {
462-
runningPrebuilds[row.CurrentPresetID.UUID] = row
463-
464-
agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, row.ID)
465-
if err != nil {
466-
return false
467-
}
468-
469-
// Workspaces are eligible once its agent is marked "ready".
470-
for _, agent := range agents {
471-
err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
472-
ID: agent.ID,
473-
LifecycleState: database.WorkspaceAgentLifecycleStateReady,
474-
StartedAt: sql.NullTime{Time: time.Now().Add(time.Hour), Valid: true},
475-
ReadyAt: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true},
476-
})
477-
if err != nil {
478-
return false
479-
}
480-
}
481-
}
482-
483-
t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), expectedPrebuildsCount)
484-
485-
return len(runningPrebuilds) == expectedPrebuildsCount
486-
}, testutil.WaitSuperLong, testutil.IntervalSlow)
487-
488-
tc.checkFn(t, ctx, errorStore, userClient, user, version.ID, presets[0].ID)
489-
})
490-
}
491-
}
492-
493344
func templateWithAgentAndPresetsWithPrebuilds(desiredInstances int32) *echo.Responses {
494345
return &echo.Responses{
495346
Parse: echo.ParseComplete,

0 commit comments

Comments
 (0)