diff --git a/rivertest/work_unit_wrapper.go b/rivertest/work_unit_wrapper.go new file mode 100644 index 00000000..1dd3f293 --- /dev/null +++ b/rivertest/work_unit_wrapper.go @@ -0,0 +1,66 @@ +package rivertest + +import ( + "context" + "encoding/json" + "time" + + "github.com/riverqueue/river" + "github.com/riverqueue/river/internal/hooklookup" + "github.com/riverqueue/river/internal/workunit" + "github.com/riverqueue/river/rivertype" +) + +// Everything in this file is duplicated from `./work_unit_wrapper.go`. I looked +// into deduplicating it, but unfortuntately it's quite difficult/impossible. +// +// We can't put it in `internal/` because it depends on top level `river` types +// so that'd immediately produce a cyclic dependency. The only approach that +// would work would be to export it from the top-level `river` which `rivertest` +// could then use, but then we'd be exporting some low-level plumbing for +// internal use, which isn't worth it. +// +// The original sin is that basic types like `Job`, `JobArgs`, `Worker` +// should've been in `rivertype` instead of the top-level `river`, then anything +// (including these types) could easily have been refactored to anywhere, but +// it's too late to do that now. +// +// The best thing to do for now is probably just to live with the duplication, +// so I've removed a todo that used to be here and replaced it with this +// explanatory comment. + +// workUnitFactoryWrapper wraps a Worker to implement workUnitFactory. +type workUnitFactoryWrapper[T river.JobArgs] struct { + worker river.Worker[T] +} + +func (w *workUnitFactoryWrapper[T]) MakeUnit(jobRow *rivertype.JobRow) workunit.WorkUnit { + return &wrapperWorkUnit[T]{jobRow: jobRow, worker: w.worker} +} + +// wrapperWorkUnit implements workUnit for a job and Worker. +type wrapperWorkUnit[T river.JobArgs] struct { + job *river.Job[T] // not set until after UnmarshalJob is invoked + jobRow *rivertype.JobRow + worker river.Worker[T] +} + +func (w *wrapperWorkUnit[T]) HookLookup(lookup *hooklookup.JobHookLookup) hooklookup.HookLookupInterface { + var job T + return lookup.ByJobArgs(job) +} + +func (w *wrapperWorkUnit[T]) Middleware() []rivertype.WorkerMiddleware { + return w.worker.Middleware(w.jobRow) +} +func (w *wrapperWorkUnit[T]) NextRetry() time.Time { return w.worker.NextRetry(w.job) } +func (w *wrapperWorkUnit[T]) Timeout() time.Duration { return w.worker.Timeout(w.job) } +func (w *wrapperWorkUnit[T]) Work(ctx context.Context) error { return w.worker.Work(ctx, w.job) } + +func (w *wrapperWorkUnit[T]) UnmarshalJob() error { + w.job = &river.Job[T]{ + JobRow: w.jobRow, + } + + return json.Unmarshal(w.jobRow.EncodedArgs, &w.job.Args) +} diff --git a/rivertest/worker.go b/rivertest/worker.go index c0b9a652..7b8e48fd 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -2,11 +2,9 @@ package rivertest import ( "context" - "encoding/json" "errors" "fmt" "testing" - "time" "github.com/riverqueue/river" "github.com/riverqueue/river/internal/execution" @@ -15,7 +13,6 @@ import ( "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/maintenance" "github.com/riverqueue/river/internal/middlewarelookup" - "github.com/riverqueue/river/internal/workunit" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/rivershared/baseservice" "github.com/riverqueue/river/rivershared/riversharedtest" @@ -294,41 +291,3 @@ func (h *errorHandlerWrapper) HandleError(ctx context.Context, job *rivertype.Jo func (h *errorHandlerWrapper) HandlePanic(ctx context.Context, job *rivertype.JobRow, panicVal any, trace string) *jobexecutor.ErrorHandlerResult { return h.HandlePanicFunc(ctx, job, panicVal, trace) } - -// TODO: move work_unit_wrapper.go so I don't need to copy paste it here: - -// workUnitFactoryWrapper wraps a Worker to implement workUnitFactory. -type workUnitFactoryWrapper[T river.JobArgs] struct { - worker river.Worker[T] -} - -func (w *workUnitFactoryWrapper[T]) MakeUnit(jobRow *rivertype.JobRow) workunit.WorkUnit { - return &wrapperWorkUnit[T]{jobRow: jobRow, worker: w.worker} -} - -// wrapperWorkUnit implements workUnit for a job and Worker. -type wrapperWorkUnit[T river.JobArgs] struct { - job *river.Job[T] // not set until after UnmarshalJob is invoked - jobRow *rivertype.JobRow - worker river.Worker[T] -} - -func (w *wrapperWorkUnit[T]) HookLookup(lookup *hooklookup.JobHookLookup) hooklookup.HookLookupInterface { - var job T - return lookup.ByJobArgs(job) -} - -func (w *wrapperWorkUnit[T]) Middleware() []rivertype.WorkerMiddleware { - return w.worker.Middleware(w.jobRow) -} -func (w *wrapperWorkUnit[T]) NextRetry() time.Time { return w.worker.NextRetry(w.job) } -func (w *wrapperWorkUnit[T]) Timeout() time.Duration { return w.worker.Timeout(w.job) } -func (w *wrapperWorkUnit[T]) Work(ctx context.Context) error { return w.worker.Work(ctx, w.job) } - -func (w *wrapperWorkUnit[T]) UnmarshalJob() error { - w.job = &river.Job[T]{ - JobRow: w.jobRow, - } - - return json.Unmarshal(w.jobRow.EncodedArgs, &w.job.Args) -}