diff --git a/periodic_job_test.go b/periodic_job_test.go index 5dc98ec1..9f55eca2 100644 --- a/periodic_job_test.go +++ b/periodic_job_test.go @@ -24,7 +24,7 @@ func TestNeverSchedule(t *testing.T) { require.False(t, next.Before(now)) // use an arbitrary duration to check that // the next schedule is far in the future - require.True(t, next.Year()-now.Year() > 1000) + require.Greater(t, next.Year()-now.Year(), 1000) }) } diff --git a/rivertest/worker.go b/rivertest/worker.go index 55a623c9..947904de 100644 --- a/rivertest/worker.go +++ b/rivertest/worker.go @@ -3,6 +3,7 @@ package rivertest import ( "context" "encoding/json" + "errors" "sync/atomic" "testing" "time" @@ -77,12 +78,62 @@ func (w *Worker[T, TTx]) Work(ctx context.Context, tb testing.TB, args T, opts * return w.WorkJob(ctx, tb, job) } +// WorkTx allocates a synthetic job with the provided arguments and optional +// insert options, then works it in the given transaction. +func (w *Worker[T, TTx]) WorkTx(ctx context.Context, tb testing.TB, tx TTx, args T, opts *river.InsertOpts) error { + tb.Helper() + job := makeJobFromArgs(tb, w.config, args, opts) + return w.WorkJobTx(ctx, tb, tx, job) +} + // WorkJob works the provided job. The job must be constructed to be a realistic // job using external logic prior to calling this method. Unlike the other // variants, this method offers total control over the job's attributes. func (w *Worker[T, TTx]) WorkJob(ctx context.Context, tb testing.TB, job *river.Job[T]) error { tb.Helper() + var exec riverdriver.Executor + if w.client.Driver().HasPool() { + exec = w.client.Driver().GetExecutor() + } + + return w.workJobExec(ctx, tb, exec, job) +} + +// WorkJobTx works the provided job in the given transaction. The job must be +// constructed to be a realistic job using external logic prior to calling this +// method. Unlike the other variants, this method offers total control over the +// job's attributes. +func (w *Worker[T, TTx]) WorkJobTx(ctx context.Context, tb testing.TB, tx TTx, job *river.Job[T]) error { + tb.Helper() + return w.workJobExec(ctx, tb, w.client.Driver().UnwrapExecutor(tx), job) +} + +func (w *Worker[T, TTx]) workJobExec(ctx context.Context, tb testing.TB, exec riverdriver.Executor, job *river.Job[T]) error { + tb.Helper() + + // Try to transition an existing job row to running, but also tolerate the + // row not existing in the database. Most of the time you don't need to + // insert a real job row to use this function (it's only necessary if you + // want to use `JobCompleteTx`). + if exec != nil { + updatedJobRow, err := exec.JobUpdate(ctx, &riverdriver.JobUpdateParams{ + ID: job.ID, + StateDoUpdate: true, + State: rivertype.JobStateRunning, + }) + if err != nil && !errors.Is(err, rivertype.ErrNotFound) { + return err + } + if updatedJobRow != nil { + job.JobRow = updatedJobRow + } + } + + // Regardless of whether the job was actually in the database, transition it + // to running. + job.JobRow.State = rivertype.JobStateRunning + // populate river client into context: ctx = WorkContext(ctx, w.client) ctx = context.WithValue(ctx, execution.ContextKeyInsideTestWorker{}, true) diff --git a/rivertest/worker_test.go b/rivertest/worker_test.go index 66a9f7f3..9114be04 100644 --- a/rivertest/worker_test.go +++ b/rivertest/worker_test.go @@ -5,11 +5,14 @@ import ( "testing" "time" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/require" "github.com/riverqueue/river" "github.com/riverqueue/river/internal/dbunique" "github.com/riverqueue/river/internal/execution" + "github.com/riverqueue/river/internal/riverinternaltest" "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivershared/baseservice" "github.com/riverqueue/river/rivershared/riversharedtest" @@ -118,7 +121,7 @@ func TestWorker_Work(t *testing.T) { }) tw := NewWorker(t, driver, config, worker) - // You can also use the WorkOpts method to pass in custom insert options: + // You can also pass in custom insert options: require.NoError(t, tw.Work(context.Background(), t, testArgs{Value: "test3"}, &river.InsertOpts{ MaxAttempts: 420, Metadata: []byte(`{"key": "value"}`), @@ -163,24 +166,224 @@ func TestWorker_Work(t *testing.T) { }) } +func TestWorker_WorkTx(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type testBundle struct { + client *river.Client[pgx.Tx] + driver *riverpgxv5.Driver + tx pgx.Tx + workFunc func(ctx context.Context, job *river.Job[testArgs]) error + } + + setup := func(t *testing.T) (*Worker[testArgs, pgx.Tx], *testBundle) { + t.Helper() + + var ( + config = &river.Config{} + driver = riverpgxv5.New(nil) + ) + + client, err := river.NewClient(driver, config) + require.NoError(t, err) + + bundle := &testBundle{ + client: client, + driver: driver, + tx: riverinternaltest.TestTx(ctx, t), + workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, + } + + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { + return bundle.workFunc(ctx, job) + }) + + return NewWorker(t, driver, config, worker), bundle + } + + t.Run("Success", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + require.Equal(t, rivertype.JobStateRunning, job.State) + return nil + } + + require.NoError(t, testWorker.WorkTx(ctx, t, bundle.tx, testArgs{}, nil)) + }) +} + func TestWorker_WorkJob(t *testing.T) { t.Parallel() - config := &river.Config{} - driver := riverpgxv5.New(nil) + ctx := context.Background() + + type testBundle struct { + client *river.Client[pgx.Tx] + dbPool *pgxpool.Pool + driver *riverpgxv5.Driver + workFunc func(ctx context.Context, job *river.Job[testArgs]) error + } + + setup := func(t *testing.T) (*Worker[testArgs, pgx.Tx], *testBundle) { + t.Helper() + + var ( + config = &river.Config{} + dbPool = riverinternaltest.TestDB(ctx, t) + driver = riverpgxv5.New(dbPool) + ) + + client, err := river.NewClient(driver, config) + require.NoError(t, err) + + bundle := &testBundle{ + client: client, + driver: driver, + dbPool: dbPool, + workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, + } + + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { + return bundle.workFunc(ctx, job) + }) + + return NewWorker(t, driver, config, worker), bundle + } + + t.Run("Success", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + require.Equal(t, []string{"worker123"}, job.JobRow.AttemptedBy) + require.Equal(t, rivertype.JobStateRunning, job.State) + return nil + } + + now := time.Now() + require.NoError(t, testWorker.WorkJob(ctx, t, makeJobFromFactoryBuild(t, testArgs{}, &testfactory.JobOpts{ + AttemptedAt: &now, + AttemptedBy: []string{"worker123"}, + CreatedAt: &now, + EncodedArgs: []byte(`{"value": "test"}`), + Errors: nil, + }))) + }) + + t.Run("JobCompleteTxWithInsertedJobRow", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) - worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { - require.Equal(t, []string{"worker123"}, job.JobRow.AttemptedBy) - return nil + args := testArgs{} + insertRes, err := bundle.client.Insert(ctx, args, nil) + require.NoError(t, err) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + tx, err := bundle.dbPool.Begin(ctx) + require.NoError(t, err) + + updatedJob, err := bundle.driver.GetExecutor().JobGetByID(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, updatedJob.State) + + _, err = river.JobCompleteTx[*riverpgxv5.Driver](ctx, tx, job) + require.NoError(t, err) + + err = tx.Commit(ctx) + require.NoError(t, err) + + return nil + } + + require.NoError(t, testWorker.WorkJob(ctx, t, &river.Job[testArgs]{Args: args, JobRow: insertRes.Job})) + + updatedJob, err := bundle.driver.GetExecutor().JobGetByID(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateCompleted, updatedJob.State) + }) +} + +func TestWorker_WorkJobTx(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type testBundle struct { + client *river.Client[pgx.Tx] + driver *riverpgxv5.Driver + tx pgx.Tx + workFunc func(ctx context.Context, job *river.Job[testArgs]) error + } + + setup := func(t *testing.T) (*Worker[testArgs, pgx.Tx], *testBundle) { + t.Helper() + + var ( + config = &river.Config{} + driver = riverpgxv5.New(nil) + ) + + client, err := river.NewClient(driver, config) + require.NoError(t, err) + + bundle := &testBundle{ + client: client, + driver: driver, + tx: riverinternaltest.TestTx(ctx, t), + workFunc: func(ctx context.Context, job *river.Job[testArgs]) error { return nil }, + } + + worker := river.WorkFunc(func(ctx context.Context, job *river.Job[testArgs]) error { + return bundle.workFunc(ctx, job) + }) + + return NewWorker(t, driver, config, worker), bundle + } + + t.Run("Success", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + require.Equal(t, rivertype.JobStateRunning, job.State) + return nil + } + + require.NoError(t, testWorker.WorkJobTx(ctx, t, bundle.tx, makeJobFromFactoryBuild(t, testArgs{}, &testfactory.JobOpts{}))) + }) + + t.Run("JobCompleteTxWithInsertedJobRow", func(t *testing.T) { + t.Parallel() + + testWorker, bundle := setup(t) + + args := testArgs{} + insertRes, err := bundle.client.InsertTx(ctx, bundle.tx, args, nil) + require.NoError(t, err) + + bundle.workFunc = func(ctx context.Context, job *river.Job[testArgs]) error { + updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, updatedJob.State) + + _, err = river.JobCompleteTx[*riverpgxv5.Driver](ctx, bundle.tx, job) + require.NoError(t, err) + + return nil + } + + require.NoError(t, testWorker.WorkJobTx(ctx, t, bundle.tx, &river.Job[testArgs]{Args: args, JobRow: insertRes.Job})) + + updatedJob, err := bundle.driver.UnwrapExecutor(bundle.tx).JobGetByID(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateCompleted, updatedJob.State) }) - tw := NewWorker(t, driver, config, worker) - - now := time.Now() - require.NoError(t, tw.WorkJob(context.Background(), t, makeJobFromFactoryBuild(t, testArgs{Value: "test"}, &testfactory.JobOpts{ - AttemptedAt: &now, - AttemptedBy: []string{"worker123"}, - CreatedAt: &now, - EncodedArgs: []byte(`{"value": "test"}`), - Errors: nil, - }))) }