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

Lines changed: 57 additions & 206 deletions
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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy