pgqueue_test.go

package pgqueue_test

import (
	"context"
	"fmt"
	"os"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/goware/pgkit/v2"
	"github.com/goware/pgqueue"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var db *pgkit.DB

func TestMain(m *testing.M) {
	host := envOr("PGQUEUE_TEST_HOST", "localhost")
	user := envOr("PGQUEUE_TEST_USER", "postgres")
	pass := envOr("PGQUEUE_TEST_PASS", "postgres")
	name := envOr("PGQUEUE_TEST_DB", "pgqueue_test")

	var err error
	db, err = pgkit.Connect("pgqueue_test", pgkit.Config{
		Database: name,
		Host:     host,
		Username: user,
		Password: pass,
	})
	if err != nil {
		fmt.Fprintf(os.Stderr, "pgqueue_test: connect: %v (set PGQUEUE_TEST_* env to configure)\n", err)
		os.Exit(1)
	}
	if err := db.Conn.Ping(context.Background()); err != nil {
		fmt.Fprintf(os.Stderr, "pgqueue_test: ping: %v\n", err)
		os.Exit(1)
	}

	os.Exit(m.Run())
}

func envOr(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func setupQueue(t *testing.T) *pgqueue.Queue {
	t.Helper()
	ctx := t.Context()
	q := pgqueue.New(db)
	require.NoError(t, pgqueue.Migrate(ctx, q))
	_, err := db.Conn.Exec(ctx, "TRUNCATE pgqueue_tasks")
	require.NoError(t, err)
	return q
}

// startWorker launches a worker in a goroutine and returns a function to
// collect the Start error. Replaces bare `go func() { _ = w.Start(...) }()`.
func startWorker(t *testing.T, w *pgqueue.Worker, ctx context.Context) (startErr func() error) {
	t.Helper()
	ch := make(chan error, 1)
	go func() { ch <- w.Start(ctx) }()
	// Wait for worker to reach running state.
	require.Eventually(t, w.IsRunning, 5*time.Second, 10*time.Millisecond, "worker did not start in time")
	return func() error { return <-ch }
}

// --- Test payloads ---

type EmailPayload struct {
	To      string `json:"to"`
	Subject string `json:"subject"`
}

var emailSpec = pgqueue.JobSpec[EmailPayload]{
	Queue: "send-email",
	HashFn: func(p EmailPayload) *string {
		h := p.To + ":" + p.Subject
		return &h
	},
	MaxRetries: 2,
	RetryDelay: 100 * time.Millisecond,
}

type CounterPayload struct {
	Count int `json:"count"`
}

var counterTickerSpec = pgqueue.TickerSpec[CounterPayload]{
	Queue:          "counter",
	Key:            "counter-singleton",
	InitialPayload: CounterPayload{Count: 0},
	Every:          100 * time.Millisecond,
	MaxRetries:     2,
	RetryDelay:     50 * time.Millisecond,
	LeaseDuration:  10 * time.Second,
	FinalizeBuffer: 1 * time.Second,
}

// --- Enqueue tests ---

func TestEnqueue(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	t.Run("enqueue and get", func(t *testing.T) {
		id, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "a@b.com", Subject: "hi"})
		require.NoError(t, err)
		require.NotZero(t, id)

		task, err := q.Get(ctx, id)
		require.NoError(t, err)
		assert.Equal(t, "send-email", task.Queue)
		assert.Equal(t, pgqueue.TaskStatusPending, task.Status)

		var p EmailPayload
		require.NoError(t, task.DecodePayload(&p))
		assert.Equal(t, "a@b.com", p.To)
	})

	t.Run("dedup with ignore duplicates", func(t *testing.T) {
		_, err := db.Conn.Exec(ctx, "TRUNCATE pgqueue_tasks")
		require.NoError(t, err)

		id1, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "x@y.com", Subject: "dup"}, pgqueue.WithIgnoreDuplicates())
		require.NoError(t, err)

		id2, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "x@y.com", Subject: "dup"}, pgqueue.WithIgnoreDuplicates())
		require.NoError(t, err)
		assert.Equal(t, id1, id2, "should return same ID for duplicate")
	})

	t.Run("ignore duplicates requires hash", func(t *testing.T) {
		noHashSpec := pgqueue.JobSpec[EmailPayload]{Queue: "no-hash"}
		_, err := pgqueue.Enqueue(ctx, q, noHashSpec, EmailPayload{To: "a@b.com", Subject: "hi"}, pgqueue.WithIgnoreDuplicates())
		require.Error(t, err)
		assert.Contains(t, err.Error(), "WithIgnoreDuplicates requires")
	})

	t.Run("find by queue and hash", func(t *testing.T) {
		_, err := db.Conn.Exec(ctx, "TRUNCATE pgqueue_tasks")
		require.NoError(t, err)

		_, err = pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "find@me.com", Subject: "test"})
		require.NoError(t, err)

		task, err := q.Find(ctx, "send-email", "find@me.com:test")
		require.NoError(t, err)
		assert.Equal(t, "send-email", task.Queue)
	})
}

// --- Cancel tests ---

func TestCancel(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	t.Run("cancel pending task", func(t *testing.T) {
		id, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "c@d.com", Subject: "cancel"})
		require.NoError(t, err)

		require.NoError(t, q.Cancel(ctx, id))

		_, err = q.Get(ctx, id)
		require.ErrorIs(t, err, pgqueue.ErrTaskNotFound, "cancelled task should be deleted")
	})

	t.Run("cancel non-pending fails", func(t *testing.T) {
		err := q.Cancel(ctx, 99999999)
		require.ErrorIs(t, err, pgqueue.ErrTaskNotFound)
	})
}

// --- Recovery tests ---

func TestRecovery(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	t.Run("enable failed task", func(t *testing.T) {
		id, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "r@r.com", Subject: "fail"})
		require.NoError(t, err)

		_, err = db.Conn.Exec(ctx, "UPDATE pgqueue_tasks SET status = $1, try = 3 WHERE id = $2", pgqueue.TaskStatusFailed, id)
		require.NoError(t, err)

		require.NoError(t, q.Enable(ctx, id))

		task, err := q.Get(ctx, id)
		require.NoError(t, err)
		assert.Equal(t, pgqueue.TaskStatusPending, task.Status)
		assert.Equal(t, int16(0), task.Try)
	})

	t.Run("enable running task fails", func(t *testing.T) {
		id, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "r@r.com", Subject: "run"})
		require.NoError(t, err)

		now := time.Now().UTC()
		token := "test-token"
		_, err = db.Conn.Exec(ctx,
			"UPDATE pgqueue_tasks SET status = $1, claimed_at = $2, lease_until = $3, claim_token = $4 WHERE id = $5",
			pgqueue.TaskStatusRunning, now, now.Add(5*time.Minute), token, id,
		)
		require.NoError(t, err)

		err = q.Enable(ctx, id)
		require.ErrorIs(t, err, pgqueue.ErrTaskNotRecoverable)
	})

	t.Run("enable pending task fails", func(t *testing.T) {
		id, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "r@r.com", Subject: "pend"})
		require.NoError(t, err)

		err = q.Enable(ctx, id)
		require.ErrorIs(t, err, pgqueue.ErrTaskNotRecoverable)
	})

	t.Run("replace payload on disabled task", func(t *testing.T) {
		id, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "old@old.com", Subject: "old"})
		require.NoError(t, err)

		_, err = db.Conn.Exec(ctx, "UPDATE pgqueue_tasks SET status = $1 WHERE id = $2", pgqueue.TaskStatusDisabled, id)
		require.NoError(t, err)

		require.NoError(t, pgqueue.ReplacePayloadJSON(ctx, q, id, EmailPayload{To: "new@new.com", Subject: "fixed"}))

		task, err := q.Get(ctx, id)
		require.NoError(t, err)
		var p EmailPayload
		require.NoError(t, task.DecodePayload(&p))
		assert.Equal(t, "new@new.com", p.To)
	})

	t.Run("replace payload rejects invalid json", func(t *testing.T) {
		id, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "x@x.com", Subject: "x"})
		require.NoError(t, err)

		_, err = db.Conn.Exec(ctx, "UPDATE pgqueue_tasks SET status = $1 WHERE id = $2", pgqueue.TaskStatusFailed, id)
		require.NoError(t, err)

		err = q.ReplacePayload(ctx, id, []byte(`{invalid`))
		require.Error(t, err)
		assert.Contains(t, err.Error(), "not valid JSON")
	})
}

// --- Worker: job handler ---

type countingHandler struct {
	count   atomic.Int32
	started chan struct{} // closed on first RunTask call
	once    sync.Once
}

func newCountingHandler() *countingHandler {
	return &countingHandler{started: make(chan struct{})}
}

