Skip to content

Commit bcfbb04

Browse files
test: add integration test
1 parent 63ffcc7 commit bcfbb04

File tree

5 files changed

+158
-13
lines changed

5 files changed

+158
-13
lines changed

coderd/prebuilds/api.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ type ReconciliationOrchestrator interface {
3131

3232
// TrackResourceReplacement handles a pathological situation whereby a terraform resource is replaced due to drift,
3333
// which can obviate the whole point of pre-provisioning a prebuilt workspace.
34-
// See more detail at https://coder.com/docs/admin/templates/extending-templates/prebuilt-workspaces#preventing-resource-replacement.
34+
// See more detail now https://coder.com/docs/admin/templates/extending-templates/prebuilt-workspaces#preventing-resource-replacement.
3535
TrackResourceReplacement(ctx context.Context, workspaceID, buildID uuid.UUID, replacements []*sdkproto.ResourceReplacement)
3636
}
3737

coderd/prebuilds/global_snapshot.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package prebuilds
22

33
import (
4+
"github.com/coder/quartz"
45
"time"
56

67
"github.com/google/uuid"
@@ -18,6 +19,7 @@ type GlobalSnapshot struct {
1819
PrebuildsInProgress []database.CountInProgressPrebuildsRow
1920
Backoffs []database.GetPresetsBackoffRow
2021
HardLimitedPresetsMap map[uuid.UUID]database.GetPresetsAtFailureLimitRow
22+
clock quartz.Clock
2123
}
2224

2325
func NewGlobalSnapshot(
@@ -27,6 +29,7 @@ func NewGlobalSnapshot(
2729
prebuildsInProgress []database.CountInProgressPrebuildsRow,
2830
backoffs []database.GetPresetsBackoffRow,
2931
hardLimitedPresets []database.GetPresetsAtFailureLimitRow,
32+
clock quartz.Clock,
3033
) GlobalSnapshot {
3134
hardLimitedPresetsMap := make(map[uuid.UUID]database.GetPresetsAtFailureLimitRow, len(hardLimitedPresets))
3235
for _, preset := range hardLimitedPresets {
@@ -40,6 +43,7 @@ func NewGlobalSnapshot(
4043
PrebuildsInProgress: prebuildsInProgress,
4144
Backoffs: backoffs,
4245
HardLimitedPresetsMap: hardLimitedPresetsMap,
46+
clock: clock,
4347
}
4448
}
4549

@@ -88,6 +92,7 @@ func (s GlobalSnapshot) FilterByPreset(presetID uuid.UUID) (*PresetSnapshot, err
8892
InProgress: inProgress,
8993
Backoff: backoffPtr,
9094
IsHardLimited: isHardLimited,
95+
clock: s.clock,
9196
}, nil
9297
}
9398

coderd/prebuilds/preset_snapshot.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ type PresetSnapshot struct {
4646
InProgress []database.CountInProgressPrebuildsRow
4747
Backoff *database.GetPresetsBackoffRow
4848
IsHardLimited bool
49+
clock quartz.Clock
4950
}
5051

5152
// ReconciliationState represents the processed state of a preset's prebuilds,
@@ -155,9 +156,10 @@ func (p PresetSnapshot) CalculateState() *ReconciliationState {
155156

156157
if p.isActive() {
157158
var err error
158-
desired, err = p.CalculateDesiredInstances(time.Now())
159+
desired, err = p.CalculateDesiredInstances(p.clock.Now())
159160
if err != nil {
160161
// TODO: handle error
162+
panic(err)
161163
}
162164
eligible = p.countEligible()
163165
extraneous = max(actual-expired-desired, 0)
@@ -318,7 +320,7 @@ func (p PresetSnapshot) countEligible() int32 {
318320
}
319321

320322
// countInProgress returns counts of prebuilds in transition states (starting, stopping, deleting).
321-
// These counts are tracked at the template level, so all presets sharing the same template see the same values.
323+
// These counts are tracked now the template level, so all presets sharing the same template see the same values.
322324
func (p PresetSnapshot) countInProgress() (starting int32, stopping int32, deleting int32) {
323325
for _, progress := range p.InProgress {
324326
num := progress.Count

coderd/prebuilds/preset_snapshot_test.go

Lines changed: 147 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ func TestNoPrebuilds(t *testing.T) {
8484
preset(true, 0, current),
8585
}
8686

87-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, nil, nil, nil, nil)
87+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, nil, nil, nil, nil, quartz.NewMock(t))
8888
ps, err := snapshot.FilterByPreset(current.presetID)
8989
require.NoError(t, err)
9090

@@ -106,7 +106,7 @@ func TestNetNew(t *testing.T) {
106106
preset(true, 1, current),
107107
}
108108

109-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, nil, nil, nil, nil)
109+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, nil, nil, nil, nil, quartz.NewMock(t))
110110
ps, err := snapshot.FilterByPreset(current.presetID)
111111
require.NoError(t, err)
112112

@@ -148,7 +148,7 @@ func TestOutdatedPrebuilds(t *testing.T) {
148148
var inProgress []database.CountInProgressPrebuildsRow
149149

150150
// WHEN: calculating the outdated preset's state.
151-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil)
151+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil, quartz.NewMock(t))
152152
ps, err := snapshot.FilterByPreset(outdated.presetID)
153153
require.NoError(t, err)
154154

@@ -214,7 +214,7 @@ func TestDeleteOutdatedPrebuilds(t *testing.T) {
214214
}
215215

216216
// WHEN: calculating the outdated preset's state.
217-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil)
217+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil, quartz.NewMock(t))
218218
ps, err := snapshot.FilterByPreset(outdated.presetID)
219219
require.NoError(t, err)
220220

