MealQuest/internal/jobs/jobs_test.go

package jobs

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/internal/store"
	"code.cubecraftlabs.com/Maniacs_Mansion/Meal-Quest/migrations"
)

// sentinelSecret stands in for a credential a handler holds while it runs.
const sentinelSecret = "sk-sentinel-5f2c9a"

type countPayload struct {
	Steps int `json:"steps"`
}

type countResult struct {
	Done int `json:"done"`
}

func migratedStore(t *testing.T) *store.Store {
	t.Helper()
	ctx := context.Background()
	st, err := store.Open(ctx, t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		if err := st.Close(); err != nil {
			t.Error(err)
		}
	})
	if err := store.Migrate(ctx, st.Writer, migrations.FS); err != nil {
		t.Fatal(err)
	}
	return st
}

func newRunner(t *testing.T, st *store.Store, workers int) *Runner {
	t.Helper()
	r, err := NewRunner(st, slog.New(slog.NewTextHandler(io.Discard, nil)), workers)
	if err != nil {
		t.Fatal(err)
	}
	return r
}

// start starts r and stops it when the test ends. The returned function
// stops it early and waits for its workers.
func start(t *testing.T, r *Runner) (stop func()) {
	t.Helper()
	ctx, cancel := context.WithCancel(context.Background())
	if err := r.Start(ctx); err != nil {
		cancel()
		t.Fatal(err)
	}
	stop = func() {
		cancel()
		r.Wait()
	}
	t.Cleanup(stop)
	return stop
}

func register[P, R any](t *testing.T, r *Runner, kind string, h Handler[P, R]) *Kind[P] {
	t.Helper()
	k, err := Register(r, kind, h)
	if err != nil {
		t.Fatal(err)
	}
	return k
}

func enqueue[P any](t *testing.T, k *Kind[P], payload P) string {
	t.Helper()
	id, err := k.Enqueue(context.Background(), payload)
	if err != nil {
		t.Fatal(err)
	}
	return id
}

// waitFor polls Get until done accepts the job, failing after five seconds.
func waitFor(t *testing.T, r *Runner, id string, done func(Job) bool) Job {
	t.Helper()
	deadline := time.Now().Add(5 * time.Second)
	for {
		job, err := r.Get(context.Background(), id)
		if err != nil {
			t.Fatal(err)
		}
		if done(job) {
			return job
		}
		if time.Now().After(deadline) {
			t.Fatalf("job %s never reached the expected state; last %+v (result %s)", id, job, job.Result)
		}
		time.Sleep(5 * time.Millisecond)
	}
}

func finished(j Job) bool { return j.Status == StatusSucceeded || j.Status == StatusFailed }

// rawRow returns every column of the job's row joined, for leak checks.
func rawRow(t *testing.T, st *store.Store, id string) string {
	t.Helper()
	var cols [10]any
	ptrs := make([]any, len(cols))
	for i := range cols {
		ptrs[i] = &cols[i]
	}
	if err := st.Reader.QueryRow("SELECT * FROM jobs WHERE id = ?", id).Scan(ptrs...); err != nil {
		t.Fatal(err)
	}
	return fmt.Sprint(cols[:]...)
}

func TestJobSucceeds(t *testing.T) {
	t.Parallel()
	st := migratedStore(t)
	r := newRunner(t, st, 2)
	count := register(t, r, "count", func(_ context.Context, p countPayload, _ func(countResult) error) (countResult, error) {
		return countResult{Done: p.Steps}, nil
	})
	start(t, r)

	id := enqueue(t, count, countPayload{Steps: 3})
	job := waitFor(t, r, id, finished)
	if job.ID != id || job.Kind != "count" || job.Status != StatusSucceeded || job.Error != nil {
		t.Fatalf("job = %+v, want succeeded count job %s with no error", job, id)
	}
	if string(job.Result) != `{"done":3}` {
		t.Errorf("result = %s, want {\"done\":3}", job.Result)
	}

	var startedAt, finishedAt, createdAt, updatedAt, payload string
	if err := st.Reader.QueryRow("SELECT started_at, finished_at, created_at, updated_at, payload FROM jobs WHERE id = ?", id).
		Scan(&startedAt, &finishedAt, &createdAt, &updatedAt, &payload); err != nil {
		t.Fatal(err)
	}
	for name, v := range map[string]string{"started_at": startedAt, "finished_at": finishedAt, "created_at": createdAt, "updated_at": updatedAt} {
		var ts store.Timestamp
		if err := ts.Scan(v); err != nil {
			t.Errorf("%s = %q: %v", name, v, err)
		}
	}
	if payload != `{"steps":3}` {
		t.Errorf("payload = %s, want {\"steps\":3}", payload)
	}
}