func (h *countingHandler) RunTask(_ context.Context, _ *pgqueue.Job[EmailPayload]) pgqueue.Result {
	h.once.Do(func() { close(h.started) })
	h.count.Add(1)
	return pgqueue.Done()
}

func TestWorkerJobHappyPath(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	handler := newCountingHandler()
	w := pgqueue.NewWorker(q)
	spec := emailSpec
	spec.PollInterval = 50 * time.Millisecond
	spec.LeaseDuration = 10 * time.Second
	spec.FinalizeBuffer = 1 * time.Second
	require.NoError(t, pgqueue.Register(w, spec, handler))

	_, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "w@w.com", Subject: "work"})
	require.NoError(t, err)

	workerCtx, workerCancel := context.WithCancel(ctx)
	defer workerCancel()
	collectErr := startWorker(t, w, workerCtx)

	require.Eventually(t, func() bool { return handler.count.Load() >= 1 }, 5*time.Second, 50*time.Millisecond)

	drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
	defer drainCancel()
	require.NoError(t, w.Stop(drainCtx))
	workerCancel()
	require.NoError(t, collectErr())

	tasks, _, err := q.List(ctx, nil, nil, nil)
	require.NoError(t, err)
	require.Len(t, tasks, 1)
	assert.Equal(t, pgqueue.TaskStatusCompleted, tasks[0].Status)
}

// --- Worker: no double-processing ---

func TestWorkerNoDoubleProcess(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	var processed sync.Map
	var dupes atomic.Int32

	spec := emailSpec
	spec.PollInterval = 50 * time.Millisecond
	spec.LeaseDuration = 10 * time.Second
	spec.FinalizeBuffer = 1 * time.Second

	for i := range 20 {
		noHashSpec := pgqueue.JobSpec[EmailPayload]{Queue: "send-email"}
		_, err := pgqueue.Enqueue(ctx, q, noHashSpec, EmailPayload{To: fmt.Sprintf("u%d@test.com", i), Subject: "test"})
		require.NoError(t, err)
	}

	workers := make([]*pgqueue.Worker, 3)
	workerCtx, workerCancel := context.WithCancel(ctx)
	defer workerCancel()
	var collectErrs []func() error

	for i := range workers {
		h := pgqueue.JobHandlerFunc[EmailPayload](func(_ context.Context, job *pgqueue.Job[EmailPayload]) pgqueue.Result {
			if _, loaded := processed.LoadOrStore(job.Task.ID, true); loaded {
				dupes.Add(1)
			}
			time.Sleep(10 * time.Millisecond)
			return pgqueue.Done()
		})
		workers[i] = pgqueue.NewWorker(q)
		require.NoError(t, pgqueue.Register(workers[i], spec, h))
		collectErrs = append(collectErrs, startWorker(t, workers[i], workerCtx))
	}

	require.Eventually(t, func() bool {
		count := 0
		processed.Range(func(_, _ any) bool { count++; return true })
		return count >= 20
	}, 10*time.Second, 100*time.Millisecond)

	drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
	defer drainCancel()
	for _, w := range workers {
		require.NoError(t, w.Stop(drainCtx))
	}
	workerCancel()
	for _, f := range collectErrs {
		require.NoError(t, f())
	}

	assert.Zero(t, dupes.Load(), "no task should be processed twice")
}

// --- Worker: retry and max retries ---

type failingHandler struct {
	attempts atomic.Int32
}

func (h *failingHandler) RunTask(_ context.Context, _ *pgqueue.Job[EmailPayload]) pgqueue.Result {
	h.attempts.Add(1)
	return pgqueue.Retry(fmt.Errorf("transient error"))
}

func TestWorkerRetry(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	handler := &failingHandler{}
	spec := emailSpec
	spec.PollInterval = 50 * time.Millisecond
	spec.MaxRetries = 2
	spec.RetryDelay = 50 * time.Millisecond
	spec.LeaseDuration = 10 * time.Second
	spec.FinalizeBuffer = 1 * time.Second

	w := pgqueue.NewWorker(q)
	require.NoError(t, pgqueue.Register(w, spec, handler))

	_, err := pgqueue.Enqueue(ctx, q, emailSpec, EmailPayload{To: "retry@test.com", Subject: "retry"})
	require.NoError(t, err)

	workerCtx, workerCancel := context.WithCancel(ctx)
	defer workerCancel()
	collectErr := startWorker(t, w, workerCtx)

	require.Eventually(t, func() bool {
		tasks, _, _ := q.List(ctx, nil, nil, nil)
		return len(tasks) == 1 && tasks[0].Status == pgqueue.TaskStatusFailed
	}, 5*time.Second, 100*time.Millisecond)

	drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
	defer drainCancel()
	require.NoError(t, w.Stop(drainCtx))
	workerCancel()
	require.NoError(t, collectErr())

	assert.GreaterOrEqual(t, handler.attempts.Load(), int32(2))
}