@@ -459,7 +459,7 @@ func TestInProgressActions(t *testing.T) {
459459
}
460460

461461
// WHEN: calculating the current preset's state.
462-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil)
462+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil, quartz.NewMock(t))
463463
ps, err := snapshot.FilterByPreset(current.presetID)
464464
require.NoError(t, err)
465465

@@ -502,7 +502,7 @@ func TestExtraneous(t *testing.T) {
502502
var inProgress []database.CountInProgressPrebuildsRow
503503

504504
// WHEN: calculating the current preset's state.
505-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil)
505+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil, quartz.NewMock(t))
506506
ps, err := snapshot.FilterByPreset(current.presetID)
507507
require.NoError(t, err)
508508

@@ -683,7 +683,7 @@ func TestExpiredPrebuilds(t *testing.T) {
683683
}
684684

685685
// WHEN: calculating the current preset's state.
686-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, nil, nil, nil)
686+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, nil, nil, nil, quartz.NewMock(t))
687687
ps, err := snapshot.FilterByPreset(current.presetID)
688688
require.NoError(t, err)
689689

@@ -719,7 +719,7 @@ func TestDeprecated(t *testing.T) {
719719
var inProgress []database.CountInProgressPrebuildsRow
720720

721721
// WHEN: calculating the current preset's state.
722-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil)
722+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, nil, nil, quartz.NewMock(t))
723723
ps, err := snapshot.FilterByPreset(current.presetID)
724724
require.NoError(t, err)
725725

@@ -772,7 +772,7 @@ func TestLatestBuildFailed(t *testing.T) {
772772
}
773773

774774
// WHEN: calculating the current preset's state.
775-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, backoffs, nil)
775+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, running, inProgress, backoffs, nil, quartz.NewMock(t))
776776
psCurrent, err := snapshot.FilterByPreset(current.presetID)
777777
require.NoError(t, err)
778778

@@ -865,7 +865,7 @@ func TestMultiplePresetsPerTemplateVersion(t *testing.T) {
865865
},
866866
}
867867

868-
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, nil, inProgress, nil, nil)
868+
snapshot := prebuilds.NewGlobalSnapshot(presets, nil, nil, inProgress, nil, nil, quartz.NewMock(t))
869869

870870
// Nothing has to be created for preset 1.
871871
{
@@ -905,6 +905,134 @@ func TestMultiplePresetsPerTemplateVersion(t *testing.T) {
905905
}
906906
}
907907

