Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit c3b57cf

Browse files
evgeniy-scherbinaevgeniy-scherbina
committed
refactor: simplify prebuild claiming test
1 parent 64084d2 commit c3b57cf

File tree

1 file changed

+13
-52
lines changed

1 file changed

+13
-52
lines changed

‎enterprise/coderd/prebuilds/claim_test.go

Lines changed: 13 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -30,34 +30,30 @@ import (
3030
"github.com/coder/coder/v2/testutil"
3131
)
3232

33-
type storeType int
34-
35-
const (
36-
spyStoreType storeType = iota
37-
errorStoreType
38-
)
39-
4033
type storeSpy struct {
4134
database.Store
4235

4336
claims *atomic.Int32
4437
claimParams *atomic.Pointer[database.ClaimPrebuiltWorkspaceParams]
4538
claimedWorkspace *atomic.Pointer[database.ClaimPrebuiltWorkspaceRow]
39+
40+
claimingErr error
4641
}
4742

48-
func newStoreSpy(db database.Store) *storeSpy {
43+
func newStoreSpy(db database.Store, claimingErr error) *storeSpy {
4944
return &storeSpy{
5045
Store: db,
5146
claims: &atomic.Int32{},
5247
claimParams: &atomic.Pointer[database.ClaimPrebuiltWorkspaceParams]{},
5348
claimedWorkspace: &atomic.Pointer[database.ClaimPrebuiltWorkspaceRow]{},
49+
claimingErr: claimingErr,
5450
}
5551
}
5652

5753
func (m *storeSpy) InTx(fn func(store database.Store) error, opts *database.TxOptions) error {
5854
// Pass spy down into transaction store.
5955
return m.Store.InTx(func(store database.Store) error {
60-
spy := newStoreSpy(store)
56+
spy := newStoreSpy(store, m.claimingErr)
6157
spy.claims = m.claims
6258
spy.claimParams = m.claimParams
6359
spy.claimedWorkspace = m.claimedWorkspace
@@ -67,6 +63,10 @@ func (m *storeSpy) InTx(fn func(store database.Store) error, opts *database.TxOp
6763
}
6864

6965
func (m *storeSpy) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
66+
if m.claimingErr != nil {
67+
return database.ClaimPrebuiltWorkspaceRow{}, m.claimingErr
68+
}
69+
7070
m.claims.Add(1)
7171
m.claimParams.Store(&arg)
7272
result, err := m.Store.ClaimPrebuiltWorkspace(ctx, arg)
@@ -76,32 +76,6 @@ func (m *storeSpy) ClaimPrebuiltWorkspace(ctx context.Context, arg database.Clai
7676
return result, err
7777
}
7878

79-
type errorStore struct {
80-
claimingErr error
81-
82-
database.Store
83-
}
84-
85-
func newErrorStore(db database.Store, claimingErr error) *errorStore {
86-
return &errorStore{
87-
Store: db,
88-
claimingErr: claimingErr,
89-
}
90-
}
91-
92-
func (es *errorStore) InTx(fn func(store database.Store) error, opts *database.TxOptions) error {
93-
// Pass failure store down into transaction store.
94-
return es.Store.InTx(func(store database.Store) error {
95-
newES := newErrorStore(store, es.claimingErr)
96-
97-
return fn(newES)
98-
}, opts)
99-
}
100-
101-
func (es *errorStore) ClaimPrebuiltWorkspace(ctx context.Context, arg database.ClaimPrebuiltWorkspaceParams) (database.ClaimPrebuiltWorkspaceRow, error) {
102-
return database.ClaimPrebuiltWorkspaceRow{}, es.claimingErr
103-
}
104-
10579
func TestClaimPrebuild(t *testing.T) {
10680
t.Parallel()
10781

@@ -119,30 +93,25 @@ func TestClaimPrebuild(t *testing.T) {
11993
cases := map[string]struct {
12094
expectPrebuildClaimed bool
12195
markPrebuildsClaimable bool
122-
storeType storeType
12396
storeError error // should be set only for errorStoreType
12497
}{
12598
"no eligible prebuilds to claim": {
12699
expectPrebuildClaimed: false,
127100
markPrebuildsClaimable: false,
128-
storeType: spyStoreType,
129101
},
130102
"claiming an eligible prebuild should succeed": {
131103
expectPrebuildClaimed: true,
132104
markPrebuildsClaimable: true,
133-
storeType: spyStoreType,
134105
},
135106

136107
"no claimable prebuilt workspaces error is returned": {
137108
expectPrebuildClaimed: false,
138109
markPrebuildsClaimable: true,
139-
storeType: errorStoreType,
140110
storeError: agplprebuilds.ErrNoClaimablePrebuiltWorkspaces,
141111
},
142112
"unexpected claiming error is returned": {
143113
expectPrebuildClaimed: false,
144114
markPrebuildsClaimable: true,
145-
storeType: errorStoreType,
146115
storeError: unexpectedClaimingError,
147116
},
148117
}
@@ -157,15 +126,7 @@ func TestClaimPrebuild(t *testing.T) {
157126
ctx := testutil.Context(t, testutil.WaitSuperLong)
158127
db, pubsub := dbtestutil.NewDB(t)
159128

160-
var wrappedStore database.Store
161-
switch tc.storeType {
162-
case spyStoreType:
163-
wrappedStore = newStoreSpy(db)
164-
case errorStoreType:
165-
wrappedStore = newErrorStore(db, tc.storeError)
166-
default:
167-
t.Fatal("unknown store type")
168-
}
129+
wrappedStore := newStoreSpy(db, tc.storeError)
169130
expectedPrebuildsCount := desiredInstances * presetCount
170131

171132
logger := testutil.Logger(t)
@@ -262,7 +223,7 @@ func TestClaimPrebuild(t *testing.T) {
262223
})
263224

264225
switch {
265-
case tc.storeType == errorStoreType && errors.Is(tc.storeError, agplprebuilds.ErrNoClaimablePrebuiltWorkspaces):
226+
case tc.storeError != nil && errors.Is(tc.storeError, agplprebuilds.ErrNoClaimablePrebuiltWorkspaces):
266227
require.NoError(t, err)
267228
coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID)
268229

@@ -272,7 +233,7 @@ func TestClaimPrebuild(t *testing.T) {
272233
require.Equal(t, expectedPrebuildsCount, len(currentPrebuilds))
273234
return
274235

275-
case tc.storeType == errorStoreType && errors.Is(tc.storeError, unexpectedClaimingError):
236+
case tc.storeError != nil && errors.Is(tc.storeError, unexpectedClaimingError):
276237
// Then: unexpected error happened and was propagated all the way to the caller
277238
require.Error(t, err)
278239
require.ErrorContains(t, err, unexpectedClaimingError.Error())
@@ -290,7 +251,7 @@ func TestClaimPrebuild(t *testing.T) {
290251
}
291252

292253
// at this point we know that wrappedStore has *storeSpy type
293-
spy := wrappedStore.(*storeSpy)
254+
spy := wrappedStore
294255

295256
// Then: a prebuild should have been claimed.
296257
require.EqualValues(t, spy.claims.Load(), 1)

0 commit comments

Comments
 (0)