// --- Worker: ticker ---

type counterTickerHandler struct {
	runs atomic.Int32
}

func (h *counterTickerHandler) RunTick(_ context.Context, job *pgqueue.Job[CounterPayload]) pgqueue.Result {
	job.Payload.Count++
	h.runs.Add(1)
	return pgqueue.Done()
}

func TestWorkerTicker(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	handler := &counterTickerHandler{}
	w := pgqueue.NewWorker(q)
	require.NoError(t, pgqueue.RegisterTicker(w, counterTickerSpec, handler))

	workerCtx, workerCancel := context.WithCancel(ctx)
	defer workerCancel()
	collectErr := startWorker(t, w, workerCtx)

	require.Eventually(t, func() bool { return handler.runs.Load() >= 3 }, 5*time.Second, 50*time.Millisecond)

	drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
	defer drainCancel()
	require.NoError(t, w.Stop(drainCtx))
	workerCancel()
	require.NoError(t, collectErr())

	task, err := q.Find(ctx, "counter", "counter-singleton")
	require.NoError(t, err)
	var p CounterPayload
	require.NoError(t, task.DecodePayload(&p))
	assert.GreaterOrEqual(t, p.Count, 3, "ticker payload should persist across runs")
}

// --- Worker: stale finalize after reap ---

