orbit/sdk/go/worker_runtime_test.go

package orbit

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"runtime"
	"strings"
	"testing"
	"time"
)

func guardLease(duration time.Duration) WorkerLease {
	return WorkerLease{Active: true, State: "RUNNING", deadline: time.Now().Add(duration)}
}
func TestExecutionGuardRenewsShortensAndNeverRevives(t *testing.T) {
	parent, cancel := context.WithCancel(context.Background())
	defer cancel()
	g, ok := newExecutionGuard(parent, guardLease(80*time.Millisecond))
	if !ok {
		t.Fatal("missing guard")
	}
	defer g.stop(context.Canceled)
	if !g.update(guardLease(500 * time.Millisecond)) {
		t.Fatal("renew rejected")
	}
	select {
	case <-g.ctx.Done():
		t.Fatal("old lease cancelled a renewed handler")
	case <-time.After(100 * time.Millisecond):
	}
	if !g.update(guardLease(35 * time.Millisecond)) {
		t.Fatal("new short session rejected")
	}
	select {
	case <-g.ctx.Done():
	case <-time.After(200 * time.Millisecond):
		t.Fatal("shorter session did not cancel handler")
	}
	if !errors.Is(context.Cause(g.ctx), ErrExecutionLeaseLost) {
		t.Fatal(context.Cause(g.ctx))
	}
	if g.update(guardLease(time.Minute)) {
		t.Fatal("expired handler revived")
	}
}

func TestHandlerOutcomesAndAbort(t *testing.T) {
	ref := WorkerAttempt{TaskID: workerTestTask, AttemptID: workerTestAttempt, Epoch: 1}
	for _, test := range []struct {
		name string
		fn   func(context.Context, WorkerExecution) (json.RawMessage, error)
		out  WorkerOutcome
		code string
	}{
		{"success", func(context.Context, WorkerExecution) (json.RawMessage, error) {
			return json.RawMessage(`{"n":9007199254740993}`), nil
		}, WorkerSucceeded, ""},
		{"known_failure", func(context.Context, WorkerExecution) (json.RawMessage, error) {
			return nil, &WorkerFailure{Code: "INPUT_REJECTED", Details: json.RawMessage(`{"field":"name"}`)}
		}, WorkerFailed, "INPUT_REJECTED"},
		{"ordinary_error", func(context.Context, WorkerExecution) (json.RawMessage, error) {
			return nil, errors.New("private error secret")
		}, WorkerUnknown, "HANDLER_ERROR"},
		{"panic", func(context.Context, WorkerExecution) (json.RawMessage, error) { panic("private panic secret") }, WorkerUnknown, "HANDLER_PANIC"},
		{"goexit", func(context.Context, WorkerExecution) (json.RawMessage, error) { runtime.Goexit(); return nil, nil }, WorkerUnknown, "HANDLER_ABORTED"},
		{"bad_json", func(context.Context, WorkerExecution) (json.RawMessage, error) {
			return json.RawMessage(`{"n":1,"n":2}`), nil
		}, WorkerUnknown, "INVALID_HANDLER_RESULT"},
		{"number_range", func(context.Context, WorkerExecution) (json.RawMessage, error) { return json.RawMessage(`1e999`), nil }, WorkerUnknown, "INVALID_HANDLER_RESULT"},
	} {
		t.Run(test.name, func(t *testing.T) {
			w := &Worker{}
			w.running.Add(1)
			done := make(chan handlerCompletion, 1)
			go w.invoke(context.Background(), WorkerHandler{Handle: test.fn}, WorkerExecution{}, ref, done, make(chan struct{}))
			select {
			case got := <-done:
				if got.result.Outcome != test.out || got.result.ErrorCode != test.code || strings.Contains(string(got.result.Result), "secret") {
					t.Fatalf("unexpected classification: %s/%s", got.result.Outcome, got.result.ErrorCode)
				}
			case <-time.After(time.Second):
				t.Fatal("handler lost result")
			}
			if w.running.Load() != 0 || w.started.Load() != 1 {
				t.Fatal("handler counters leaked")
			}
		})
	}
	for _, raw := range []string{`null`, `0.1e-308`, strings.Repeat("[", 31) + "1" + strings.Repeat("]", 31), strings.Repeat("9", 257)} {
		if validHandlerResult(json.RawMessage(raw)) {
			t.Fatal("unsupported result passed runtime validation")
		}
	}
}

func TestWorkerOptionsAndSingleRun(t *testing.T) {
	c, _ := NewClient("http://127.0.0.1:1", Options{Token: "test-token"})
	opts := WorkerOptions{QueueID: workerTestQueue, Handlers: []WorkerHandler{{Capability: WorkerCapability{Handler: "test", HandlerVersion: "v1", PayloadVersion: "v1"}, Handle: func(context.Context, WorkerExecution) (json.RawMessage, error) { return json.RawMessage(`{}`), nil }}}}
	w, err := NewWorker(c, opts)
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if err := w.Run(ctx); err != nil {
		t.Fatal(err)
	}
	if err := w.Run(ctx); !errors.Is(err, ErrWorkerAlreadyRun) {
		t.Fatal("Worker identity reused after Run")
	}
	if _, open := <-w.Events(); open {
		t.Fatal("event channel not closed")
	}
	for _, invalid := range []WorkerOptions{{QueueID: workerTestQueue}, {QueueID: workerTestQueue, Handlers: append(opts.Handlers, opts.Handlers...)}, {QueueID: workerTestQueue, Handlers: opts.Handlers, ShutdownTimeout: -1}} {
		if _, err := NewWorker(c, invalid); err == nil {
			t.Fatal("invalid Worker options accepted")
		}
	}
}

func TestWorkerShutdownBoundsPendingHandshake(t *testing.T) {
	entered := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { close(entered); <-r.Context().Done() }))
	defer server.Close()
	c, _ := NewClient(server.URL, Options{Token: "test-token"})
	w, err := NewWorker(c, WorkerOptions{QueueID: workerTestQueue, ShutdownTimeout: 50 * time.Millisecond, Handlers: []WorkerHandler{runtimeHandler("test", func(context.Context, WorkerExecution) (json.RawMessage, error) {
		t.Error("unexpected handler")
		return nil, nil
	})}})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- w.Run(ctx) }()
	select {
	case <-entered:
	case <-time.After(time.Second):
		t.Fatal("handshake did not start")
	}
	began := time.Now()
	cancel()
	select {
	case err := <-done:
		if !errors.Is(err, ErrWorkerShutdownTimeout) {
			t.Fatal(err)
		}
	case <-time.After(time.Second):
		t.Fatal("shutdown blocked on pending handshake")
	}
	if time.Since(began) > 500*time.Millisecond {
		t.Fatal("handshake exceeded shutdown bound")
	}
}

func TestWorkerPermanentAuthenticationFailureStops(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "private upstream response", http.StatusUnauthorized)
	}))
	defer server.Close()
	c, _ := NewClient(server.URL, Options{Token: "test-token"})
	w, err := NewWorker(c, WorkerOptions{QueueID: workerTestQueue, Handlers: []WorkerHandler{runtimeHandler("test", func(context.Context, WorkerExecution) (json.RawMessage, error) {
		t.Error("unauthorized handler invoked")
		return nil, nil
	})}})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	err = w.Run(ctx)
	var stopped *WorkerRunError
	var status *HTTPError
	if !errors.As(err, &stopped) || !errors.As(err, &status) || status.StatusCode != 401 || stopped.Unconfirmed != 0 || strings.Contains(err.Error(), "private") {
		t.Fatalf("unexpected authentication failure: %v", err)
	}
}