add an integration test for prebuilt workspace agent reinitialization

coder · SasSwart · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025 · Apr 24, 2025
commit cebd5db97cba54976d84b8232c97fe2fe884668c
diff --git a/cli/agent_test.go b/cli/agent_test.go
@@ -22,6 +22,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/prebuilds"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/codersdk/workspacesdk"
@@ -332,14 +333,21 @@ func TestAgent_Prebuild(t *testing.T) {
 		Pubsub:   pubsub,
 	})
 	user := coderdtest.CreateFirstUser(t, client)
-	r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
+	presetID := uuid.New()
+	tv := dbfake.TemplateVersion(t, db).Seed(database.TemplateVersion{
 		OrganizationID: user.OrganizationID,
-		OwnerID:        user.UserID,
+		CreatedBy:      user.UserID,
+	}).Preset(database.TemplateVersionPreset{
+		ID: presetID,
+	}).Do()
+	r := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{
+		OwnerID:    prebuilds.SystemUserID,
+		TemplateID: tv.Template.ID,
 	}).WithAgent(func(a []*proto.Agent) []*proto.Agent {
 		a[0].Scripts = []*proto.Script{
 			{
 				DisplayName: "Prebuild Test Script",
-				Script:      "sleep 5", // Make reinitiazation take long enough to assert that it happened
+				Script:      "sleep 5", // Make reinitialization take long enough to assert that it happened
 				RunOnStart:  true,
 			},
 		}

diff --git a/coderd/database/dbfake/dbfake.go b/coderd/database/dbfake/dbfake.go
@@ -294,6 +294,8 @@ type TemplateVersionBuilder struct {
 	ps                 pubsub.Pubsub
 	resources          []*sdkproto.Resource
 	params             []database.TemplateVersionParameter
+	presets            []database.TemplateVersionPreset
+	presetParams       []database.TemplateVersionPresetParameter
 	promote            bool
 	autoCreateTemplate bool
 }
@@ -339,6 +341,13 @@ func (t TemplateVersionBuilder) Params(ps ...database.TemplateVersionParameter)
 	return t
 }
 
+func (t TemplateVersionBuilder) Preset(preset database.TemplateVersionPreset, params ...database.TemplateVersionPresetParameter) TemplateVersionBuilder {
+	// nolint: revive // returns modified struct
+	t.presets = append(t.presets, preset)
+	t.presetParams = append(t.presetParams, params...)
+	return t
+}
+
 func (t TemplateVersionBuilder) SkipCreateTemplate() TemplateVersionBuilder {
 	// nolint: revive // returns modified struct
 	t.autoCreateTemplate = false
@@ -378,6 +387,25 @@ func (t TemplateVersionBuilder) Do() TemplateVersionResponse {
 		require.NoError(t.t, err)
 	}
 
+	for _, preset := range t.presets {
+		dbgen.Preset(t.t, t.db, database.InsertPresetParams{
+			ID:                  preset.ID,
+			TemplateVersionID:   version.ID,
+			Name:                preset.Name,
+			CreatedAt:           version.CreatedAt,
+			DesiredInstances:    preset.DesiredInstances,
+			InvalidateAfterSecs: preset.InvalidateAfterSecs,
+		})
+	}
+
+	for _, presetParam := range t.presetParams {
+		dbgen.PresetParameter(t.t, t.db, database.InsertPresetParametersParams{
+			TemplateVersionPresetID: presetParam.TemplateVersionPresetID,
+			Names:                   []string{presetParam.Name},
+			Values:                  []string{presetParam.Value},
+		})
+	}
+
 	payload, err := json.Marshal(provisionerdserver.TemplateVersionImportJob{
 		TemplateVersionID: t.seed.ID,
 	})

diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go
@@ -1198,6 +1198,7 @@ func TelemetryItem(t testing.TB, db database.Store, seed database.TelemetryItem)
 
 func Preset(t testing.TB, db database.Store, seed database.InsertPresetParams) database.TemplateVersionPreset {
 	preset, err := db.InsertPreset(genCtx, database.InsertPresetParams{
+		ID:                  takeFirst(seed.ID, uuid.New()),
 		TemplateVersionID:   takeFirst(seed.TemplateVersionID, uuid.New()),
 		Name:                takeFirst(seed.Name, testutil.GetRandomName(t)),
 		CreatedAt:           takeFirst(seed.CreatedAt, dbtime.Now()),

diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/presets.sql b/coderd/database/queries/presets.sql
@@ -1,12 +1,14 @@
 -- name: InsertPreset :one
 INSERT INTO template_version_presets (
+	id,
 	template_version_id,
 	name,
 	created_at,
 	desired_instances,
 	invalidate_after_secs
 )
 VALUES (
+	@id,
 	@template_version_id,
 	@name,
 	@created_at,
@@ -16,8 +18,9 @@ VALUES (
 
 -- name: InsertPresetParameters :many
 INSERT INTO
-	template_version_preset_parameters (template_version_preset_id, name, value)
+	template_version_preset_parameters (id, template_version_preset_id, name, value)
 SELECT
+	@id,
 	@template_version_preset_id,
 	unnest(@names :: TEXT[]),
 	unnest(@values :: TEXT[])