add tests for prebuilds in the provisionerdserver

coder · SasSwart · Apr 7, 2025 · Mar 12, 2025 · Mar 13, 2025 · Mar 13, 2025
commit 4540a55e322e118ad6df55b7126665ff5fab1b24
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -2131,14 +2131,14 @@ func (q *querier) GetPresetByWorkspaceBuildID(ctx context.Context, workspaceID u
 	return q.db.GetPresetByWorkspaceBuildID(ctx, workspaceID)
 }
 
-func (q *querier) GetPresetParametersByTemplateVersionID(ctx context.Context, templateVersionID uuid.UUID) ([]database.TemplateVersionPresetParameter, error) {
+func (q *querier) GetPresetParametersByTemplateVersionID(ctx context.Context, args database.GetPresetParametersByTemplateVersionIDParams) ([]database.TemplateVersionPresetParameter, error) {
 	// An actor can read template version presets if they can read the related template version.
-	_, err := q.GetTemplateVersionByID(ctx, templateVersionID)
+	_, err := q.GetTemplateVersionByID(ctx, args.TemplateVersionID)
 	if err != nil {
 		return nil, err
 	}
 
-	return q.db.GetPresetParametersByTemplateVersionID(ctx, templateVersionID)
+	return q.db.GetPresetParametersByTemplateVersionID(ctx, args)
 }
 
 func (q *querier) GetPresetsBackoff(ctx context.Context, lookback time.Time) ([]database.GetPresetsBackoffRow, error) {

diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go
@@ -1748,7 +1748,7 @@ func (s *MethodTestSuite) TestUser() {
 		check.Args(database.DeleteCustomRoleParams{
 			Name: customRole.Name,
 		}).Asserts(
-			// fails immediately, missing organization id
+		// fails immediately, missing organization id
 		).Errors(dbauthz.NotAuthorizedError{Err: xerrors.New("custom roles must belong to an organization")})
 	}))
 	s.Run("Blank/UpdateCustomRole", s.Subtest(func(db database.Store, check *expects) {
@@ -1779,7 +1779,7 @@ func (s *MethodTestSuite) TestUser() {
 				codersdk.ResourceWorkspace: {codersdk.ActionRead},
 			}), convertSDKPerm),
 		}).Asserts(
-			// fails immediately, missing organization id
+		// fails immediately, missing organization id
 		).Errors(dbauthz.NotAuthorizedError{Err: xerrors.New("custom roles must belong to an organization")})
 	}))
 	s.Run("OrgPermissions/UpdateCustomRole", s.Subtest(func(db database.Store, check *expects) {
@@ -1832,7 +1832,7 @@ func (s *MethodTestSuite) TestUser() {
 				codersdk.ResourceWorkspace: {codersdk.ActionRead},
 			}), convertSDKPerm),
 		}).Asserts(
-			// fails immediately, missing organization id
+		// fails immediately, missing organization id
 		).Errors(dbauthz.NotAuthorizedError{Err: xerrors.New("custom roles must belong to an organization")})
 	}))
 	s.Run("OrgPermissions/InsertCustomRole", s.Subtest(func(db database.Store, check *expects) {
@@ -3843,7 +3843,7 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 	s.Run("GetProvisionerJobsCreatedAfter", s.Subtest(func(db database.Store, check *expects) {
 		// TODO: add provisioner job resource type
 		_ = dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{CreatedAt: time.Now().Add(-time.Hour)})
-		check.Args(time.Now()).Asserts( /*rbac.ResourceSystem, policy.ActionRead*/)
+		check.Args(time.Now()).Asserts( /*rbac.ResourceSystem, policy.ActionRead*/ )
 	}))
 	s.Run("GetTemplateVersionsByIDs", s.Subtest(func(db database.Store, check *expects) {
 		dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)
@@ -3934,10 +3934,14 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 			Values:                  []string{"test"},
 		})
 		require.NoError(s.T(), err)
-		presetParameters, err := db.GetPresetParametersByTemplateVersionID(ctx, templateVersion.ID)
+		presetParameters, err := db.GetPresetParametersByTemplateVersionID(ctx, database.GetPresetParametersByTemplateVersionIDParams{
+			TemplateVersionID: templateVersion.ID,
+		})
 		require.NoError(s.T(), err)
 
-		check.Args(templateVersion.ID).Asserts(template.RBACObject(), policy.ActionRead).Returns(presetParameters)
+		check.Args(database.GetPresetParametersByTemplateVersionIDParams{
+			TemplateVersionID: templateVersion.ID,
+		}).Asserts(template.RBACObject(), policy.ActionRead).Returns(presetParameters)
 	}))
 	s.Run("GetPresetsByTemplateVersionID", s.Subtest(func(db database.Store, check *expects) {
 		ctx := context.Background()
@@ -4020,7 +4024,7 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 		a := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{})
 		b := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{})
 		check.Args([]uuid.UUID{a.ID, b.ID}).
-			Asserts( /*rbac.ResourceSystem, policy.ActionRead*/).
+			Asserts( /*rbac.ResourceSystem, policy.ActionRead*/ ).
 			Returns(slice.New(a, b))
 	}))
 	s.Run("InsertWorkspaceAgent", s.Subtest(func(db database.Store, check *expects) {
@@ -4065,22 +4069,22 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 			OrganizationID:  j.OrganizationID,
 			Types:           []database.ProvisionerType{j.Provisioner},
 			ProvisionerTags: must(json.Marshal(j.Tags)),
-		}).Asserts( /*rbac.ResourceSystem, policy.ActionUpdate*/)
+		}).Asserts( /*rbac.ResourceSystem, policy.ActionUpdate*/ )
 	}))
 	s.Run("UpdateProvisionerJobWithCompleteByID", s.Subtest(func(db database.Store, check *expects) {
 		// TODO: we need to create a ProvisionerJob resource
 		j := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{})
 		check.Args(database.UpdateProvisionerJobWithCompleteByIDParams{
 			ID: j.ID,
-		}).Asserts( /*rbac.ResourceSystem, policy.ActionUpdate*/)
+		}).Asserts( /*rbac.ResourceSystem, policy.ActionUpdate*/ )
 	}))
 	s.Run("UpdateProvisionerJobByID", s.Subtest(func(db database.Store, check *expects) {
 		// TODO: we need to create a ProvisionerJob resource
 		j := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{})
 		check.Args(database.UpdateProvisionerJobByIDParams{
 			ID:        j.ID,
 			UpdatedAt: time.Now(),
-		}).Asserts( /*rbac.ResourceSystem, policy.ActionUpdate*/)
+		}).Asserts( /*rbac.ResourceSystem, policy.ActionUpdate*/ )
 	}))
 	s.Run("InsertProvisionerJob", s.Subtest(func(db database.Store, check *expects) {
 		dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)
@@ -4091,21 +4095,21 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 			StorageMethod: database.ProvisionerStorageMethodFile,
 			Type:          database.ProvisionerJobTypeWorkspaceBuild,
 			Input:         json.RawMessage("{}"),
-		}).Asserts( /*rbac.ResourceSystem, policy.ActionCreate*/)
+		}).Asserts( /*rbac.ResourceSystem, policy.ActionCreate*/ )
 	}))
 	s.Run("InsertProvisionerJobLogs", s.Subtest(func(db database.Store, check *expects) {
 		// TODO: we need to create a ProvisionerJob resource
 		j := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{})
 		check.Args(database.InsertProvisionerJobLogsParams{
 			JobID: j.ID,
-		}).Asserts( /*rbac.ResourceSystem, policy.ActionCreate*/)
+		}).Asserts( /*rbac.ResourceSystem, policy.ActionCreate*/ )
 	}))
 	s.Run("InsertProvisionerJobTimings", s.Subtest(func(db database.Store, check *expects) {
 		// TODO: we need to create a ProvisionerJob resource
 		j := dbgen.ProvisionerJob(s.T(), db, nil, database.ProvisionerJob{})
 		check.Args(database.InsertProvisionerJobTimingsParams{
 			JobID: j.ID,
-		}).Asserts( /*rbac.ResourceSystem, policy.ActionCreate*/)
+		}).Asserts( /*rbac.ResourceSystem, policy.ActionCreate*/ )
 	}))
 	s.Run("UpsertProvisionerDaemon", s.Subtest(func(db database.Store, check *expects) {
 		dbtestutil.DisableForeignKeysAndTriggers(s.T(), db)

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -4188,14 +4188,14 @@ func (q *FakeQuerier) GetPresetByWorkspaceBuildID(_ context.Context, workspaceBu
 	return database.TemplateVersionPreset{}, sql.ErrNoRows
 }
 
-func (q *FakeQuerier) GetPresetParametersByTemplateVersionID(_ context.Context, templateVersionID uuid.UUID) ([]database.TemplateVersionPresetParameter, error) {
+func (q *FakeQuerier) GetPresetParametersByTemplateVersionID(_ context.Context, args database.GetPresetParametersByTemplateVersionIDParams) ([]database.TemplateVersionPresetParameter, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()
 
 	presets := make([]database.TemplateVersionPreset, 0)
 	parameters := make([]database.TemplateVersionPresetParameter, 0)
 	for _, preset := range q.presets {
-		if preset.TemplateVersionID != templateVersionID {
+		if preset.TemplateVersionID != args.TemplateVersionID {
 			continue
 		}
 		presets = append(presets, preset)
@@ -8739,6 +8739,11 @@ func (q *FakeQuerier) InsertPreset(_ context.Context, arg database.InsertPresetP
 		TemplateVersionID: arg.TemplateVersionID,
 		Name:              arg.Name,
 		CreatedAt:         arg.CreatedAt,
+		DesiredInstances:  arg.DesiredInstances,
+		InvalidateAfterSecs: sql.NullInt32{
+			Int32: 0,
+			Valid: true,
+		},
 	}
 	q.presets = append(q.presets, preset)
 	return preset, nil

diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/...base/migrations/000312_prebuilds.down.sql → ...base/migrations/000313_prebuilds.down.sql b/...base/migrations/000312_prebuilds.down.sql → ...base/migrations/000313_prebuilds.down.sql
diff --git a/...tabase/migrations/000312_prebuilds.up.sql → ...tabase/migrations/000313_prebuilds.up.sql b/...tabase/migrations/000312_prebuilds.up.sql → ...tabase/migrations/000313_prebuilds.up.sql
diff --git a/...grations/000313_preset_prebuilds.down.sql → ...grations/000314_preset_prebuilds.down.sql b/...grations/000313_preset_prebuilds.down.sql → ...grations/000314_preset_prebuilds.down.sql
diff --git a/...migrations/000313_preset_prebuilds.up.sql → ...migrations/000314_preset_prebuilds.up.sql b/...migrations/000313_preset_prebuilds.up.sql → ...migrations/000314_preset_prebuilds.up.sql
diff --git a/...a/fixtures/000313_preset_prebuilds.up.sql → ...a/fixtures/000314_preset_prebuilds.up.sql b/...a/fixtures/000313_preset_prebuilds.up.sql → ...a/fixtures/000314_preset_prebuilds.up.sql
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
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
@@ -47,4 +47,5 @@ FROM
 	template_version_preset_parameters
 	INNER JOIN template_version_presets ON template_version_preset_parameters.template_version_preset_id = template_version_presets.id
 WHERE
-	template_version_presets.template_version_id = @template_version_id;
+	template_version_presets.template_version_id = @template_version_id
+	AND (sqlc.narg('preset_id')::uuid IS NULL OR template_version_presets.id = sqlc.narg('preset_id'));
diff --git a/coderd/presets.go b/coderd/presets.go
@@ -3,6 +3,7 @@ package coderd
 import (
 	"net/http"
 
+	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/codersdk"
@@ -29,7 +30,9 @@ func (api *API) templateVersionPresets(rw http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	presetParams, err := api.Database.GetPresetParametersByTemplateVersionID(ctx, templateVersion.ID)
+	presetParams, err := api.Database.GetPresetParametersByTemplateVersionID(ctx, database.GetPresetParametersByTemplateVersionIDParams{
+		TemplateVersionID: templateVersion.ID,
+	})
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error fetching template version presets.",

diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -1733,6 +1733,34 @@ func TestInsertWorkspacePresetsAndParameters(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "one preset, no parameters, requesting prebuilds",
+			givenPresets: []*sdkproto.Preset{
+				{
+					Name: "preset1",
+					Prebuild: &sdkproto.Prebuild{
+						Instances: 1,
+					},
+				},
+			},
+		},
+		{
+			name: "one preset with multiple parameters, requesting 0 prebuilds",
+			givenPresets: []*sdkproto.Preset{
+				{
+					Name: "preset1",
+					Parameters: []*sdkproto.PresetParameter{
+						{
+							Name:  "param1",
+							Value: "value1",
+						},
+					},
+					Prebuild: &sdkproto.Prebuild{
+						Instances: 0,
+					},
+				},
+			},
+		},
 		{
 			name: "one preset with multiple parameters",
 			givenPresets: []*sdkproto.Preset{
@@ -1751,6 +1779,27 @@ func TestInsertWorkspacePresetsAndParameters(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "one preset, multiple parameters, requesting prebuilds",
+			givenPresets: []*sdkproto.Preset{
+				{
+					Name: "preset1",
+					Parameters: []*sdkproto.PresetParameter{
+						{
+							Name:  "param1",
+							Value: "value1",
+						},
+						{
+							Name:  "param2",
+							Value: "value2",
+						},
+					},
+					Prebuild: &sdkproto.Prebuild{
+						Instances: 1,
+					},
+				},
+			},
+		},
 		{
 			name: "multiple presets with parameters",
 			givenPresets: []*sdkproto.Preset{
@@ -1766,6 +1815,9 @@ func TestInsertWorkspacePresetsAndParameters(t *testing.T) {
 							Value: "value2",
 						},
 					},
+					Prebuild: &sdkproto.Prebuild{
+						Instances: 1,
+					},
 				},
 				{
 					Name: "preset2",
@@ -1820,42 +1872,40 @@ func TestInsertWorkspacePresetsAndParameters(t *testing.T) {
 			require.Len(t, gotPresets, len(c.givenPresets))
 
 			for _, givenPreset := range c.givenPresets {
-				foundMatch := false
+				var foundPreset *database.TemplateVersionPreset = nil
 				for _, gotPreset := range gotPresets {
 					if givenPreset.Name == gotPreset.Name {
-						foundMatch = true
+						foundPreset = &gotPreset
 						break
 					}
 				}
-				require.True(t, foundMatch, "preset %s not found in parameters", givenPreset.Name)
-			}
+				require.NotNil(t, foundPreset, "preset %s not found in parameters", givenPreset.Name)
 
-			gotPresetParameters, err := db.GetPresetParametersByTemplateVersionID(ctx, templateVersion.ID)
-			require.NoError(t, err)
+				gotPresetParameters, err := db.GetPresetParametersByTemplateVersionID(ctx, database.GetPresetParametersByTemplateVersionIDParams{
+					TemplateVersionID: templateVersion.ID,
+					PresetID:          uuid.NullUUID{UUID: foundPreset.ID, Valid: true},
+				})
+				require.NoError(t, err)
+				require.Len(t, gotPresetParameters, len(givenPreset.Parameters))
 
-			for _, givenPreset := range c.givenPresets {
 				for _, givenParameter := range givenPreset.Parameters {
 					foundMatch := false
 					for _, gotParameter := range gotPresetParameters {
 						nameMatches := givenParameter.Name == gotParameter.Name
 						valueMatches := givenParameter.Value == gotParameter.Value
-
-						// ensure that preset parameters are matched to the correct preset:
-						var gotPreset database.TemplateVersionPreset
-						for _, preset := range gotPresets {
-							if preset.ID == gotParameter.TemplateVersionPresetID {
-								gotPreset = preset
-								break
-							}
-						}
-						presetMatches := gotPreset.Name == givenPreset.Name
-
-						if nameMatches && valueMatches && presetMatches {
+						if nameMatches && valueMatches {
 							foundMatch = true
 							break
 						}
 					}
-					require.True(t, foundMatch, "preset parameter %s not found in presets", givenParameter.Name)
+					require.True(t, foundMatch, "preset parameter %s not found in parameters", givenParameter.Name)
+				}
+				if givenPreset.Prebuild == nil {
+					require.False(t, foundPreset.DesiredInstances.Valid)
+				}
+				if givenPreset.Prebuild != nil {
+					require.True(t, foundPreset.DesiredInstances.Valid)
+					require.Equal(t, givenPreset.Prebuild.Instances, foundPreset.DesiredInstances.Int32)
 				}
 			}
 		})