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

Skip to content

Commit c459533

Browse files
committed
Refactoring
Signed-off-by: Danny Kopping <[email protected]>
1 parent 217e46f commit c459533

File tree

2 files changed

+21
-80
lines changed

2 files changed

+21
-80
lines changed

coderd/workspaces.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -640,7 +640,7 @@ func createWorkspace(
640640
provisionerDaemons []database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow
641641
)
642642

643-
prebuilds := (*api.PrebuildsClaimer.Load()).(prebuilds.Claimer)
643+
prebuilds := *api.PrebuildsClaimer.Load()
644644

645645
err = api.Database.InTx(func(db database.Store) error {
646646
var (

enterprise/coderd/prebuilds/claim_test.go

Lines changed: 20 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,13 @@ import (
1010

1111
"github.com/coder/quartz"
1212
"github.com/google/uuid"
13-
"github.com/stretchr/testify/assert"
1413
"github.com/stretchr/testify/require"
1514

1615
"github.com/coder/coder/v2/coderd/coderdtest"
1716
"github.com/coder/coder/v2/coderd/database"
1817
"github.com/coder/coder/v2/coderd/database/dbauthz"
1918
"github.com/coder/coder/v2/coderd/database/dbtestutil"
19+
agplprebuilds "github.com/coder/coder/v2/coderd/prebuilds"
2020
"github.com/coder/coder/v2/coderd/rbac"
2121
"github.com/coder/coder/v2/codersdk"
2222
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
@@ -78,40 +78,20 @@ func TestClaimPrebuild(t *testing.T) {
7878
)
7979

8080
cases := map[string]struct {
81-
entitlementEnabled bool
82-
experimentEnabled bool
83-
attemptPrebuildClaim bool
8481
expectPrebuildClaimed bool
8582
markPrebuildsClaimable bool
86-
expectedPrebuildsCount int
8783
}{
88-
"without the experiment enabled, prebuilds will not provisioned": {
89-
experimentEnabled: false,
90-
entitlementEnabled: true,
91-
attemptPrebuildClaim: false,
92-
expectedPrebuildsCount: 0,
93-
},
94-
"without the entitlement, prebuilds will not provisioned": {
95-
experimentEnabled: true,
96-
entitlementEnabled: false,
97-
attemptPrebuildClaim: false,
98-
expectedPrebuildsCount: 0,
99-
},
100-
"with everything enabled, but no eligible prebuilds to claim": {
101-
entitlementEnabled: true,
102-
experimentEnabled: true,
103-
attemptPrebuildClaim: true,
84+
"no eligible prebuilds to claim": {
10485
expectPrebuildClaimed: false,
10586
markPrebuildsClaimable: false,
106-
expectedPrebuildsCount: desiredInstances * presetCount,
10787
},
108-
"with everything enabled, claiming an eligible prebuild should succeed": {
109-
entitlementEnabled: true,
110-
experimentEnabled: true,
111-
attemptPrebuildClaim: true,
88+
"claiming an eligible prebuild should succeed": {
89+
expectPrebuildClaimed: true,
90+
markPrebuildsClaimable: true,
91+
},
92+
"claiming an eligible prebuild results in error": {
11293
expectPrebuildClaimed: true,
11394
markPrebuildsClaimable: true,
114-
expectedPrebuildsCount: desiredInstances * presetCount,
11595
},
11696
}
11797

@@ -121,49 +101,26 @@ func TestClaimPrebuild(t *testing.T) {
121101
t.Run(name, func(t *testing.T) {
122102
t.Parallel()
123103

124-
// Setup. // TODO: abstract?
125-
104+
// Setup.
126105
ctx := testutil.Context(t, testutil.WaitMedium)
127106
db, pubsub := dbtestutil.NewDB(t)
128107
spy := newStoreSpy(db)
129-
130-
//var prebuildsEntitled int64
131-
//if tc.entitlementEnabled {
132-
// prebuildsEntitled = 1
133-
//}
108+
expectedPrebuildsCount := desiredInstances * presetCount
134109

135110
logger := testutil.Logger(t)
136-
client, _, _, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
111+
client, _, api, owner := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
137112
Options: &coderdtest.Options{
138113
IncludeProvisionerDaemon: true,
139114
Database: spy,
140115
Pubsub: pubsub,
141-
DeploymentValues: coderdtest.DeploymentValues(t, func(values *codersdk.DeploymentValues) {
142-
//values.Prebuilds.ReconciliationInterval = serpent.Duration(time.Hour) // We will kick off a reconciliation manually.
143-
//
144-
//if tc.experimentEnabled {
145-
// values.Experiments = serpent.StringArray{string(codersdk.ExperimentWorkspacePrebuilds)}
146-
//}
147-
}),
148116
},
149117

150118
EntitlementsUpdateInterval: time.Second,
151-
//LicenseOptions: &coderdenttest.LicenseOptions{
152-
// Features: license.Features{
153-
// codersdk.FeatureWorkspacePrebuilds: prebuildsEntitled,
154-
// },
155-
//},
156119
})
157-
reconciler := prebuilds.NewStoreReconciler(spy, pubsub, codersdk.PrebuildsConfig{}, logger, quartz.NewMock(t))
158-
159-
// The entitlements will need to refresh before the reconciler is set.
160-
require.Eventually(t, func() bool {
161-
if tc.entitlementEnabled && tc.experimentEnabled {
162-
assert.IsType(t, &prebuilds.StoreReconciler{}, reconciler)
163-
}
164120

165-
return reconciler != nil
166-
}, testutil.WaitSuperLong, testutil.IntervalFast)
121+
reconciler := prebuilds.NewStoreReconciler(spy, pubsub, codersdk.PrebuildsConfig{}, logger, quartz.NewMock(t))
122+
var claimer agplprebuilds.Claimer = &prebuilds.EnterpriseClaimer{}
123+
api.AGPL.PrebuildsClaimer.Store(&claimer)
167124

168125
version := coderdtest.CreateTemplateVersion(t, client, owner.OrganizationID, templateWithAgentAndPresetsWithPrebuilds(desiredInstances))
169126
_ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
@@ -179,14 +136,8 @@ func TestClaimPrebuild(t *testing.T) {
179136
// Given: the reconciliation state is snapshot.
180137
state, err := reconciler.SnapshotState(ctx, spy)
181138
require.NoError(t, err)
182-
183-
// When: the experiment or entitlement is not preset, there should be nothing to reconcile.
184-
if !tc.entitlementEnabled || !tc.experimentEnabled {
185-
require.Len(t, state.Presets, 0)
186-
return
187-
}
188-
189139
require.Len(t, state.Presets, presetCount)
140+
190141
// When: a reconciliation is setup for each preset.
191142
for _, preset := range presets {
192143
ps, err := state.FilterByPreset(preset.ID)
@@ -215,6 +166,7 @@ func TestClaimPrebuild(t *testing.T) {
215166
agents, err := db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, row.ID)
216167
require.NoError(t, err)
217168

169+
// Workspaces are eligible once its agent is marked "ready".
218170
for _, agent := range agents {
219171
require.NoError(t, db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
220172
ID: agent.ID,
@@ -225,9 +177,9 @@ func TestClaimPrebuild(t *testing.T) {
225177
}
226178
}
227179

228-
t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), tc.expectedPrebuildsCount)
180+
t.Logf("found %d running prebuilds so far, want %d", len(runningPrebuilds), expectedPrebuildsCount)
229181

230-
return len(runningPrebuilds) == tc.expectedPrebuildsCount
182+
return len(runningPrebuilds) == expectedPrebuildsCount
231183
}, testutil.WaitSuperLong, testutil.IntervalSlow)
232184

233185
// When: a user creates a new workspace with a preset for which prebuilds are configured.
@@ -243,21 +195,10 @@ func TestClaimPrebuild(t *testing.T) {
243195
TemplateVersionPresetID: presets[0].ID,
244196
ClaimPrebuildIfAvailable: true, // TODO: doesn't do anything yet; it probably should though.
245197
})
198+
246199
require.NoError(t, err)
247200
coderdtest.AwaitWorkspaceBuildJobCompleted(t, userClient, userWorkspace.LatestBuild.ID)
248201

249-
// Then: if we're not expecting any prebuild claims to succeed, handle this specifically.
250-
if !tc.attemptPrebuildClaim {
251-
require.EqualValues(t, spy.claims.Load(), 0)
252-
require.Nil(t, spy.claimedWorkspace.Load())
253-
254-
currentPrebuilds, err := spy.GetRunningPrebuiltWorkspaces(ctx)
255-
require.NoError(t, err)
256-
// The number of prebuilds should NOT change.
257-
require.Equal(t, len(currentPrebuilds), len(runningPrebuilds))
258-
return
259-
}
260-
261202
// Then: a prebuild should have been claimed.
262203
require.EqualValues(t, spy.claims.Load(), 1)
263204
require.NotNil(t, spy.claims.Load())
@@ -315,9 +256,9 @@ func TestClaimPrebuild(t *testing.T) {
315256
rows, err := spy.GetRunningPrebuiltWorkspaces(ctx)
316257
require.NoError(t, err)
317258

318-
t.Logf("found %d running prebuilds so far, want %d", len(rows), tc.expectedPrebuildsCount)
259+
t.Logf("found %d running prebuilds so far, want %d", len(rows), expectedPrebuildsCount)
319260

320-
return len(runningPrebuilds) == tc.expectedPrebuildsCount
261+
return len(runningPrebuilds) == expectedPrebuildsCount
321262
}, testutil.WaitSuperLong, testutil.IntervalSlow)
322263

323264
// Then: when restarting the created workspace (which claimed a prebuild), it should not try and claim a new prebuild.

0 commit comments

Comments
 (0)