Skip to content

Commit 25a0c80

Browse files
johnstcnCopilot
andauthored
chore(coderd/database/dbfake): add support for provisioner job timestamp control (#21944)
Relates to #21922 / coder/internal#1259 * Adds `dbfake.BuilderOption func(*WorkspaceBuildBuilder)` * Adds `BuilderOption` methods for setting various provisioner job related fields on `WorkspaceBuildBuilder`. * Migrates a number of existing tests that previously dependeded on provisioner job timing to use these updated methods in the following packages: * `coderd/jobreaper` * `coderd/notifications/reports` * `enterprise/coderd/schedule` * `enterprise/coderd/prebuilds` * `scripts/workspace-runtime-audit` 🤖 Created using Mux (Opus 4.5) --------- Co-authored-by: Copilot <[email protected]>
1 parent fabb0b8 commit 25a0c80

6 files changed

Lines changed: 417 additions & 523 deletions

File tree

‎coderd/database/dbfake/dbfake.go‎

Lines changed: 140 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,61 @@ type WorkspaceBuildBuilder struct {
5858
jobStatus database.ProvisionerJobStatus
5959
taskAppID uuid.UUID
6060
taskSeed database.TaskTable
61+
62+
// Individual timestamp fields for job customization.
63+
jobCreatedAt time.Time
64+
jobStartedAt time.Time
65+
jobUpdatedAt time.Time
66+
jobCompletedAt time.Time
67+
68+
jobError string // Error message for failed jobs
69+
jobErrorCode string // Error code for failed jobs
70+
}
71+
72+
// BuilderOption is a functional option for customizing job timestamps
73+
// on status methods.
74+
type BuilderOption func(*WorkspaceBuildBuilder)
75+
76+
// WithJobCreatedAt sets the CreatedAt timestamp for the provisioner job.
77+
func WithJobCreatedAt(t time.Time) BuilderOption {
78+
return func(b *WorkspaceBuildBuilder) {
79+
b.jobCreatedAt = t
80+
}
81+
}
82+
83+
// WithJobStartedAt sets the StartedAt timestamp for the provisioner job.
84+
func WithJobStartedAt(t time.Time) BuilderOption {
85+
return func(b *WorkspaceBuildBuilder) {
86+
b.jobStartedAt = t
87+
}
88+
}
89+
90+
// WithJobUpdatedAt sets the UpdatedAt timestamp for the provisioner job.
91+
func WithJobUpdatedAt(t time.Time) BuilderOption {
92+
return func(b *WorkspaceBuildBuilder) {
93+
b.jobUpdatedAt = t
94+
}
95+
}
96+
97+
// WithJobCompletedAt sets the CompletedAt timestamp for the provisioner job.
98+
func WithJobCompletedAt(t time.Time) BuilderOption {
99+
return func(b *WorkspaceBuildBuilder) {
100+
b.jobCompletedAt = t
101+
}
102+
}
103+
104+
// WithJobError sets the error message for the provisioner job.
105+
func WithJobError(msg string) BuilderOption {
106+
return func(b *WorkspaceBuildBuilder) {
107+
b.jobError = msg
108+
}
109+
}
110+
111+
// WithJobErrorCode sets the error code for the provisioner job.
112+
func WithJobErrorCode(code string) BuilderOption {
113+
return func(b *WorkspaceBuildBuilder) {
114+
b.jobErrorCode = code
115+
}
61116
}
62117

63118
// WorkspaceBuild generates a workspace build for the provided workspace.
@@ -141,18 +196,59 @@ func (b WorkspaceBuildBuilder) WithTask(taskSeed database.TaskTable, appSeed *sd
141196
})
142197
}
143198