func TestWorkerStaleFinalize(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	// Handler that blocks on its own signal — ignores context deadline so the
	// poll goroutine stays stuck past lease expiry.
	handlerStarted := make(chan struct{})
	releaseHandler := make(chan struct{})
	blocking := pgqueue.JobHandlerFunc[EmailPayload](func(_ context.Context, _ *pgqueue.Job[EmailPayload]) pgqueue.Result {
		close(handlerStarted)
		<-releaseHandler // block until test releases
		return pgqueue.Done()
	})

	// Worker A: very short lease so we can manually reap.
	wA := pgqueue.NewWorker(q)
	specA := pgqueue.JobSpec[EmailPayload]{
		Queue:          "stale-test",
		PollInterval:   50 * time.Millisecond,
		LeaseDuration:  500 * time.Millisecond,
		FinalizeBuffer: 100 * time.Millisecond,
	}
	require.NoError(t, pgqueue.Register(wA, specA, blocking))

	// Worker B: completes tasks normally.
	completedByB := make(chan int64, 1)
	handlerB := pgqueue.JobHandlerFunc[EmailPayload](func(_ context.Context, job *pgqueue.Job[EmailPayload]) pgqueue.Result {
		completedByB <- job.Task.ID
		return pgqueue.Done()
	})
	wB := pgqueue.NewWorker(q)
	specB := pgqueue.JobSpec[EmailPayload]{
		Queue:          "stale-test",
		PollInterval:   50 * time.Millisecond,
		LeaseDuration:  10 * time.Second,
		FinalizeBuffer: 1 * time.Second,
	}
	require.NoError(t, pgqueue.Register(wB, specB, handlerB))

	// Enqueue one task.
	id, err := pgqueue.Enqueue(ctx, q, pgqueue.JobSpec[EmailPayload]{Queue: "stale-test"}, EmailPayload{To: "stale@test.com", Subject: "stale"})
	require.NoError(t, err)

	// Start worker A — it will claim the task and block in the handler.
	ctxA, cancelA := context.WithCancel(ctx)
	defer cancelA()
	collectA := startWorker(t, wA, ctxA)

	// Wait for handler A to start executing.
	select {
	case <-handlerStarted:
	case <-time.After(5 * time.Second):
		t.Fatal("handler A did not start")
	}

	// Wait for lease to actually expire in the DB.
	require.Eventually(t, func() bool {
		task, err := q.Get(ctx, id)
		return err == nil && task.LeaseUntil != nil && task.LeaseUntil.Before(time.Now().UTC())
	}, 3*time.Second, 50*time.Millisecond, "lease should expire in the DB")

	// Manually reap: reset the expired task to pending (simulates reaper).
	// Enforces lease_until < NOW() — same condition as the real reaper.
	tag, err := db.Conn.Exec(ctx,
		"UPDATE pgqueue_tasks SET status = $1, run_at = NOW(), claimed_at = NULL, lease_until = NULL, claim_token = NULL WHERE id = $2 AND status = $3 AND lease_until < NOW()",
		pgqueue.TaskStatusPending, id, pgqueue.TaskStatusRunning,
	)
	require.NoError(t, err)
	require.Equal(t, int64(1), tag.RowsAffected(), "reap should affect 1 row (lease expired)")

	// Start worker B — it will reclaim the reaped task with a new claim_token.
	ctxB, cancelB := context.WithCancel(ctx)
	defer cancelB()
	collectB := startWorker(t, wB, ctxB)

	// Wait for B to complete the task.
	select {
	case completedID := <-completedByB:
		assert.Equal(t, id, completedID)
	case <-time.After(5 * time.Second):
		t.Fatal("worker B did not pick up reaped task")
	}

	// Verify task is completed (by worker B).
	require.Eventually(t, func() bool {
		task, err := q.Get(ctx, id)
		return err == nil && task.Status == pgqueue.TaskStatusCompleted
	}, 3*time.Second, 50*time.Millisecond)

	// Snapshot task state after B completed it, before releasing A.
	// B's claim incremented Try to 2 (A had it at 1). If A's stale finalize
	// overwrites the row, Try would revert to 1.
	taskBeforeRelease, err := q.Get(ctx, id)
	require.NoError(t, err)
	require.Equal(t, pgqueue.TaskStatusCompleted, taskBeforeRelease.Status)
	require.Equal(t, int16(2), taskBeforeRelease.Try, "B's claim should have incremented Try to 2")

	// Release worker A's handler. It will return Done(), then poll() will
	// attempt to finalize with the OLD claim_token → 0 rows (fenced out).
	close(releaseHandler)

	// Cancel worker A and wait for Start to return — this guarantees the
	// poll goroutine (including the stale finalize attempt) has completed.
	cancelA()
	require.NoError(t, collectA()) // Start returns nil after context cancel + clean exit

	drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
	defer drainCancel()
	require.NoError(t, wB.Stop(drainCtx))
	cancelB()
	require.NoError(t, collectB())

	// Final verification: compare full row against snapshot.
	// If A's stale finalize clobbered B, Try would be 1 (not 2) and
	// other fields could differ.
	taskAfterRelease, err := q.Get(ctx, id)
	require.NoError(t, err)
	assert.Equal(t, pgqueue.TaskStatusCompleted, taskAfterRelease.Status, "status must still be completed")
	assert.Equal(t, taskBeforeRelease.Try, taskAfterRelease.Try, "Try must not change (stale finalize fenced)")
	assert.Equal(t, taskBeforeRelease.LastRanAt, taskAfterRelease.LastRanAt, "LastRanAt must not change")
}

// --- Worker: handler context deadline ---

func TestWorkerHandlerDeadline(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	var gotDeadline atomic.Bool
	var deadlineWithin atomic.Bool

	handler := pgqueue.JobHandlerFunc[EmailPayload](func(ctx context.Context, _ *pgqueue.Job[EmailPayload]) pgqueue.Result {
		deadline, ok := ctx.Deadline()
		gotDeadline.Store(ok)
		if ok {
			// Deadline should be roughly leaseUntil - finalizeBuffer from now.
			remaining := time.Until(deadline)
			// LeaseDuration=2s, FinalizeBuffer=500ms → ~1.5s remaining.
			deadlineWithin.Store(remaining > 500*time.Millisecond && remaining < 3*time.Second)
		}
		return pgqueue.Done()
	})

	w := pgqueue.NewWorker(q)
	spec := pgqueue.JobSpec[EmailPayload]{
		Queue:          "deadline-test",
		PollInterval:   50 * time.Millisecond,
		LeaseDuration:  2 * time.Second,
		FinalizeBuffer: 500 * time.Millisecond,
	}
	require.NoError(t, pgqueue.Register(w, spec, handler))

	_, err := pgqueue.Enqueue(ctx, q, pgqueue.JobSpec[EmailPayload]{Queue: "deadline-test"}, EmailPayload{To: "dl@test.com", Subject: "dl"})
	require.NoError(t, err)

	workerCtx, workerCancel := context.WithCancel(ctx)
	defer workerCancel()
	collectErr := startWorker(t, w, workerCtx)

	require.Eventually(t, func() bool { return gotDeadline.Load() }, 5*time.Second, 50*time.Millisecond)

	drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
	defer drainCancel()
	require.NoError(t, w.Stop(drainCtx))
	workerCancel()
	require.NoError(t, collectErr())

	assert.True(t, gotDeadline.Load(), "handler context should have a deadline")
	assert.True(t, deadlineWithin.Load(), "deadline should be ~leaseUntil-finalizeBuffer")
}

