Let test workers call JobCompleteTx + test work transaction variants

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)
 	})
 }
 

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)

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,
-	})))
 }