@@ -84,7 +84,7 @@ func TestNoPrebuilds(t *testing.T) {
84
84
preset (true ,0 ,current ),
85
85
}
86
86
87
- snapshot := prebuilds .NewGlobalSnapshot (presets ,nil ,nil ,nil ,nil ,nil )
87
+ snapshot := prebuilds .NewGlobalSnapshot (presets ,nil ,nil ,nil ,nil ,nil , quartz . NewMock ( t ) )
88
88
ps ,err := snapshot .FilterByPreset (current .presetID )
89
89
require .NoError (t ,err )
90
90
@@ -106,7 +106,7 @@ func TestNetNew(t *testing.T) {
106
106
preset (true ,1 ,current ),
107
107
}
108
108
109
- snapshot := prebuilds .NewGlobalSnapshot (presets ,nil ,nil ,nil ,nil ,nil )
109
+ snapshot := prebuilds .NewGlobalSnapshot (presets ,nil ,nil ,nil ,nil ,nil , quartz . NewMock ( t ) )
110
110
ps ,err := snapshot .FilterByPreset (current .presetID )
111
111
require .NoError (t ,err )
112
112
@@ -148,7 +148,7 @@ func TestOutdatedPrebuilds(t *testing.T) {
148
148
var inProgress []database.CountInProgressPrebuildsRow
149
149
150
150
// 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 ) )
152
152
ps ,err := snapshot .FilterByPreset (outdated .presetID )
153
153
require .NoError (t ,err )
154
154
@@ -214,7 +214,7 @@ func TestDeleteOutdatedPrebuilds(t *testing.T) {
214
214
}
215
215
216
216
// 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 ) )
218
218
ps ,err := snapshot .FilterByPreset (outdated .presetID )
219
219
require .NoError (t ,err )
220
220
@@ -459,7 +459,7 @@ func TestInProgressActions(t *testing.T) {
459
459
}
460
460
461
461
// 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 ) )
463
463
ps ,err := snapshot .FilterByPreset (current .presetID )
464
464
require .NoError (t ,err )
465
465
@@ -502,7 +502,7 @@ func TestExtraneous(t *testing.T) {
502
502
var inProgress []database.CountInProgressPrebuildsRow
503
503
504
504
// 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 ) )
506
506
ps ,err := snapshot .FilterByPreset (current .presetID )
507
507
require .NoError (t ,err )
508
508
@@ -683,7 +683,7 @@ func TestExpiredPrebuilds(t *testing.T) {
683
683
}
684
684
685
685
// 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 ) )
687
687
ps ,err := snapshot .FilterByPreset (current .presetID )
688
688
require .NoError (t ,err )
689
689
@@ -719,7 +719,7 @@ func TestDeprecated(t *testing.T) {
719
719
var inProgress []database.CountInProgressPrebuildsRow
720
720
721
721
// 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 ) )
723
723
ps ,err := snapshot .FilterByPreset (current .presetID )
724
724
require .NoError (t ,err )
725
725
@@ -772,7 +772,7 @@ func TestLatestBuildFailed(t *testing.T) {
772
772
}
773
773
774
774
// 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 ) )
776
776
psCurrent ,err := snapshot .FilterByPreset (current .presetID )
777
777
require .NoError (t ,err )
778
778
@@ -865,7 +865,7 @@ func TestMultiplePresetsPerTemplateVersion(t *testing.T) {
865
865
},
866
866
}
867
867
868
- snapshot := prebuilds .NewGlobalSnapshot (presets ,nil ,nil ,inProgress ,nil ,nil )
868
+ snapshot := prebuilds .NewGlobalSnapshot (presets ,nil ,nil ,inProgress ,nil ,nil , quartz . NewMock ( t ) )
869
869
870
870
// Nothing has to be created for preset 1.
871
871
{
@@ -905,6 +905,134 @@ func TestMultiplePresetsPerTemplateVersion(t *testing.T) {
905
905
}
906
906
}
907
907
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
+
908
1036
func TestMatchesCron (t * testing.T ) {
909
1037
t .Parallel ()
910
1038
testCases := []struct {
@@ -1294,6 +1422,15 @@ func preset(active bool, instances int32, opts options, muts ...func(row databas
1294
1422
return entry
1295
1423
}
1296
1424
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
+
1297
1434
func prebuiltWorkspace (
1298
1435
opts options ,
1299
1436
clock quartz.Clock ,