// --- Worker: shutdown ---

func TestWorkerShutdown(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	t.Run("clean drain", func(t *testing.T) {
		w := pgqueue.NewWorker(q)
		spec := pgqueue.JobSpec[EmailPayload]{Queue: "drain-test", PollInterval: 50 * time.Millisecond}
		require.NoError(t, pgqueue.Register(w, spec, newCountingHandler()))

		workerCtx, workerCancel := context.WithCancel(ctx)
		defer workerCancel()
		collectErr := startWorker(t, w, workerCtx)

		drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
		defer drainCancel()
		require.NoError(t, w.Stop(drainCtx), "should drain cleanly")

		workerCancel()
		require.NoError(t, collectErr())
	})

	t.Run("timeout drain", func(t *testing.T) {
		handlerStarted := make(chan struct{})
		blocking := pgqueue.JobHandlerFunc[EmailPayload](func(ctx context.Context, _ *pgqueue.Job[EmailPayload]) pgqueue.Result {
			close(handlerStarted)
			<-ctx.Done()
			return pgqueue.Done()
		})
		w := pgqueue.NewWorker(q)
		spec := pgqueue.JobSpec[EmailPayload]{Queue: "block-test", PollInterval: 50 * time.Millisecond, LeaseDuration: 30 * time.Second, FinalizeBuffer: 1 * time.Second}
		require.NoError(t, pgqueue.Register(w, spec, blocking))

		_, err := pgqueue.Enqueue(ctx, q, pgqueue.JobSpec[EmailPayload]{Queue: "block-test"}, EmailPayload{To: "block@test.com", Subject: "block"})
		require.NoError(t, err)

		workerCtx, workerCancel := context.WithCancel(ctx)
		defer workerCancel()
		collectErr := startWorker(t, w, workerCtx)

		// Wait for the handler to be actively blocking.
		select {
		case <-handlerStarted:
		case <-time.After(5 * time.Second):
			t.Fatal("handler did not start")
		}

		drainCtx, drainCancel := context.WithTimeout(ctx, 200*time.Millisecond)
		defer drainCancel()
		err = w.Stop(drainCtx)
		assert.ErrorIs(t, err, context.DeadlineExceeded, "should timeout")
		workerCancel()
		// Start returns nil after workerCancel — goroutines exit once workerCtx is cancelled.
		require.NoError(t, collectErr())
	})

	t.Run("start after stop fails", func(t *testing.T) {
		w := pgqueue.NewWorker(q)
		spec := pgqueue.JobSpec[EmailPayload]{Queue: "once-test", PollInterval: 50 * time.Millisecond}
		require.NoError(t, pgqueue.Register(w, spec, newCountingHandler()))

		workerCtx, workerCancel := context.WithCancel(ctx)
		defer workerCancel()
		collectErr := startWorker(t, w, workerCtx)

		drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
		defer drainCancel()
		require.NoError(t, w.Stop(drainCtx))
		workerCancel()
		require.NoError(t, collectErr())

		err := w.Start(ctx)
		require.ErrorIs(t, err, pgqueue.ErrWorkerStopped)
	})
}

// --- Worker: state machine ---