func TestJobFailsWithFailureMessage(t *testing.T) {
	t.Parallel()
	r := newRunner(t, migratedStore(t), 1)
	k := register(t, r, "suggest", func(context.Context, countPayload, func(countResult) error) (countResult, error) {
		return countResult{}, fmt.Errorf("suggest: %w", &Failure{Message: "add some Regulars first"})
	})
	start(t, r)

	job := waitFor(t, r, enqueue(t, k, countPayload{}), finished)
	if job.Status != StatusFailed || job.Error == nil || *job.Error != "add some Regulars first" {
		t.Fatalf("job = %+v, want failed with the Failure message", job)
	}
	if job.Result != nil {
		t.Errorf("result = %s, want none", job.Result)
	}
}

// credential holds a secret the way config.Secret does: unexported, so
// encoding/json cannot see it.
type credential struct{ value string }

func (credential) String() string { return "[REDACTED]" }

type leakyResult struct {
	Dinners []string   `json:"dinners"`
	Key     credential `json:"key"`
	Token   string     `json:"-"`
	apiKey  string
}

// TestSecretsNeverStored proves a handler holding a secret cannot surface it
// through its typed result, its partial results or its error text.
func TestSecretsNeverStored(t *testing.T) {
	t.Parallel()
	st := migratedStore(t)
	r := newRunner(t, st, 1)
	secretResult := leakyResult{Dinners: []string{"Tacos"}, Key: credential{sentinelSecret}, Token: sentinelSecret, apiKey: sentinelSecret}
	ok := register(t, r, "typed", func(_ context.Context, _ countPayload, publish func(leakyResult) error) (leakyResult, error) {
		if err := publish(secretResult); err != nil {
			return leakyResult{}, err
		}
		return secretResult, nil
	})
	failing := register(t, r, "leaky-error", func(context.Context, countPayload, func(leakyResult) error) (leakyResult, error) {
		return leakyResult{}, fmt.Errorf("GET https://api.example.com/v1?key=%s: 401 %s", sentinelSecret, sentinelSecret)
	})
	start(t, r)

	typed := waitFor(t, r, enqueue(t, ok, countPayload{}), finished)
	if typed.Status != StatusSucceeded || string(typed.Result) != `{"dinners":["Tacos"],"key":{}}` {
		t.Errorf("typed job = %+v (result %s), want succeeded with only the declared data", typed, typed.Result)
	}
	failed := waitFor(t, r, enqueue(t, failing, countPayload{}), finished)
	if failed.Status != StatusFailed || failed.Error == nil || *failed.Error != MessageInternal {
		t.Errorf("leaky-error job = %+v, want failed with %q", failed, MessageInternal)
	}
	for _, id := range []string{typed.ID, failed.ID} {
		if row := rawRow(t, st, id); strings.Contains(row, sentinelSecret) {
			t.Errorf("job %s row contains the secret: %s", id, row)
		}
	}
}

func TestProgressiveResultsVisibleMidRun(t *testing.T) {
	t.Parallel()
	r := newRunner(t, migratedStore(t), 1)
	release := make(chan struct{})
	k := register(t, r, "draft", func(ctx context.Context, p countPayload, publish func(countResult) error) (countResult, error) {
		for i := 1; i < p.Steps; i++ {
			if err := publish(countResult{Done: i}); err != nil {
				return countResult{}, err
			}
			select {
			case <-release:
			case <-ctx.Done():
				return countResult{}, ctx.Err()
			}
		}
		return countResult{Done: p.Steps}, nil
	})
	start(t, r)

	id := enqueue(t, k, countPayload{Steps: 3})
	for step := 1; step < 3; step++ {
		want := fmt.Sprintf(`{"done":%d}`, step)
		job := waitFor(t, r, id, func(j Job) bool { return string(j.Result) == want })
		if job.Status != StatusRunning || job.Error != nil {
			t.Fatalf("mid-run job = %+v, want running with no error", job)
		}
		release <- struct{}{}
	}
	job := waitFor(t, r, id, finished)
	if job.Status != StatusSucceeded || string(job.Result) != `{"done":3}` {
		t.Fatalf("job = %+v (result %s), want succeeded with {\"done\":3}", job, job.Result)
	}
}

