chore: refactor ResolveAutostart tests to use dbfake (#10603)

coderd/database/dbfake/dbfake.go

@@ -16,37 +16,71 @@ import (
 	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/provisionerdserver"
+	"github.com/coder/coder/v2/coderd/rbac"
 	"github.com/coder/coder/v2/coderd/telemetry"
 	"github.com/coder/coder/v2/codersdk"
 	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
 )
 
-type WorkspaceBuilder struct {
-	t          testing.TB
-	db         database.Store
-	seed       database.Workspace
-	resources  []*sdkproto.Resource
-	agentToken string
-}
+var ownerCtx = dbauthz.As(context.Background(), rbac.Subject{
+	ID:     "owner",
+	Roles:  rbac.Roles(must(rbac.RoleNames{rbac.RoleOwner()}.Expand())),
+	Groups: []string{},
+	Scope:  rbac.ExpandableScope(rbac.ScopeAll),
+})
 
 type WorkspaceResponse struct {
 	Workspace  database.Workspace
-	Template   database.Template
 	Build      database.WorkspaceBuild
 	AgentToken string
+	TemplateVersionResponse
 }
 
-func Workspace(t testing.TB, db database.Store) WorkspaceBuilder {
-	return WorkspaceBuilder{t: t, db: db}
+// WorkspaceBuildBuilder generates workspace builds and associated
+// resources.
+type WorkspaceBuildBuilder struct {
+	t          testing.TB
+	db         database.Store
+	ps         pubsub.Pubsub
+	ws         database.Workspace
+	seed       database.WorkspaceBuild
+	resources  []*sdkproto.Resource
+	params     []database.WorkspaceBuildParameter
+	agentToken string
 }
 
-func (b WorkspaceBuilder) Seed(seed database.Workspace) WorkspaceBuilder {
+// WorkspaceBuild generates a workspace build for the provided workspace.
+// Pass a database.Workspace{} with a nil ID to also generate a new workspace.
+// Omitting the template ID on a workspace will also generate a new template
+// with a template version.
+func WorkspaceBuild(t testing.TB, db database.Store, ws database.Workspace) WorkspaceBuildBuilder {
+	return WorkspaceBuildBuilder{t: t, db: db, ws: ws}
+}
+
+func (b WorkspaceBuildBuilder) Pubsub(ps pubsub.Pubsub) WorkspaceBuildBuilder {
+	// nolint: revive // returns modified struct
+	b.ps = ps
+	return b
+}
+
+func (b WorkspaceBuildBuilder) Seed(seed database.WorkspaceBuild) WorkspaceBuildBuilder {
 	//nolint: revive // returns modified struct
 	b.seed = seed
 	return b
 }
 
-func (b WorkspaceBuilder) WithAgent(mutations ...func([]*sdkproto.Agent) []*sdkproto.Agent) WorkspaceBuilder {
+func (b WorkspaceBuildBuilder) Resource(resource ...*sdkproto.Resource) WorkspaceBuildBuilder {
+	//nolint: revive // returns modified struct
+	b.resources = append(b.resources, resource...)
+	return b
+}
+
+func (b WorkspaceBuildBuilder) Params(params ...database.WorkspaceBuildParameter) WorkspaceBuildBuilder {
+	b.params = params
+	return b
+}
+
+func (b WorkspaceBuildBuilder) WithAgent(mutations ...func([]*sdkproto.Agent) []*sdkproto.Agent) WorkspaceBuildBuilder {
 	//nolint: revive // returns modified struct
 	b.agentToken = uuid.NewString()
 	agents := []*sdkproto.Agent{{
@@ -66,75 +100,57 @@ func (b WorkspaceBuilder) WithAgent(mutations ...func([]*sdkproto.Agent) []*sdkp
 	return b
 }
 
-func (b WorkspaceBuilder) Do() WorkspaceResponse {
-	var r WorkspaceResponse
-	// This intentionally fulfills the minimum requirements of the schema.
-	// Tests can provide a custom template ID if necessary.
-	if b.seed.TemplateID == uuid.Nil {
-		r.Template = dbgen.Template(b.t, b.db, database.Template{
-			OrganizationID: b.seed.OrganizationID,
-			CreatedBy:      b.seed.OwnerID,
-		})
-		b.seed.TemplateID = r.Template.ID
-		b.seed.OwnerID = r.Template.CreatedBy
-		b.seed.OrganizationID = r.Template.OrganizationID
-	}
-	r.Workspace = dbgen.Workspace(b.t, b.db, b.seed)
-	if b.agentToken != "" {
-		r.AgentToken = b.agentToken
-		r.Build = WorkspaceBuild(b.t, b.db, r.Workspace).
-			Resource(b.resources...).
-			Do()
-	}
-	return r
-}
-
-type WorkspaceBuildBuilder struct {
-	t         testing.TB
-	db        database.Store
-	ps        pubsub.Pubsub
-	ws        database.Workspace
-	seed      database.WorkspaceBuild
-	resources []*sdkproto.Resource
-}
-
-func WorkspaceBuild(t testing.TB, db database.Store, ws database.Workspace) WorkspaceBuildBuilder {
-	return WorkspaceBuildBuilder{t: t, db: db, ws: ws}
-}
-
-func (b WorkspaceBuildBuilder) Pubsub(ps pubsub.Pubsub) WorkspaceBuildBuilder {
-	//nolint: revive // returns modified struct
-	b.ps = ps
-	return b
-}
-
-func (b WorkspaceBuildBuilder) Seed(seed database.WorkspaceBuild) WorkspaceBuildBuilder {
-	//nolint: revive // returns modified struct
-	b.seed = seed
-	return b
-}
-
-func (b WorkspaceBuildBuilder) Resource(resource ...*sdkproto.Resource) WorkspaceBuildBuilder {
-	//nolint: revive // returns modified struct
-	b.resources = append(b.resources, resource...)
-	return b
-}
-
-func (b WorkspaceBuildBuilder) Do() database.WorkspaceBuild {
+// Do generates all the resources associated with a workspace build.
+// Template and TemplateVersion will be optionally populated if no
+// TemplateID is set on the provided workspace.
+// Workspace will be optionally populated if no ID is set on the provided
+// workspace.
+func (b WorkspaceBuildBuilder) Do() WorkspaceResponse {
 	b.t.Helper()
 	jobID := uuid.New()
 	b.seed.ID = uuid.New()
 	b.seed.JobID = jobID
+
+	resp := WorkspaceResponse{
+		AgentToken: b.agentToken,
+	}
+	if b.ws.TemplateID == uuid.Nil {
+		resp.TemplateVersionResponse = TemplateVersion(b.t, b.db).
+			Resources(b.resources...).
+			Pubsub(b.ps).
+			Seed(database.TemplateVersion{
+				OrganizationID: b.ws.OrganizationID,
+				CreatedBy:      b.ws.OwnerID,
+			}).
+			Do()
+		b.ws.TemplateID = resp.Template.ID
+		b.seed.TemplateVersionID = resp.TemplateVersion.ID
+	}
+
+	// If no template version is set assume the active version.
+	if b.seed.TemplateVersionID == uuid.Nil {
+		template, err := b.db.GetTemplateByID(ownerCtx, b.ws.TemplateID)
+		require.NoError(b.t, err)
+		require.NotNil(b.t, template.ActiveVersionID, "active version ID unexpectedly nil")
+		b.seed.TemplateVersionID = template.ActiveVersionID
+	}
+
+	// No ID on the workspace implies we should generate an entry.
+	if b.ws.ID == uuid.Nil {
+		// nolint: revive
+		b.ws = dbgen.Workspace(b.t, b.db, b.ws)
+		resp.Workspace = b.ws
+	}
 	b.seed.WorkspaceID = b.ws.ID
+	b.seed.InitiatorID = takeFirst(b.seed.InitiatorID, b.ws.OwnerID)
 
 	// Create a provisioner job for the build!
 	payload, err := json.Marshal(provisionerdserver.WorkspaceProvisionJob{
 		WorkspaceBuildID: b.seed.ID,
 	})
 	require.NoError(b.t, err)
-	//nolint:gocritic // This is only used by tests.
-	ctx := dbauthz.AsSystemRestricted(context.Background())
-	job, err := b.db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
+
+	job, err := b.db.InsertProvisionerJob(ownerCtx, database.InsertProvisionerJobParams{
 		ID:             jobID,
 		CreatedAt:      dbtime.Now(),
 		UpdatedAt:      dbtime.Now(),
@@ -149,7 +165,8 @@ func (b WorkspaceBuildBuilder) Do() database.WorkspaceBuild {
 		TraceMetadata:  pqtype.NullRawMessage{},
 	})
 	require.NoError(b.t, err, "insert job")
-	err = b.db.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
+
+	err = b.db.UpdateProvisionerJobWithCompleteByID(ownerCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
 		ID:        job.ID,
 		UpdatedAt: dbtime.Now(),
 		Error:     sql.NullString{},
@@ -161,42 +178,20 @@ func (b WorkspaceBuildBuilder) Do() database.WorkspaceBuild {
 	})
 	require.NoError(b.t, err, "complete job")
 
-	// This intentionally fulfills the minimum requirements of the schema.
-	// Tests can provide a custom version ID if necessary.
-	if b.seed.TemplateVersionID == uuid.Nil {
-		jobID := uuid.New()
-		templateVersion := dbgen.TemplateVersion(b.t, b.db, database.TemplateVersion{
-			JobID:          jobID,
-			OrganizationID: b.ws.OrganizationID,
-			CreatedBy:      b.ws.OwnerID,
-			TemplateID: uuid.NullUUID{
-				UUID:  b.ws.TemplateID,
-				Valid: true,
-			},
-		})
-		payload, _ := json.Marshal(provisionerdserver.TemplateVersionImportJob{
-			TemplateVersionID: templateVersion.ID,
-		})
-		dbgen.ProvisionerJob(b.t, b.db, nil, database.ProvisionerJob{
-			ID:             jobID,
-			OrganizationID: b.ws.OrganizationID,
-			Input:          payload,
-			Type:           database.ProvisionerJobTypeTemplateVersionImport,
-			CompletedAt: sql.NullTime{
-				Time:  dbtime.Now(),
-				Valid: true,
-			},
-		})
-		ProvisionerJobResources(b.t, b.db, jobID, b.seed.Transition, b.resources...).Do()
-		b.seed.TemplateVersionID = templateVersion.ID
-	}
-	build := dbgen.WorkspaceBuild(b.t, b.db, b.seed)
+	resp.Build = dbgen.WorkspaceBuild(b.t, b.db, b.seed)
 	ProvisionerJobResources(b.t, b.db, job.ID, b.seed.Transition, b.resources...).Do()
+
+	for i := range b.params {
+		b.params[i].WorkspaceBuildID = resp.Build.ID
+	}
+	_ = dbgen.WorkspaceBuildParameters(b.t, b.db, b.params)
+
 	if b.ps != nil {
-		err = b.ps.Publish(codersdk.WorkspaceNotifyChannel(build.WorkspaceID), []byte{})
+		err = b.ps.Publish(codersdk.WorkspaceNotifyChannel(resp.Build.WorkspaceID), []byte{})
 		require.NoError(b.t, err)
 	}
-	return build
+
+	return resp
 }
 
 type ProvisionerJobResourcesBuilder struct {
@@ -229,7 +224,148 @@ func (b ProvisionerJobResourcesBuilder) Do() {
 	}
 	for _, resource := range b.resources {
 		//nolint:gocritic // This is only used by tests.
-		err := provisionerdserver.InsertWorkspaceResource(dbauthz.AsSystemRestricted(context.Background()), b.db, b.jobID, transition, resource, &telemetry.Snapshot{})
+		err := provisionerdserver.InsertWorkspaceResource(ownerCtx, b.db, b.jobID, transition, resource, &telemetry.Snapshot{})
 		require.NoError(b.t, err)
 	}
 }
+
+type TemplateVersionResponse struct {
+	Template        database.Template
+	TemplateVersion database.TemplateVersion
+}
+
+type TemplateVersionBuilder struct {
+	t         testing.TB
+	db        database.Store
+	seed      database.TemplateVersion
+	ps        pubsub.Pubsub
+	resources []*sdkproto.Resource
+	params    []database.TemplateVersionParameter
+	promote   bool
+}
+
+// TemplateVersion generates a template version and optionally a parent
+// template if no template ID is set on the seed.
+func TemplateVersion(t testing.TB, db database.Store) TemplateVersionBuilder {
+	return TemplateVersionBuilder{
+		t:       t,
+		db:      db,
+		promote: true,
+	}
+}
+
+func (t TemplateVersionBuilder) Seed(v database.TemplateVersion) TemplateVersionBuilder {
+	// nolint: revive // returns modified struct
+	t.seed = v
+	return t
+}
+
+func (t TemplateVersionBuilder) Pubsub(ps pubsub.Pubsub) TemplateVersionBuilder {
+	// nolint: revive // returns modified struct
+	t.ps = ps
+	return t
+}
+
+func (t TemplateVersionBuilder) Resources(rs ...*sdkproto.Resource) TemplateVersionBuilder {
+	// nolint: revive // returns modified struct
+	t.resources = rs
+	return t
+}
+
+func (t TemplateVersionBuilder) Params(ps ...database.TemplateVersionParameter) TemplateVersionBuilder {
+	// nolint: revive // returns modified struct
+	t.params = ps
+	return t
+}
+
+func (t TemplateVersionBuilder) Do() TemplateVersionResponse {
+	t.t.Helper()
+
+	t.seed.OrganizationID = takeFirst(t.seed.OrganizationID, uuid.New())
+	t.seed.ID = takeFirst(t.seed.ID, uuid.New())
+	t.seed.CreatedBy = takeFirst(t.seed.CreatedBy, uuid.New())
+
+	var resp TemplateVersionResponse
+	if t.seed.TemplateID.UUID == uuid.Nil {
+		resp.Template = dbgen.Template(t.t, t.db, database.Template{
+			ActiveVersionID: t.seed.ID,
+			OrganizationID:  t.seed.OrganizationID,
+			CreatedBy:       t.seed.CreatedBy,
+		})
+		t.seed.TemplateID = uuid.NullUUID{
+			Valid: true,
+			UUID:  resp.Template.ID,
+		}
+	}
+
+	version := dbgen.TemplateVersion(t.t, t.db, t.seed)
+
+	// Always make this version the active version. We can easily
+	// add a conditional to the builder to opt out of this when
+	// necessary.
+	err := t.db.UpdateTemplateActiveVersionByID(ownerCtx, database.UpdateTemplateActiveVersionByIDParams{
+		ID:              t.seed.TemplateID.UUID,
+		ActiveVersionID: t.seed.ID,
+		UpdatedAt:       dbtime.Now(),
+	})
+	require.NoError(t.t, err)
+
+	payload, err := json.Marshal(provisionerdserver.TemplateVersionImportJob{
+		TemplateVersionID: t.seed.ID,
+	})
+	require.NoError(t.t, err)
+
+	job := dbgen.ProvisionerJob(t.t, t.db, t.ps, database.ProvisionerJob{
+		ID:             version.JobID,
+		OrganizationID: t.seed.OrganizationID,
+		InitiatorID:    t.seed.CreatedBy,
+		Type:           database.ProvisionerJobTypeTemplateVersionImport,
+		Input:          payload,
+		CompletedAt: sql.NullTime{
+			Time:  dbtime.Now(),
+			Valid: true,
+		},
+	})
+
+	t.seed.JobID = job.ID
+
+	ProvisionerJobResources(t.t, t.db, job.ID, "", t.resources...).Do()
+
+	for i, param := range t.params {
+		param.TemplateVersionID = version.ID
+		t.params[i] = dbgen.TemplateVersionParameter(t.t, t.db, param)
+	}
+
+	resp.TemplateVersion = version
+	return resp
+}
+
+func must[V any](v V, err error) V {
+	if err != nil {
+		panic(err)
+	}
+	return v
+}
+
+// takeFirstF takes the first value that returns true
+func takeFirstF[Value any](values []Value, take func(v Value) bool) Value {
+	for _, v := range values {
+		if take(v) {
+			return v
+		}
+	}
+	// If all empty, return the last element
+	if len(values) > 0 {
+		return values[len(values)-1]
+	}
+	var empty Value
+	return empty
+}
+
+// takeFirst will take the first non-empty value.
+func takeFirst[Value comparable](values ...Value) Value {
+	var empty Value
+	return takeFirstF(values, func(v Value) bool {
+		return v != empty
+	})
+}

coderd/database/dbgen/dbgen.go

@@ -167,6 +167,8 @@ func WorkspaceAgent(t testing.TB, db database.Store, orig database.WorkspaceAgen
 }
 
 func Workspace(t testing.TB, db database.Store, orig database.Workspace) database.Workspace {
+	t.Helper()
+
 	workspace, err := db.InsertWorkspace(genCtx, database.InsertWorkspaceParams{
 		ID:                takeFirst(orig.ID, uuid.New()),
 		OwnerID:           takeFirst(orig.OwnerID, uuid.New()),
@@ -197,6 +199,8 @@ func WorkspaceAgentLogSource(t testing.TB, db database.Store, orig database.Work
 }
 
 func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuild) database.WorkspaceBuild {
+	t.Helper()
+
 	buildID := takeFirst(orig.ID, uuid.New())
 	var build database.WorkspaceBuild
 	err := db.InTx(func(db database.Store) error {
@@ -229,6 +233,38 @@ func WorkspaceBuild(t testing.TB, db database.Store, orig database.WorkspaceBuil
 	return build
 }
 
+func WorkspaceBuildParameters(t testing.TB, db database.Store, orig []database.WorkspaceBuildParameter) []database.WorkspaceBuildParameter {
+	if len(orig) == 0 {
+		return nil
+	}
+
+	var (
+		names  = make([]string, 0, len(orig))
+		values = make([]string, 0, len(orig))
+		params []database.WorkspaceBuildParameter
+	)
+	for _, param := range orig {
+		names = append(names, param.Name)
+		values = append(values, param.Value)
+	}
+	err := db.InTx(func(tx database.Store) error {
+		id := takeFirst(orig[0].WorkspaceBuildID, uuid.New())
+		err := tx.InsertWorkspaceBuildParameters(genCtx, database.InsertWorkspaceBuildParametersParams{
+			WorkspaceBuildID: id,
+			Name:             names,
+			Value:            values,
+		})
+		if err != nil {
+			return err
+		}
+
+		params, err = tx.GetWorkspaceBuildParameters(genCtx, id)
+		return err
+	}, nil)
+	require.NoError(t, err)
+	return params
+}
+
 func User(t testing.TB, db database.Store, orig database.User) database.User {
 	user, err := db.InsertUser(genCtx, database.InsertUserParams{
 		ID:             takeFirst(orig.ID, uuid.New()),
@@ -335,6 +371,8 @@ func GroupMember(t testing.TB, db database.Store, orig database.GroupMember) dat
 // ProvisionerJob is a bit more involved to get the values such as "completedAt", "startedAt", "cancelledAt" set.  ps
 // can be set to nil if you are SURE that you don't require a provisionerdaemon to acquire the job in your test.
 func ProvisionerJob(t testing.TB, db database.Store, ps pubsub.Pubsub, orig database.ProvisionerJob) database.ProvisionerJob {
+	t.Helper()
+
 	jobID := takeFirst(orig.ID, uuid.New())
 	// Always set some tags to prevent Acquire from grabbing jobs it should not.
 	if !orig.StartedAt.Time.IsZero() {
@@ -545,7 +583,7 @@ func TemplateVersion(t testing.TB, db database.Store, orig database.TemplateVers
 		versionID := takeFirst(orig.ID, uuid.New())
 		err := db.InsertTemplateVersion(genCtx, database.InsertTemplateVersionParams{
 			ID:             versionID,
-			TemplateID:     orig.TemplateID,
+			TemplateID:     takeFirst(orig.TemplateID, uuid.NullUUID{}),
 			OrganizationID: takeFirst(orig.OrganizationID, uuid.New()),
 			CreatedAt:      takeFirst(orig.CreatedAt, dbtime.Now()),
 			UpdatedAt:      takeFirst(orig.UpdatedAt, dbtime.Now()),
@@ -585,6 +623,32 @@ func TemplateVersionVariable(t testing.TB, db database.Store, orig database.Temp
 	return version
 }
 
+func TemplateVersionParameter(t testing.TB, db database.Store, orig database.TemplateVersionParameter) database.TemplateVersionParameter {
+	t.Helper()
+
+	version, err := db.InsertTemplateVersionParameter(genCtx, database.InsertTemplateVersionParameterParams{
+		TemplateVersionID:   takeFirst(orig.TemplateVersionID, uuid.New()),
+		Name:                takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
+		Description:         takeFirst(orig.Description, namesgenerator.GetRandomName(1)),
+		Type:                takeFirst(orig.Type, "string"),
+		Mutable:             takeFirst(orig.Mutable, false),
+		DefaultValue:        takeFirst(orig.DefaultValue, namesgenerator.GetRandomName(1)),
+		Icon:                takeFirst(orig.Icon, namesgenerator.GetRandomName(1)),
+		Options:             takeFirstSlice(orig.Options, []byte("[]")),
+		ValidationRegex:     takeFirst(orig.ValidationRegex, ""),
+		ValidationMin:       takeFirst(orig.ValidationMin, sql.NullInt32{}),
+		ValidationMax:       takeFirst(orig.ValidationMax, sql.NullInt32{}),
+		ValidationError:     takeFirst(orig.ValidationError, ""),
+		ValidationMonotonic: takeFirst(orig.ValidationMonotonic, ""),
+		Required:            takeFirst(orig.Required, false),
+		DisplayName:         takeFirst(orig.DisplayName, namesgenerator.GetRandomName(1)),
+		DisplayOrder:        takeFirst(orig.DisplayOrder, 0),
+		Ephemeral:           takeFirst(orig.Ephemeral, false),
+	})
+	require.NoError(t, err, "insert template version parameter")
+	return version
+}
+
 func WorkspaceAgentStat(t testing.TB, db database.Store, orig database.WorkspaceAgentStat) database.WorkspaceAgentStat {
 	if orig.ConnectionsByProto == nil {
 		orig.ConnectionsByProto = json.RawMessage([]byte("{}"))

coderd/database/dbgen/dbgen_test.go

@@ -173,6 +173,22 @@ func TestGenerator(t *testing.T) {
 		exp := dbgen.GitSSHKey(t, db, database.GitSSHKey{})
 		require.Equal(t, exp, must(db.GetGitSSHKey(context.Background(), exp.UserID)))
 	})
+
+	t.Run("WorkspaceBuildParameters", func(t *testing.T) {
+		t.Parallel()
+		db := dbmem.New()
+		exp := dbgen.WorkspaceBuildParameters(t, db, []database.WorkspaceBuildParameter{{}, {}, {}})
+		require.Equal(t, exp, must(db.GetWorkspaceBuildParameters(context.Background(), exp[0].WorkspaceBuildID)))
+	})
+
+	t.Run("TemplateVersionParameter", func(t *testing.T) {
+		t.Parallel()
+		db := dbmem.New()
+		exp := dbgen.TemplateVersionParameter(t, db, database.TemplateVersionParameter{})
+		actual := must(db.GetTemplateVersionParameters(context.Background(), exp.TemplateVersionID))
+		require.Len(t, actual, 1)
+		require.Equal(t, exp, actual[0])
+	})
 }
 
 func must[T any](value T, err error) T {