Vellum/internal/jobs/jobs_test.go

package jobs

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"path/filepath"
	"testing"

	"github.com/mithral/vellum/internal/core"
	"github.com/mithral/vellum/migrations"
)

func testDB(t *testing.T) *sql.DB {
	t.Helper()
	db, err := core.OpenDB(filepath.Join(t.TempDir(), "jobs.db"))
	if err != nil {
		t.Fatalf("OpenDB: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })
	if err := core.Migrate(context.Background(), db, migrations.FS, discard()); err != nil {
		t.Fatalf("Migrate: %v", err)
	}
	return db
}

func discard() *slog.Logger { return slog.New(slog.NewTextHandler(io.Discard, nil)) }

func stateOf(t *testing.T, db *sql.DB, id int64) (string, int, string) {
	t.Helper()
	var state, errMsg string
	var attempts int
	if err := db.QueryRow(
		`SELECT state, attempts, error FROM jobs WHERE id = ?`, id).Scan(&state, &attempts, &errMsg); err != nil {
		t.Fatal(err)
	}
	return state, attempts, errMsg
}

func TestRunOnceExecutesHandler(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())

	var got json.RawMessage
	q.Register("greet", func(_ context.Context, payload json.RawMessage) error {
		got = payload
		return nil
	})

	if err := q.Enqueue(context.Background(), "greet", map[string]string{"name": "ada"}); err != nil {
		t.Fatal(err)
	}

	ran, err := q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if !ran {
		t.Fatal("RunOnce reported nothing to do")
	}
	if string(got) != `{"name":"ada"}` {
		t.Errorf("payload = %s", got)
	}

	// Nothing left to run.
	ran, err = q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if ran {
		t.Error("RunOnce ran a job that was already done")
	}

	var state string
	var n int
	if err := db.QueryRow(`SELECT state, count(*) FROM jobs GROUP BY state`).Scan(&state, &n); err != nil {
		t.Fatal(err)
	}
	if state != StateDone || n != 1 {
		t.Errorf("state=%s count=%d, want done/1", state, n)
	}
}

func TestUnknownKindFailsImmediately(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())

	if err := q.Enqueue(context.Background(), "nobody-handles-this", nil); err != nil {
		t.Fatal(err)
	}
	if _, err := q.RunOnce(context.Background()); err != nil {
		t.Fatal(err)
	}

	var state, errMsg string
	if err := db.QueryRow(`SELECT state, error FROM jobs`).Scan(&state, &errMsg); err != nil {
		t.Fatal(err)
	}
	if state != StateFailed {
		t.Errorf("state = %s, want failed (a missing handler must not retry)", state)
	}
	if errMsg == "" {
		t.Error("failure reason was not recorded")
	}
}

func TestFailureBacksOffThenFails(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())
	q.maxAttempts = 2
	q.baseBackoff = 30 * 1000 * 1000 * 1000 // 30s

	boom := errors.New("transient")
	q.Register("flaky", func(context.Context, json.RawMessage) error { return boom })

	if err := q.Enqueue(context.Background(), "flaky", nil); err != nil {
		t.Fatal(err)
	}

	// First run fails and reschedules into the future.
	if _, err := q.RunOnce(context.Background()); err != nil {
		t.Fatal(err)
	}
	state, attempts, errMsg := stateOf(t, db, 1)
	if state != StateQueued || attempts != 1 {
		t.Fatalf("after first failure: state=%s attempts=%d", state, attempts)
	}
	if errMsg == "" {
		t.Error("failure reason was not recorded")
	}

	// Not due yet, so nothing runs.
	ran, err := q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if ran {
		t.Error("a backed-off job ran before its run_after")
	}

	// Force it due and run again: attempts is now at the ceiling.
	if _, err := db.Exec(`UPDATE jobs SET run_after = '1970-01-01T00:00:00Z'`); err != nil {
		t.Fatal(err)
	}
	if _, err := q.RunOnce(context.Background()); err != nil {
		t.Fatal(err)
	}
	state, attempts, _ = stateOf(t, db, 1)
	if state != StateFailed {
		t.Errorf("after exhausting attempts: state=%s, want failed", state)
	}
	if attempts != 2 {
		t.Errorf("attempts = %d, want 2", attempts)
	}
}

func TestEnqueueAfterDelays(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())
	q.Register("later", func(context.Context, json.RawMessage) error { return nil })

	if err := q.EnqueueAfter(context.Background(), "later", nil, 60*1000*1000*1000); err != nil {
		t.Fatal(err)
	}
	ran, err := q.RunOnce(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if ran {
		t.Error("a delayed job ran immediately")
	}
}

func TestStats(t *testing.T) {
	db := testDB(t)
	q := New(db, discard())
	q.Register("noop", func(context.Context, json.RawMessage) error { return nil })

	for i := 0; i < 3; i++ {
		if err := q.Enqueue(context.Background(), "noop", nil); err != nil {
			t.Fatal(err)
		}
	}
	if _, err := q.RunOnce(context.Background()); err != nil {
		t.Fatal(err)
	}

	stats, err := q.Stats(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if stats[StateDone] != 1 || stats[StateQueued] != 2 {
		t.Errorf("stats = %v, want done:1 queued:2", stats)
	}
}