@@ -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