Tests: Implement in-memory worker pattern via simple WaitGroup

Author: VojtechVitek

table.go

@@ -180,11 +180,43 @@ func (t *Table[T, PT, IDT]) WithTx(tx pgx.Tx) *Table[T, PT, IDT] {
 	}
 }
 
-// LockForUpdates locks and processes records using PostgreSQL's FOR UPDATE SKIP LOCKED pattern
-// for safe concurrent processing. Each record is processed exactly once across multiple workers.
-// Records are automatically updated after updateFn() completes. Keep updateFn() fast to avoid
-// holding the transaction. For long-running work, update status to "processing" and return early,
-// then process asynchronously. Use defer LockOneForUpdate() to update status to "completed" or "failed".
+// LockForUpdate locks and updates one record using PostgreSQL's FOR UPDATE SKIP LOCKED pattern
+// within a database transaction for safe concurrent processing. The record is processed exactly
+// once across multiple workers. The record is automatically updated after updateFn() completes.
+//
+// Keep updateFn() fast to avoid holding the transaction. For long-running work, update status
+// to "processing" and return early, then process asynchronously. Use defer LockForUpdate()
+// to update status to "completed" or "failed".
+//
+// Returns ErrNoRows if no matching records are available for locking.
+func (t *Table[T, PT, IDT]) LockForUpdate(ctx context.Context, where sq.Sqlizer, orderBy []string, updateFn func(record PT)) error {
+	var noRows bool
+
+	err := t.LockForUpdates(ctx, where, orderBy, 1, func(records []PT) {
+		if len(records) > 0 {
+			updateFn(records[0])
+		} else {
+			noRows = true
+		}
+	})
+	if err != nil {
+		return fmt.Errorf("lock for update one: %w", err)
+	}
+
+	if noRows {
+		return ErrNoRows
+	}
+
+	return nil
+}
+
+// LockForUpdates locks and updates records using PostgreSQL's FOR UPDATE SKIP LOCKED pattern
+// within a database transaction for safe concurrent processing. Each record is processed exactly
+// once across multiple workers. Records are automatically updated after updateFn() completes.
+//
+// Keep updateFn() fast to avoid holding the transaction. For long-running work, update status
+// to "processing" and return early, then process asynchronously. Use defer LockForUpdate()
+// to update status to "completed" or "failed".
 func (t *Table[T, PT, IDT]) LockForUpdates(ctx context.Context, where sq.Sqlizer, orderBy []string, limit uint64, updateFn func(records []PT)) error {
 	// Check if we're already in a transaction
 	if t.DB.Query.Tx != nil {
@@ -227,31 +259,3 @@ func (t *Table[T, PT, IDT]) lockForUpdatesWithTx(ctx context.Context, pgTx pgx.T
 
 	return nil
 }
-
-// LockForUpdate locks and processes one record using PostgreSQL's FOR UPDATE SKIP LOCKED pattern
-// for safe concurrent processing. The record is processed exactly once across multiple workers.
-// The record is automatically updated after updateFn() completes. Keep updateFn() fast to avoid
-// holding the transaction. For long-running work, update status to "processing" and return early,
-// then process asynchronously. Use defer LockForUpdate() to update status to "completed" or "failed".
-//
-// Returns ErrNoRows if no matching records are available for locking.
-func (t *Table[T, PT, IDT]) LockForUpdate(ctx context.Context, where sq.Sqlizer, orderBy []string, updateFn func(record PT)) error {
-	var noRows bool
-
-	err := t.LockForUpdates(ctx, where, orderBy, 1, func(records []PT) {
-		if len(records) > 0 {
-			updateFn(records[0])
-		} else {
-			noRows = true
-		}
-	})
-	if err != nil {
-		return fmt.Errorf("lock for update one: %w", err)
-	}
-
-	if noRows {
-		return ErrNoRows
-	}
-
-	return nil
-}

tests/database_test.go

@@ -41,6 +41,8 @@ func (db *Database) WithTx(tx pgx.Tx) *Database {
 	return initDB(pgkitDB)
 }
 
+func (db *Database) Close() { db.DB.Conn.Close() }
+
 type accountsTable struct {
 	*pgkit.Table[Account, *Account, int64]
 }

tests/table_test.go

@@ -1,10 +1,7 @@
 package pgkit_test
 
 import (
-	"context"
 	"fmt"
-	"log"
-	"math/rand"
 	"slices"
 	"sync"
 	"testing"
@@ -160,6 +157,7 @@ func TestLockForUpdates(t *testing.T) {
 
 	ctx := t.Context()
 	db := initDB(DB)
+	worker := &Worker{DB: db}
 
 	t.Run("TestLockForUpdates", func(t *testing.T) {
 		// Create account.
@@ -216,7 +214,7 @@ func TestLockForUpdates(t *testing.T) {
 				require.NoError(t, err, "lock for update")
 
 				for _, review := range processReviews {
-					go processReviewAsynchronously(ctx, db, review)
+					go worker.ProcessReview(ctx, review)
 				}
 
 				for i, review := range processReviews {
@@ -233,50 +231,11 @@ func TestLockForUpdates(t *testing.T) {
 		require.Equal(t, 100, len(uniqueIDs), "number of unique reviews picked up for processing should be 100")
 
 		// Wait for all reviews to be processed asynchronously.
-		time.Sleep(2 * time.Second)
+		worker.Wait()
 
 		// Double check there's no reviews stuck in "processing" status.
 		count, err := db.Reviews.Count(ctx, sq.Eq{"status": ReviewStatusProcessing})
 		require.NoError(t, err, "count reviews")
 		require.Zero(t, count, "there should be no reviews stuck in 'processing' status")
 	})
 }
-
-func processReviewAsynchronously(ctx context.Context, db *Database, review *Review) (err error) {
-	defer func() {
-		// Always update status to "approved", "rejected" or "failed".
-		noCtx := context.Background()
-		err = db.Reviews.LockForUpdate(noCtx, sq.Eq{"id": review.ID}, []string{"id DESC"}, func(update *Review) {
-			now := time.Now().UTC()
-			update.ProcessedAt = &now
-			if err != nil {
-				update.Status = ReviewStatusFailed
-				return
-			}
-			update.Status = review.Status
-		})
-		if err != nil {
-			log.Printf("failed to save review: %v", err)
-		}
-	}()
-
-	// Simulate long-running work.
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case <-time.After(1 * time.Second):
-	}
-
-	// Simulate external API call to an LLM.
-	if rand.Intn(2) == 0 {
-		return fmt.Errorf("failed to process review: <some underlying error>")
-	}
-
-	review.Status = ReviewStatusApproved
-	if rand.Intn(2) == 0 {
-		review.Status = ReviewStatusRejected
-	}
-	now := time.Now().UTC()
-	review.ProcessedAt = &now
-	return nil
-}

tests/worker_test.go

@@ -0,0 +1,64 @@
+package pgkit_test
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"math/rand"
+	"sync"
+	"time"
+
+	sq "github.com/Masterminds/squirrel"
+)
+
+type Worker struct {
+	DB *Database
+
+	wg sync.WaitGroup
+}
+
+func (w *Worker) Wait() {
+	w.wg.Wait()
+}
+
+func (w *Worker) ProcessReview(ctx context.Context, review *Review) (err error) {
+	w.wg.Add(1)
+	defer w.wg.Done()
+
+	defer func() {
+		// Always update review status to "approved", "rejected" or "failed".
+		noCtx := context.Background()
+		err = w.DB.Reviews.LockForUpdate(noCtx, sq.Eq{"id": review.ID}, []string{"id DESC"}, func(update *Review) {
+			now := time.Now().UTC()
+			update.ProcessedAt = &now
+			if err != nil {
+				update.Status = ReviewStatusFailed
+				return
+			}
+			update.Status = review.Status
+		})
+		if err != nil {
+			log.Printf("failed to save review: %v", err)
+		}
+	}()
+
+	// Simulate long-running work.
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-time.After(1 * time.Second):
+	}
+
+	// Simulate external API call to an LLM.
+	if rand.Intn(2) == 0 {
+		return fmt.Errorf("failed to process review: <some underlying error>")
+	}
+
+	review.Status = ReviewStatusApproved
+	if rand.Intn(2) == 0 {
+		review.Status = ReviewStatusRejected
+	}
+	now := time.Now().UTC()
+	review.ProcessedAt = &now
+	return nil
+}