144-
func (b WorkspaceBuildBuilder) Starting() WorkspaceBuildBuilder {
199+
// Starting sets the job to running status.
200+
func (b WorkspaceBuildBuilder) Starting(opts ...BuilderOption) WorkspaceBuildBuilder {
201+
//nolint: revive // returns modified struct
145202
b.jobStatus = database.ProvisionerJobStatusRunning
203+
for _, opt := range opts {
204+
opt(&b)
205+
}
146206
return b
147207
}
148208

149-
func (b WorkspaceBuildBuilder) Pending() WorkspaceBuildBuilder {
209+
// Pending sets the job to pending status.
210+
func (b WorkspaceBuildBuilder) Pending(opts ...BuilderOption) WorkspaceBuildBuilder {
211+
//nolint: revive // returns modified struct
150212
b.jobStatus = database.ProvisionerJobStatusPending
213+
for _, opt := range opts {
214+
opt(&b)
215+
}
151216
return b
152217
}
153218

154-
func (b WorkspaceBuildBuilder) Canceled() WorkspaceBuildBuilder {
219+
// Canceled sets the job to canceled status.
220+
func (b WorkspaceBuildBuilder) Canceled(opts ...BuilderOption) WorkspaceBuildBuilder {
221+
//nolint: revive // returns modified struct
155222
b.jobStatus = database.ProvisionerJobStatusCanceled
223+
for _, opt := range opts {
224+
opt(&b)
225+
}
226+
return b
227+
}
228+
229+
// Succeeded sets the job to succeeded status.
230+
// This is the default status.
231+
func (b WorkspaceBuildBuilder) Succeeded(opts ...BuilderOption) WorkspaceBuildBuilder {
232+
//nolint: revive // returns modified struct
233+
b.jobStatus = database.ProvisionerJobStatusSucceeded
234+
for _, opt := range opts {
235+
opt(&b)
236+
}
237+
return b
238+
}
239+
240+
// Failed sets the provisioner job to a failed state. Use WithJobError and
241+
// WithJobErrorCode options to set the error message and code. If no error
242+
// message is provided, "failed" is used as the default.
243+
func (b WorkspaceBuildBuilder) Failed(opts ...BuilderOption) WorkspaceBuildBuilder {
244+
//nolint: revive // returns modified struct
245+
b.jobStatus = database.ProvisionerJobStatusFailed
246+
for _, opt := range opts {
247+
opt(&b)
248+
}
249+
if b.jobError == "" {
250+
b.jobError = "failed"
251+
}
156252
return b
157253
}
158254

@@ -267,8 +363,8 @@ func (b WorkspaceBuildBuilder) doInTX() WorkspaceResponse {
267363

268364
job, err := b.db.InsertProvisionerJob(ownerCtx, database.InsertProvisionerJobParams{
269365
ID: jobID,
270-
CreatedAt: dbtime.Now(),
271-
UpdatedAt: dbtime.Now(),
366+
CreatedAt: takeFirstTime(b.jobCreatedAt, b.ws.CreatedAt, dbtime.Now()),
367+
UpdatedAt: takeFirstTime(b.jobCreatedAt, b.ws.CreatedAt, dbtime.Now()),
272368
OrganizationID: b.ws.OrganizationID,
273369
InitiatorID: b.ws.OwnerID,
274370
Provisioner: database.ProvisionerTypeEcho,
@@ -291,11 +387,12 @@ func (b WorkspaceBuildBuilder) doInTX() WorkspaceResponse {
291387
// might need to do this multiple times if we got a template version
292388
// import job as well
293389
b.logger.Debug(context.Background(), "looping to acquire provisioner job")
390+
startedAt := takeFirstTime(b.jobStartedAt, dbtime.Now())
294391
for {
295392
j, err := b.db.AcquireProvisionerJob(ownerCtx, database.AcquireProvisionerJobParams{
296393
OrganizationID: job.OrganizationID,
297394
StartedAt: sql.NullTime{
298-
Time: dbtime.Now(),
395+
Time: startedAt,
299396
Valid: true,
300397
},
301398
WorkerID: uuid.NullUUID{
@@ -311,32 +408,54 @@ func (b WorkspaceBuildBuilder) doInTX() WorkspaceResponse {
311408
break
312409
}
313410
}
411+
if !b.jobUpdatedAt.IsZero() {
412+
err = b.db.UpdateProvisionerJobByID(ownerCtx, database.UpdateProvisionerJobByIDParams{
413+
ID: job.ID,
414+
UpdatedAt: b.jobUpdatedAt,
415+
})
416+
require.NoError(b.t, err, "update job updated_at")
417+
}
314418
case database.ProvisionerJobStatusCanceled:
315419
// Set provisioner job status to 'canceled'
316420
b.logger.Debug(context.Background(), "canceling the provisioner job")
317-
now := dbtime.Now()
421+
completedAt := takeFirstTime(b.jobCompletedAt, dbtime.Now())
318422
err = b.db.UpdateProvisionerJobWithCancelByID(ownerCtx, database.UpdateProvisionerJobWithCancelByIDParams{
319423
ID: jobID,
320424
CanceledAt: sql.NullTime{
321-
Time: now,
425+
Time: completedAt,
322426
Valid: true,
323427
},
324428
CompletedAt: sql.NullTime{
325-
Time: now,
429+
Time: completedAt,
326430
Valid: true,
327431
},
328432
})
329433
require.NoError(b.t, err, "cancel job")
434+
case database.ProvisionerJobStatusFailed:
435+
b.logger.Debug(context.Background(), "failing the provisioner job")
436+
completedAt := takeFirstTime(b.jobCompletedAt, dbtime.Now())
437+
err = b.db.UpdateProvisionerJobWithCompleteByID(ownerCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
438+
ID: job.ID,
439+
UpdatedAt: completedAt,
440+
Error: sql.NullString{String: b.jobError, Valid: b.jobError != ""},
441+
ErrorCode: sql.NullString{String: b.jobErrorCode, Valid: b.jobErrorCode != ""},
442+
CompletedAt: sql.NullTime{
443+
Time: completedAt,
444+
Valid: true,
445+
},
446+
})
447+
require.NoError(b.t, err, "fail job")
330448
default:
331449
// By default, consider jobs in 'succeeded' status
332450
b.logger.Debug(context.Background(), "completing the provisioner job")
451+
completedAt := takeFirstTime(b.jobCompletedAt, dbtime.Now())
333452
err = b.db.UpdateProvisionerJobWithCompleteByID(ownerCtx, database.UpdateProvisionerJobWithCompleteByIDParams{
334453
ID: job.ID,
335-
UpdatedAt: dbtime.Now(),
454+
UpdatedAt: completedAt,
336455
Error: sql.NullString{},
337456
ErrorCode: sql.NullString{},
338457
CompletedAt: sql.NullTime{
339-
Time: dbtime.Now(),
458+
Time: completedAt,
340459
Valid: true,
341460
},
342461
})
@@ -751,6 +870,16 @@ func takeFirst[Value comparable](values ...Value) Value {
751870
})
752871
}
753872

873+
// takeFirstTime returns the first non-zero time.Time.
874+
func takeFirstTime(values ...time.Time) time.Time {
875+
for _, v := range values {
876+
if !v.IsZero() {
877+
return v
878+
}
879+
}
880+
return time.Time{}
881+
}
882+
754883
// mustWorkspaceAppByWorkspaceAndBuildAndAppID finds a workspace app by
755884
// workspace ID, build number, and app ID. It returns the workspace app
756885
// if found, otherwise fails the test.

0 commit comments

Comments
 (0)