func TestWorkerStateMachine(t *testing.T) {
	q := setupQueue(t)

	t.Run("register after start fails", func(t *testing.T) {
		w := pgqueue.NewWorker(q)
		spec := pgqueue.JobSpec[EmailPayload]{Queue: "sm-test", PollInterval: 50 * time.Millisecond}
		require.NoError(t, pgqueue.Register(w, spec, newCountingHandler()))

		workerCtx, workerCancel := context.WithCancel(t.Context())
		defer workerCancel()
		collectErr := startWorker(t, w, workerCtx)

		err := pgqueue.Register(w, pgqueue.JobSpec[EmailPayload]{Queue: "late"}, newCountingHandler())
		require.Error(t, err)
		assert.Contains(t, err.Error(), "cannot register")

		drainCtx, drainCancel := context.WithTimeout(t.Context(), 5*time.Second)
		defer drainCancel()
		require.NoError(t, w.Stop(drainCtx))
		workerCancel()
		require.NoError(t, collectErr())
	})

	t.Run("duplicate queue fails", func(t *testing.T) {
		w := pgqueue.NewWorker(q)
		spec := pgqueue.JobSpec[EmailPayload]{Queue: "dup-queue", PollInterval: 50 * time.Millisecond}
		require.NoError(t, pgqueue.Register(w, spec, newCountingHandler()))
		err := pgqueue.Register(w, spec, newCountingHandler())
		require.Error(t, err)
		assert.Contains(t, err.Error(), "already registered")
	})
}

// --- Worker: real reaper reclaims expired lease ---

func TestWorkerReaper(t *testing.T) {
	q := setupQueue(t)
	ctx := t.Context()

	// Worker A: handler blocks on signal, short lease, fast reaper.
	handlerStarted := make(chan struct{})
	releaseHandler := make(chan struct{})
	blockingHandler := pgqueue.JobHandlerFunc[EmailPayload](func(_ context.Context, _ *pgqueue.Job[EmailPayload]) pgqueue.Result {
		close(handlerStarted)
		<-releaseHandler
		return pgqueue.Done()
	})

	wA := pgqueue.NewWorker(q, pgqueue.WithReapInterval(200*time.Millisecond))
	specA := pgqueue.JobSpec[EmailPayload]{
		Queue:          "reaper-test",
		PollInterval:   50 * time.Millisecond,
		LeaseDuration:  500 * time.Millisecond,
		FinalizeBuffer: 100 * time.Millisecond,
	}
	require.NoError(t, pgqueue.Register(wA, specA, blockingHandler))

	// Worker B: same queue, completes normally.
	completedByB := make(chan int64, 1)
	handlerB := pgqueue.JobHandlerFunc[EmailPayload](func(_ context.Context, job *pgqueue.Job[EmailPayload]) pgqueue.Result {
		completedByB <- job.Task.ID
		return pgqueue.Done()
	})
	wB := pgqueue.NewWorker(q, pgqueue.WithReapInterval(200*time.Millisecond))
	specB := pgqueue.JobSpec[EmailPayload]{
		Queue:          "reaper-test",
		PollInterval:   50 * time.Millisecond,
		LeaseDuration:  10 * time.Second,
		FinalizeBuffer: 1 * time.Second,
	}
	require.NoError(t, pgqueue.Register(wB, specB, handlerB))

	// Enqueue one task.
	id, err := pgqueue.Enqueue(ctx, q, pgqueue.JobSpec[EmailPayload]{Queue: "reaper-test"}, EmailPayload{To: "reap@test.com", Subject: "reap"})
	require.NoError(t, err)

	// Start worker A — it claims the task and blocks.
	ctxA, cancelA := context.WithCancel(ctx)
	defer cancelA()
	collectA := startWorker(t, wA, ctxA)

	select {
	case <-handlerStarted:
	case <-time.After(5 * time.Second):
		t.Fatal("handler A did not start")
	}

	// Start worker B — its reaper will reclaim the expired task.
	ctxB, cancelB := context.WithCancel(ctx)
	defer cancelB()
	collectB := startWorker(t, wB, ctxB)

	// Wait for B to complete the reaped task. The reaper runs every 200ms,
	// lease is 500ms, so within ~1s the task should be reaped and reclaimed.
	select {
	case completedID := <-completedByB:
		assert.Equal(t, id, completedID, "B should process the same task")
	case <-time.After(5 * time.Second):
		t.Fatal("worker B did not pick up reaped task")
	}

	// Release A and shut down both workers.
	close(releaseHandler)
	cancelA()
	require.NoError(t, collectA())

	drainCtx, drainCancel := context.WithTimeout(ctx, 5*time.Second)
	defer drainCancel()
	require.NoError(t, wB.Stop(drainCtx))
	cancelB()
	require.NoError(t, collectB())

	// Verify task completed.
	task, err := q.Get(ctx, id)
	require.NoError(t, err)
	assert.Equal(t, pgqueue.TaskStatusCompleted, task.Status)
}