From 73e561e68421bf44e45a95821885877d0bbdf9fe Mon Sep 17 00:00:00 2001 From: Brandur Date: Sat, 15 Feb 2025 14:46:53 -0800 Subject: [PATCH] Let test workers call `JobCompleteTx` + test work transaction variants Follows up #753 to make it possible to call `JobCompleteTx` inside test worker implementations. The test worker tries to update an equivalent job in the database's state to `running` so `JobCompleteTx` will be able to find it when called (it only considers jobs in `running` state). We also augment the API so that it's workable for users who want to use test transactions in their tests (which is hopefully the vast majority of users) by adding transaction variants `WorkTx` and `WorkJobTx`. --- periodic_job_test.go | 2 +- rivertest/worker.go | 51 +++++++++ rivertest/worker_test.go | 235 ++++++++++++++++++++++++++++++++++++--- 3 files changed, 271 insertions(+), 17 deletions(-) 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, - }))) }