908+
func TestPrebuildAutoscaling(t *testing.T) {
909+
t.Parallel()
910+
911+
// The test includes 2 presets, each with 2 schedules.
912+
// It checks that the calculated actions match expectations for various provided times,
913+
// based on the corresponding schedules.
914+
testCases := []struct {
915+
name string
916+
// now specifies the current time.
917+
now time.Time
918+
// expected instances for preset1 and preset2, respectively.
919+
expectedInstances []int32
920+
}{
921+
{
922+
name: "Before the 1st schedule",
923+
now: mustParseTime(t, time.RFC1123, "Mon, 02 Jun 2025 01:00:00 UTC"),
924+
expectedInstances: []int32{1, 1},
925+
},
926+
{
927+
name: "1st schedule",
928+
now: mustParseTime(t, time.RFC1123, "Mon, 02 Jun 2025 03:00:00 UTC"),
929+
expectedInstances: []int32{2, 1},
930+
},
931+
{
932+
name: "2nd schedule",
933+
now: mustParseTime(t, time.RFC1123, "Mon, 02 Jun 2025 07:00:00 UTC"),
934+
expectedInstances: []int32{3, 1},
935+
},
936+
{
937+
name: "3rd schedule",
938+
now: mustParseTime(t, time.RFC1123, "Mon, 02 Jun 2025 11:00:00 UTC"),
939+
expectedInstances: []int32{1, 4},
940+
},
941+
{
942+
name: "4th schedule",
943+
now: mustParseTime(t, time.RFC1123, "Mon, 02 Jun 2025 15:00:00 UTC"),
944+
expectedInstances: []int32{1, 5},
945+
},
946+
}
947+
948+
for _, tc := range testCases {
949+
t.Run(tc.name, func(t *testing.T) {
950+
t.Parallel()
951+
952+
templateID := uuid.New()
953+
templateVersionID := uuid.New()
954+
presetOpts1 := options{
955+
templateID: templateID,
956+
templateVersionID: templateVersionID,
957+
presetID: uuid.New(),
958+
presetName: "my-preset-1",
959+
prebuiltWorkspaceID: uuid.New(),
960+
workspaceName: "prebuilds1",
961+
}
962+
presetOpts2 := options{
963+
templateID: templateID,
964+
templateVersionID: templateVersionID,
965+
presetID: uuid.New(),
966+
presetName: "my-preset-2",
967+
prebuiltWorkspaceID: uuid.New(),
968+
workspaceName: "prebuilds2",
969+
}
970+
971+
clock := quartz.NewMock(t)
972+
clock.Set(tc.now)
973+
enableAutoscaling := func(preset database.GetTemplatePresetsWithPrebuildsRow) database.GetTemplatePresetsWithPrebuildsRow {
974+
preset.AutoscalingEnabled = true
975+
preset.AutoscalingTimezone = "UTC"
976+
return preset
977+
}
978+
presets := []database.GetTemplatePresetsWithPrebuildsRow{
979+
preset(true, 1, presetOpts1, enableAutoscaling),
980+
preset(true, 1, presetOpts2, enableAutoscaling),
981+
}
982+
schedules := []database.TemplateVersionPresetPrebuildSchedule{
983+
schedule(presets[0].ID, "* 2-4 * * 1-5", 2),
984+
schedule(presets[0].ID, "* 6-8 * * 1-5", 3),
985+
schedule(presets[1].ID, "* 10-12 * * 1-5", 4),
986+
schedule(presets[1].ID, "* 14-16 * * 1-5", 5),
987+
}
988+
989+
snapshot := prebuilds.NewGlobalSnapshot(presets, schedules, nil, nil, nil, nil, clock)
990+
991+
// Check 1st preset.
992+
{
993+
ps, err := snapshot.FilterByPreset(presetOpts1.presetID)
994+
require.NoError(t, err)
995+
996+
state := ps.CalculateState()
997+
actions, err := ps.CalculateActions(clock, backoffInterval)
998+
require.NoError(t, err)
999+
1000+
validateState(t, prebuilds.ReconciliationState{
1001+
Starting: 0,
1002+
Desired: tc.expectedInstances[0],
1003+
}, *state)
1004+
validateActions(t, []*prebuilds.ReconciliationActions{
1005+
{
1006+
ActionType: prebuilds.ActionTypeCreate,
1007+
Create: tc.expectedInstances[0],
1008+
},
1009+
}, actions)
1010+
}
1011+
1012+
// Check 2nd preset.
1013+
{
1014+
ps, err := snapshot.FilterByPreset(presetOpts2.presetID)
1015+
require.NoError(t, err)
1016+
1017+
state := ps.CalculateState()
1018+
actions, err := ps.CalculateActions(clock, backoffInterval)
1019+
require.NoError(t, err)
1020+
1021+
validateState(t, prebuilds.ReconciliationState{
1022+
Starting: 0,
1023+
Desired: tc.expectedInstances[1],
1024+
}, *state)
1025+
validateActions(t, []*prebuilds.ReconciliationActions{
1026+
{
1027+
ActionType: prebuilds.ActionTypeCreate,
1028+
Create: tc.expectedInstances[1],
1029+
},
1030+
}, actions)
1031+
}
1032+
})
1033+
}
1034+
}
1035+
9081036
func TestMatchesCron(t *testing.T) {
9091037
t.Parallel()
9101038
testCases := []struct {
@@ -1294,6 +1422,15 @@ func preset(active bool, instances int32, opts options, muts ...func(row databas
12941422
return entry
12951423
}
12961424

1425+
func schedule(presetID uuid.UUID, cronExpr string, instances int32) database.TemplateVersionPresetPrebuildSchedule {
1426+
return database.TemplateVersionPresetPrebuildSchedule{
1427+
ID: uuid.New(),
1428+
PresetID: presetID,
1429+
CronExpression: cronExpr,
1430+
Instances: instances,
1431+
}
1432+
}
1433+
12971434
func prebuiltWorkspace(
12981435
opts options,
12991436
clock quartz.Clock,

enterprise/coderd/prebuilds/reconcile.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -392,6 +392,7 @@ func (c *StoreReconciler) SnapshotState(ctx context.Context, store database.Stor
392392
allPrebuildsInProgress,
393393
presetsBackoff,
394394
hardLimitedPresets,
395+
c.clock,
395396
)
396397
return nil
397398
}, &database.TxOptions{

0 commit comments

Comments
 (0)