From c7307dc242ecc9ffbc556d9cf8ddece1bffed046 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Mon, 4 Aug 2025 16:10:14 +0000 Subject: [PATCH] refactor: prebuilds tests --- cli/delete_test.go | 28 +---- coderd/autobuild/lifecycle_executor_test.go | 31 +---- coderd/database/dbfake/dbfake.go | 66 ++++++++++ .../coderd/prebuilds/metricscollector_test.go | 6 +- enterprise/coderd/prebuilds/reconcile_test.go | 116 +++--------------- 5 files changed, 93 insertions(+), 154 deletions(-) diff --git a/cli/delete_test.go b/cli/delete_test.go index c01893419f80f..9d8d5002d00a8 100644 --- a/cli/delete_test.go +++ b/cli/delete_test.go @@ -9,6 +9,8 @@ import ( "testing" "time" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/google/uuid" "github.com/coder/coder/v2/coderd/database" @@ -246,7 +248,7 @@ func TestDelete(t *testing.T) { // Given a template version with a preset and a template version := coderdtest.CreateTemplateVersion(t, client, orgID, nil) coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) - preset := setupTestDBPreset(t, db, version.ID) + preset := dbfake.NewPreset(t, db, version.ID).WithDesiredInstances(1).Do() template := coderdtest.CreateTemplate(t, client, orgID, version.ID) cases := []struct { @@ -378,30 +380,6 @@ func TestDelete(t *testing.T) { }) } -func setupTestDBPreset( - t *testing.T, - db database.Store, - templateVersionID uuid.UUID, -) database.TemplateVersionPreset { - t.Helper() - - preset := dbgen.Preset(t, db, database.InsertPresetParams{ - TemplateVersionID: templateVersionID, - Name: "preset-test", - DesiredInstances: sql.NullInt32{ - Valid: true, - Int32: 1, - }, - }) - dbgen.PresetParameter(t, db, database.InsertPresetParametersParams{ - TemplateVersionPresetID: preset.ID, - Names: []string{"test"}, - Values: []string{"test"}, - }) - - return preset -} - func setupTestDBWorkspace( t *testing.T, clock quartz.Clock, diff --git a/coderd/autobuild/lifecycle_executor_test.go b/coderd/autobuild/lifecycle_executor_test.go index 0229a907cbb2e..b41175f3d8073 100644 --- a/coderd/autobuild/lifecycle_executor_test.go +++ b/coderd/autobuild/lifecycle_executor_test.go @@ -7,6 +7,8 @@ import ( "testing" "time" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/rbac" @@ -1231,7 +1233,7 @@ func TestExecutorPrebuilds(t *testing.T) { template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID) // Database setup of a preset with a prebuild instance - preset := setupTestDBPreset(t, db, version.ID, int32(1)) + preset := dbfake.NewPreset(t, db, version.ID).WithDesiredInstances(1).Do() // Given: a running prebuilt workspace with a deadline and ready to be claimed dbPrebuild := setupTestDBPrebuiltWorkspace( @@ -1307,7 +1309,7 @@ func TestExecutorPrebuilds(t *testing.T) { template := coderdtest.CreateTemplate(t, client, owner.OrganizationID, version.ID) // Database setup of a preset with a prebuild instance - preset := setupTestDBPreset(t, db, version.ID, int32(1)) + preset := dbfake.NewPreset(t, db, version.ID).WithDesiredInstances(1).Do() // Given: prebuilt workspace is stopped and set to autostart daily at midnight sched := mustSchedule(t, "CRON_TZ=UTC 0 0 * * *") @@ -1384,31 +1386,6 @@ func TestExecutorPrebuilds(t *testing.T) { }) } -func setupTestDBPreset( - t *testing.T, - db database.Store, - templateVersionID uuid.UUID, - desiredInstances int32, -) database.TemplateVersionPreset { - t.Helper() - - preset := dbgen.Preset(t, db, database.InsertPresetParams{ - TemplateVersionID: templateVersionID, - Name: "preset-test", - DesiredInstances: sql.NullInt32{ - Valid: true, - Int32: desiredInstances, - }, - }) - dbgen.PresetParameter(t, db, database.InsertPresetParametersParams{ - TemplateVersionPresetID: preset.ID, - Names: []string{"test-name"}, - Values: []string{"test-value"}, - }) - - return preset -} - type SetupPrebuiltOptions struct { AutostartSchedule sql.NullString IsStopped bool diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go index 6d99005fb3334..68ad43cacde15 100644 --- a/coderd/database/dbfake/dbfake.go +++ b/coderd/database/dbfake/dbfake.go @@ -563,6 +563,72 @@ func (b JobCompleteBuilder) Do() JobCompleteResponse { return r } +type PresetBuilder struct { + t *testing.T + db database.Store + presetName string + templateVersionID uuid.UUID + desiredInstances int32 + // Optional parameters + ttl *int32 +} + +func NewPreset(t *testing.T, db database.Store, templateVersionID uuid.UUID) PresetBuilder { + return PresetBuilder{ + t: t, + db: db, + presetName: uuid.New().String(), + templateVersionID: templateVersionID, + desiredInstances: 0, // Default to 0 prebuild instances + } +} + +func (b PresetBuilder) WithPresetName(name string) PresetBuilder { + b.presetName = name + return b +} + +func (b PresetBuilder) WithDesiredInstances(instances int32) PresetBuilder { + b.desiredInstances = instances + return b +} + +func (b PresetBuilder) WithTTL(ttl int32) PresetBuilder { + b.ttl = &ttl + return b +} + +func (b PresetBuilder) Do() database.TemplateVersionPreset { + // Using only required fields for testing; other fields will use DB defaults. + //nolint:exhaustruct + insertPresetParams := database.InsertPresetParams{ + TemplateVersionID: b.templateVersionID, + Name: b.presetName, + DesiredInstances: sql.NullInt32{ + Valid: true, + Int32: b.desiredInstances, + }, + } + + // Handle optional TTL + if b.ttl != nil { + insertPresetParams.InvalidateAfterSecs = sql.NullInt32{ + Valid: true, + Int32: *b.ttl, + } + } + + preset := dbgen.Preset(b.t, b.db, insertPresetParams) + + dbgen.PresetParameter(b.t, b.db, database.InsertPresetParametersParams{ + TemplateVersionPresetID: preset.ID, + Names: []string{"test-name"}, + Values: []string{"test-value"}, + }) + + return preset +} + func must[V any](v V, err error) V { if err != nil { panic(err) diff --git a/enterprise/coderd/prebuilds/metricscollector_test.go b/enterprise/coderd/prebuilds/metricscollector_test.go index 1e9f3f5082806..c423a7aacd2e5 100644 --- a/enterprise/coderd/prebuilds/metricscollector_test.go +++ b/enterprise/coderd/prebuilds/metricscollector_test.go @@ -5,6 +5,8 @@ import ( "slices" "testing" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/google/uuid" "github.com/stretchr/testify/require" "tailscale.com/types/ptr" @@ -222,7 +224,7 @@ func TestMetricsCollector(t *testing.T) { for i := 0; i < numTemplates; i++ { org, template := setupTestDBTemplate(t, db, ownerID, templateDeleted) templateVersionID := setupTestDBTemplateVersion(ctx, t, clock, db, pubsub, org.ID, ownerID, template.ID) - preset := setupTestDBPreset(t, db, templateVersionID, 1, uuid.New().String()) + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(1).Do() workspace, _ := setupTestDBWorkspace( t, clock, db, pubsub, transition, jobStatus, org.ID, preset, template.ID, templateVersionID, initiatorID, ownerID, @@ -350,7 +352,7 @@ func TestMetricsCollector_DuplicateTemplateNames(t *testing.T) { setupTemplateWithDeps := func() database.Template { template := setupTestDBTemplateWithinOrg(t, db, test.ownerID, false, "default-template", defaultOrg) templateVersionID := setupTestDBTemplateVersion(ctx, t, clock, db, pubsub, defaultOrg.ID, test.ownerID, template.ID) - preset := setupTestDBPreset(t, db, templateVersionID, 1, "default-preset") + preset := dbfake.NewPreset(t, db, templateVersionID).WithPresetName(presetName).WithDesiredInstances(1).Do() workspace, _ := setupTestDBWorkspace( t, clock, db, pubsub, test.transition, test.jobStatus, defaultOrg.ID, preset, template.ID, templateVersionID, test.initiatorID, test.ownerID, diff --git a/enterprise/coderd/prebuilds/reconcile_test.go b/enterprise/coderd/prebuilds/reconcile_test.go index 8d2a81e1ade83..da358ddb96994 100644 --- a/enterprise/coderd/prebuilds/reconcile_test.go +++ b/enterprise/coderd/prebuilds/reconcile_test.go @@ -10,6 +10,8 @@ import ( "testing" "time" + "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/prometheus/client_golang/prometheus" "github.com/stretchr/testify/assert" "golang.org/x/xerrors" @@ -352,13 +354,7 @@ func TestPrebuildReconciliation(t *testing.T) { ownerID, template.ID, ) - preset := setupTestDBPreset( - t, - db, - templateVersionID, - 1, - uuid.New().String(), - ) + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(1).Do() prebuild, _ := setupTestDBPrebuild( t, clock, @@ -479,20 +475,8 @@ func TestMultiplePresetsPerTemplateVersion(t *testing.T) { ownerID, template.ID, ) - preset := setupTestDBPreset( - t, - db, - templateVersionID, - 4, - uuid.New().String(), - ) - preset2 := setupTestDBPreset( - t, - db, - templateVersionID, - 10, - uuid.New().String(), - ) + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(4).Do() + preset2 := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(10).Do() prebuildIDs := make([]uuid.UUID, 0) for i := 0; i < int(preset.DesiredInstances.Int32); i++ { prebuild, _ := setupTestDBPrebuild( @@ -719,13 +703,7 @@ func TestInvalidPreset(t *testing.T) { DefaultValue: "", Required: true, }) - setupTestDBPreset( - t, - db, - templateVersionID, - 1, - uuid.New().String(), - ) + dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(1).Do() // Run the reconciliation multiple times to ensure idempotency // 8 was arbitrary, but large enough to reasonably trust the result @@ -766,7 +744,7 @@ func TestDeletionOfPrebuiltWorkspaceWithInvalidPreset(t *testing.T) { }) org, template := setupTestDBTemplate(t, db, ownerID, templateDeleted) templateVersionID := setupTestDBTemplateVersion(ctx, t, clock, db, pubSub, org.ID, ownerID, template.ID) - preset := setupTestDBPreset(t, db, templateVersionID, 1, uuid.New().String()) + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(1).Do() prebuiltWorkspace, _ := setupTestDBPrebuild( t, clock, @@ -864,7 +842,7 @@ func TestSkippingHardLimitedPresets(t *testing.T) { }) org, template := setupTestDBTemplate(t, db, ownerID, templateDeleted) templateVersionID := setupTestDBTemplateVersion(ctx, t, clock, db, pubSub, org.ID, ownerID, template.ID) - preset := setupTestDBPreset(t, db, templateVersionID, 1, uuid.New().String()) + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(1).Do() // Create a failed prebuild workspace that counts toward the hard failure limit. setupTestDBPrebuild( @@ -1008,7 +986,7 @@ func TestHardLimitedPresetShouldNotBlockDeletion(t *testing.T) { }) org, template := setupTestDBTemplate(t, db, ownerID, false) templateVersionID := setupTestDBTemplateVersion(ctx, t, clock, db, pubSub, org.ID, ownerID, template.ID) - preset := setupTestDBPreset(t, db, templateVersionID, 2, uuid.New().String()) + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(2).Do() // Create a successful prebuilt workspace. successfulWorkspace, _ := setupTestDBPrebuild( @@ -1210,20 +1188,8 @@ func TestRunLoop(t *testing.T) { ownerID, template.ID, ) - preset := setupTestDBPreset( - t, - db, - templateVersionID, - 4, - uuid.New().String(), - ) - preset2 := setupTestDBPreset( - t, - db, - templateVersionID, - 10, - uuid.New().String(), - ) + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(4).Do() + preset2 := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(10).Do() prebuildIDs := make([]uuid.UUID, 0) for i := 0; i < int(preset.DesiredInstances.Int32); i++ { prebuild, _ := setupTestDBPrebuild( @@ -1275,13 +1241,7 @@ func TestRunLoop(t *testing.T) { }, testutil.WaitShort, testutil.IntervalFast) // setup one more preset with 5 prebuilds - preset3 := setupTestDBPreset( - t, - db, - templateVersionID, - 5, - uuid.New().String(), - ) + preset3 := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(5).Do() newPrebuildCount := getNewPrebuildCount() // nothing changed, because we didn't trigger a new iteration of a loop require.Equal(t, preset2.DesiredInstances.Int32, newPrebuildCount) @@ -1335,7 +1295,7 @@ func TestFailedBuildBackoff(t *testing.T) { org, template := setupTestDBTemplate(t, db, userID, false) templateVersionID := setupTestDBTemplateVersion(ctx, t, clock, db, ps, org.ID, userID, template.ID) - preset := setupTestDBPreset(t, db, templateVersionID, desiredInstances, "test") + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(desiredInstances).Do() for range desiredInstances { _, _ = setupTestDBPrebuild(t, clock, db, ps, database.WorkspaceTransitionStart, database.ProvisionerJobStatusFailed, org.ID, preset, template.ID, templateVersionID) } @@ -1496,7 +1456,7 @@ func TestTrackResourceReplacement(t *testing.T) { dbgen.User(t, db, database.User{ID: userID}) org, template := setupTestDBTemplate(t, db, userID, false) templateVersionID := setupTestDBTemplateVersion(ctx, t, clock, db, ps, org.ID, userID, template.ID) - preset := setupTestDBPreset(t, db, templateVersionID, 1, "b0rked") + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(1).Do() prebuiltWorkspace, prebuild := setupTestDBPrebuild(t, clock, db, ps, database.WorkspaceTransitionStart, database.ProvisionerJobStatusSucceeded, org.ID, preset, template.ID, templateVersionID) // Given: no replacement has been tracked yet, we should not see a metric for it yet. @@ -1652,7 +1612,7 @@ func TestExpiredPrebuildsMultipleActions(t *testing.T) { ttlDuration := muchEarlier - time.Hour ttl := int32(-ttlDuration.Seconds()) - preset := setupTestDBPreset(t, db, templateVersionID, tc.desired, "b0rked", withTTL(ttl)) + preset := dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(tc.desired).WithTTL(ttl).Do() // The implementation uses time.Since(prebuild.CreatedAt) > ttl to check a prebuild expiration. // Since our mock clock defaults to a fixed time, we must align it with the current time @@ -1911,50 +1871,6 @@ func setupTestDBTemplateVersion( return templateVersion.ID } -// Preset optional parameters. -// presetOptions defines a function type for modifying InsertPresetParams. -type presetOptions func(*database.InsertPresetParams) - -// withTTL returns a presetOptions function that sets the invalidate_after_secs (TTL) field in InsertPresetParams. -func withTTL(ttl int32) presetOptions { - return func(p *database.InsertPresetParams) { - p.InvalidateAfterSecs = sql.NullInt32{Valid: true, Int32: ttl} - } -} - -func setupTestDBPreset( - t *testing.T, - db database.Store, - templateVersionID uuid.UUID, - desiredInstances int32, - presetName string, - opts ...presetOptions, -) database.TemplateVersionPreset { - t.Helper() - insertPresetParams := database.InsertPresetParams{ - TemplateVersionID: templateVersionID, - Name: presetName, - DesiredInstances: sql.NullInt32{ - Valid: true, - Int32: desiredInstances, - }, - } - - // Apply optional parameters to insertPresetParams (e.g., TTL). - for _, opt := range opts { - opt(&insertPresetParams) - } - - preset := dbgen.Preset(t, db, insertPresetParams) - - dbgen.PresetParameter(t, db, database.InsertPresetParametersParams{ - TemplateVersionPresetID: preset.ID, - Names: []string{"test"}, - Values: []string{"test"}, - }) - return preset -} - func setupTestDBPresetWithScheduling( t *testing.T, db database.Store, @@ -2290,7 +2206,7 @@ func TestReconciliationRespectsPauseSetting(t *testing.T) { OrganizationID: org.ID, }) templateVersionID := setupTestDBTemplateVersion(ctx, t, clock, db, ps, org.ID, user.ID, template.ID) - _ = setupTestDBPreset(t, db, templateVersionID, 2, "test") + dbfake.NewPreset(t, db, templateVersionID).WithDesiredInstances(2).Do() // Initially, reconciliation should create prebuilds err := reconciler.ReconcileAll(ctx)