scobleclaw/pool-manager/pgqueue/pgqueue_test.go

package pgqueue

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

// testPool connects to the local pool_manager DB.
// Skips the test if PGBOUNCER or PG is not available.
func testPool(t *testing.T) *pgxpool.Pool {
	t.Helper()
	dsn := os.Getenv("PGQUEUE_TEST_DSN")
	if dsn == "" {
		dsn = "postgres://sscoble@/pool_manager?host=/run/postgresql"
	}
	pool, err := pgxpool.New(context.Background(), dsn)
	if err != nil {
		t.Skipf("no postgres available: %v", err)
	}
	return pool
}

func cleanup(t *testing.T, pool *pgxpool.Pool, queue string) {
	t.Helper()
	_, err := pool.Exec(context.Background(),
		"DELETE FROM pgqueue.jobs WHERE queue = $1", queue)
	if err != nil {
		t.Logf("cleanup: %v", err)
	}
}

func TestEnqueueClaimAck(t *testing.T) {
	pool := testPool(t)
	defer pool.Close()
	ctx := context.Background()
	q := fmt.Sprintf("test_e2e_%d", time.Now().UnixNano())
	defer cleanup(t, pool, q)

	// Enqueue
	payload, _ := json.Marshal(map[string]any{"task": "hello", "n": 42})
	res, err := Enqueue(ctx, pool, q, payload)
	if err != nil {
		t.Fatalf("enqueue: %v", err)
	}
	if res.ID == 0 {
		t.Fatal("got id=0")
	}

	// Claim
	job, err := Claim(ctx, pool, q, "test-worker", 30*time.Second)
	if err != nil {
		t.Fatalf("claim: %v", err)
	}
	if job == nil {
		t.Fatal("got nil job")
	}
	if job.ID != res.ID {
		t.Fatalf("id mismatch: enqueue=%d claim=%d", res.ID, job.ID)
	}
	if job.Attempts != 1 {
		t.Fatalf("expected attempts=1, got %d", job.Attempts)
	}

	// Ack
	if err := Ack(ctx, pool, job.ID, "test-worker"); err != nil {
		t.Fatalf("ack: %v", err)
	}

	// Verify done
	var state string
	err = pool.QueryRow(ctx, "SELECT state FROM pgqueue.jobs WHERE id=$1", job.ID).Scan(&state)
	if err != nil {
		t.Fatalf("verify: %v", err)
	}
	if state != "done" {
		t.Fatalf("expected state=done, got %q", state)
	}
}

func TestNackRequeue(t *testing.T) {
	pool := testPool(t)
	defer pool.Close()
	ctx := context.Background()
	q := fmt.Sprintf("test_nack_requeue_%d", time.Now().UnixNano())
	defer cleanup(t, pool, q)

	payload, _ := json.Marshal(map[string]any{"task": "retry"})
	res, _ := Enqueue(ctx, pool, q, payload)

	job, _ := Claim(ctx, pool, q, "w1", 30*time.Second)
	if job == nil {
		t.Fatal("expected a job")
	}

	r, err := Nack(ctx, pool, job.ID, "w1", "transient", 100*time.Millisecond)
	if err != nil {
		t.Fatalf("nack: %v", err)
	}
	if r != NackRequeued {
		t.Fatalf("expected requeued, got %q", r)
	}

	// After requeue, can claim again
	job2, _ := Claim(ctx, pool, q, "w1", 30*time.Second)
	if job2 == nil {
		t.Fatal("expected requeued job to be claimable")
	}
	if job2.ID != res.ID {
		t.Fatal("requeued job has different id")
	}
}

func TestIdempotencyKey(t *testing.T) {
	pool := testPool(t)
	defer pool.Close()
	ctx := context.Background()
	q := fmt.Sprintf("test_idem_%d", time.Now().UnixNano())
	defer cleanup(t, pool, q)

	key := fmt.Sprintf("key-%d", time.Now().UnixNano())
	p1, _ := json.Marshal(map[string]any{"v": 1})
	p2, _ := json.Marshal(map[string]any{"v": 2})

	r1, _ := Enqueue(ctx, pool, q, p1, WithIdempotencyKey(key))
	r2, _ := Enqueue(ctx, pool, q, p2, WithIdempotencyKey(key))

	if r1.ID != r2.ID {
		t.Fatalf("idem: r1=%d r2=%d, expected equal", r1.ID, r2.ID)
	}
}

func TestJanitor(t *testing.T) {
	pool := testPool(t)
	defer pool.Close()
	ctx := context.Background()
	q := fmt.Sprintf("test_janitor_%d", time.Now().UnixNano())
	defer cleanup(t, pool, q)

	// Enqueue, claim with very short lease
	payload, _ := json.Marshal(map[string]any{"task": "will-stale"})
	Enqueue(ctx, pool, q, payload)
	job, _ := Claim(ctx, pool, q, "w1", 100*time.Millisecond)
	if job == nil {
		t.Fatal("expected job")
	}

	// Wait for lease to expire
	time.Sleep(200 * time.Millisecond)

	// Run janitor
	n, err := Janitor(ctx, pool)
	if err != nil {
		t.Fatalf("janitor: %v", err)
	}
	if n < 1 {
		t.Fatalf("expected janitor to reset ≥1, got %d", n)
	}

	// Verify job is back to ready
	var state string
	err = pool.QueryRow(ctx, "SELECT state FROM pgqueue.jobs WHERE id=$1", job.ID).Scan(&state)
	if err != nil {
		t.Fatalf("verify: %v", err)
	}
	if state != "ready" {
		t.Fatalf("expected state=ready, got %q", state)
	}
}

func TestHeartbeat(t *testing.T) {
	pool := testPool(t)
	defer pool.Close()
	ctx := context.Background()
	q := fmt.Sprintf("test_hb_%d", time.Now().UnixNano())
	defer cleanup(t, pool, q)

	Enqueue(ctx, pool, q, []byte(`{}`))
	job, _ := Claim(ctx, pool, q, "w1", 100*time.Millisecond)
	if job == nil {
		t.Fatal("expected job")
	}

	// Heartbeat
	ok, err := Heartbeat(ctx, pool, job.ID, "w1", 5*time.Second)
	if err != nil {
		t.Fatalf("heartbeat: %v", err)
	}
	if !ok {
		t.Fatal("heartbeat returned false")
	}

	// Verify lease extended
	var newLease time.Time
	pool.QueryRow(ctx, "SELECT lease_until FROM pgqueue.jobs WHERE id=$1", job.ID).Scan(&newLease)
	if newLease.Before(time.Now().Add(4 * time.Second)) {
		t.Fatalf("heartbeat did not extend lease: %v", newLease)
	}
}