func TestRestartFailsInterruptedJobsAndRunsQueuedOnes(t *testing.T) {
	t.Parallel()
	st := migratedStore(t)

	first := newRunner(t, st, 1)
	started := make(chan struct{})
	k := register(t, first, "slow", func(ctx context.Context, _ countPayload, publish func(countResult) error) (countResult, error) {
		if err := publish(countResult{Done: 1}); err != nil {
			return countResult{}, err
		}
		close(started)
		<-ctx.Done()
		return countResult{}, ctx.Err()
	})
	stop := start(t, first)
	interrupted := enqueue(t, k, countPayload{})
	<-started
	waiting := enqueue(t, k, countPayload{}) // the only worker is busy
	stop()

	if job := waitFor(t, first, interrupted, func(Job) bool { return true }); job.Status != StatusRunning {
		t.Fatalf("after shutdown job = %+v, want left running", job)
	}

	second := newRunner(t, st, 1)
	register(t, second, "slow", func(context.Context, countPayload, func(countResult) error) (countResult, error) {
		return countResult{Done: 9}, nil
	})
	start(t, second)

	job := waitFor(t, second, interrupted, finished)
	if job.Status != StatusFailed || job.Error == nil || *job.Error != MessageInterrupted {
		t.Errorf("interrupted job = %+v, want failed with %q", job, MessageInterrupted)
	}
	if string(job.Result) != `{"done":1}` {
		t.Errorf("interrupted job result = %s, want the last partial result kept", job.Result)
	}
	if job := waitFor(t, second, waiting, finished); job.Status != StatusSucceeded {
		t.Errorf("queued job = %+v, want run by the new runner", job)
	}
}

func TestWorkerPoolIsBounded(t *testing.T) {
	t.Parallel()
	r := newRunner(t, migratedStore(t), 2)
	var running, peak atomic.Int32
	release := make(chan struct{})
	k := register(t, r, "busy", func(context.Context, countPayload, func(countResult) error) (countResult, error) {
		n := running.Add(1)
		for {
			p := peak.Load()
			if n <= p || peak.CompareAndSwap(p, n) {
				break
			}
		}
		<-release
		running.Add(-1)
		return countResult{}, nil
	})
	start(t, r)

	ids := make([]string, 5)
	for i := range ids {
		ids[i] = enqueue(t, k, countPayload{})
	}
	deadline := time.Now().Add(5 * time.Second)
	for running.Load() < 2 {
		if time.Now().After(deadline) {
			t.Fatal("two workers never ran at once")
		}
		time.Sleep(5 * time.Millisecond)
	}
	close(release)
	for _, id := range ids {
		if job := waitFor(t, r, id, finished); job.Status != StatusSucceeded {
			t.Errorf("job = %+v, want succeeded", job)
		}
	}
	if got := peak.Load(); got != 2 {
		t.Errorf("peak concurrent jobs = %d, want 2", got)
	}
}

func TestJobWithoutHandlerFails(t *testing.T) {
	t.Parallel()
	st := migratedStore(t)
	old := newRunner(t, st, 1)
	k := register(t, old, "retired", func(context.Context, countPayload, func(countResult) error) (countResult, error) {
		return countResult{}, nil
	})
	id := enqueue(t, k, countPayload{}) // never started

	r := newRunner(t, st, 1)
	start(t, r)
	job := waitFor(t, r, id, finished)
	if job.Status != StatusFailed || job.Error == nil || *job.Error != "no handler for job kind retired" {
		t.Errorf("job = %+v, want failed for a missing handler", job)
	}
}

func TestGetUnknownID(t *testing.T) {
	t.Parallel()
	r := newRunner(t, migratedStore(t), 1)
	if _, err := r.Get(context.Background(), store.NewID()); !errors.Is(err, ErrNotFound) {
		t.Errorf("Get unknown id error = %v, want ErrNotFound", err)
	}
}

func TestRegisterRefusesBadRegistrations(t *testing.T) {
	t.Parallel()
	r := newRunner(t, migratedStore(t), 1)
	h := func(context.Context, countPayload, func(countResult) error) (countResult, error) {
		return countResult{}, nil
	}
	register(t, r, "count", h)

	if _, err := Register(r, "count", h); err == nil {
		t.Error("duplicate kind registered")
	}
	if _, err := Register(r, "", h); err == nil {
		t.Error("empty kind registered")
	}
	if _, err := Register[countPayload, countResult](r, "nil", nil); err == nil {
		t.Error("nil handler registered")
	}
	start(t, r)
	if _, err := Register(r, "late", h); err == nil {
		t.Error("kind registered after Start")
	}
	if err := r.Start(context.Background()); err == nil {
		t.Error("second Start succeeded")
	}
}

func TestNewRunnerRefusesBadConfig(t *testing.T) {
	t.Parallel()
	st := migratedStore(t)
	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
	if _, err := NewRunner(st, logger, 0); err == nil {
		t.Error("zero workers accepted")
	}
	if _, err := NewRunner(nil, logger, 1); err == nil {
		t.Error("nil store accepted")
	}
	if _, err := NewRunner(st, nil, 1); err == nil {
		t.Error("nil logger accepted")
	}
}