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