core/server/sandbox/batch_module_test.go

package sandbox

import (
	"context"
	"strings"
	"testing"
	"time"

	"go.starlark.net/starlark"
)

// execWithBatch runs a Starlark script with the batch module available.
func execWithBatch(t *testing.T, script string) (*Result, error) {
	t.Helper()
	ctx := context.Background()
	batchMod := BuildBatchModule(ctx, nil, "test-pkg", nil, nil, nil)
	sb := New(DefaultConfig())
	return sb.Exec(ctx, "test.star", script, map[string]starlark.Value{
		"batch": batchMod,
	})
}

// ─── Parallel Ordering ──────────────────────

func TestBatchExec_ParallelOrdering(t *testing.T) {
	result, err := execWithBatch(t, `
results, errors = batch.exec([
    lambda: 10,
    lambda: 20,
    lambda: 30,
])
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	// Verify results are in input order.
	globals := result.Globals
	results := globals["results"].(*starlark.List)
	errors := globals["errors"].(*starlark.List)

	if results.Len() != 3 {
		t.Fatalf("expected 3 results, got %d", results.Len())
	}

	for i, want := range []int{10, 20, 30} {
		v, _ := starlark.AsInt32(results.Index(i))
		got := int(v)
		if got != want {
			t.Errorf("results[%d] = %d, want %d", i, got, want)
		}
		if errors.Index(i) != starlark.None {
			t.Errorf("errors[%d] = %v, want None", i, errors.Index(i))
		}
	}
}

// ─── Partial Failure ────────────────────────

func TestBatchExec_PartialFailure(t *testing.T) {
	result, err := execWithBatch(t, `
def ok_a():
    return "a"

def bad():
    return 1 // 0

def ok_c():
    return "c"

results, errors = batch.exec([ok_a, bad, ok_c])
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	results := result.Globals["results"].(*starlark.List)
	errors := result.Globals["errors"].(*starlark.List)

	// results[0] = "a", results[2] = "c"
	if s, _ := starlark.AsString(results.Index(0)); s != "a" {
		t.Errorf("results[0] = %v, want 'a'", results.Index(0))
	}
	if s, _ := starlark.AsString(results.Index(2)); s != "c" {
		t.Errorf("results[2] = %v, want 'c'", results.Index(2))
	}

	// results[1] = None (failed)
	if results.Index(1) != starlark.None {
		t.Errorf("results[1] = %v, want None", results.Index(1))
	}

	// errors[0] and errors[2] = None (succeeded)
	if errors.Index(0) != starlark.None {
		t.Errorf("errors[0] = %v, want None", errors.Index(0))
	}
	if errors.Index(2) != starlark.None {
		t.Errorf("errors[2] = %v, want None", errors.Index(2))
	}

	// errors[1] should be a string containing the error
	errStr, ok := starlark.AsString(errors.Index(1))
	if !ok || errStr == "" {
		t.Errorf("errors[1] = %v, want non-empty error string", errors.Index(1))
	}
}

// ─── Timeout Per-Branch ─────────────────────

func TestBatchExec_TimeoutPerBranch(t *testing.T) {
	// Use a tight timeout. The slow lambda burns steps in an infinite loop,
	// which will hit either the step limit or the context timeout.
	result, err := execWithBatch(t, `
def fast():
    return "done"

def slow():
    x = 0
    for i in range(999999999):
        x = x + 1
    return x

results, errors = batch.exec([fast, slow], timeout=1)
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	results := result.Globals["results"].(*starlark.List)
	errors := result.Globals["errors"].(*starlark.List)

	// Fast branch should succeed.
	if s, _ := starlark.AsString(results.Index(0)); s != "done" {
		t.Errorf("results[0] = %v, want 'done'", results.Index(0))
	}
	if errors.Index(0) != starlark.None {
		t.Errorf("errors[0] = %v, want None", errors.Index(0))
	}

	// Slow branch should fail (step limit or timeout).
	if results.Index(1) != starlark.None {
		t.Errorf("results[1] = %v, want None", results.Index(1))
	}
	errStr, _ := starlark.AsString(errors.Index(1))
	if errStr == "" {
		t.Errorf("errors[1] should contain error string")
	}
}

// ─── Parent Cancellation ────────────────────

func TestBatchExec_ParentCancellation(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel() // cancel immediately

	batchMod := BuildBatchModule(ctx, nil, "test-pkg", nil, nil, nil)
	sb := New(DefaultConfig())
	result, err := sb.Exec(ctx, "test.star", `
results, errors = batch.exec([lambda: 1, lambda: 2])
`, map[string]starlark.Value{"batch": batchMod})

	// Either the Exec itself fails (context cancelled before script runs)
	// or the batch.exec branches all fail.
	if err != nil {
		// Script-level cancellation — acceptable.
		return
	}

	errors := result.Globals["errors"].(*starlark.List)
	allErrored := true
	for i := 0; i < errors.Len(); i++ {
		if errors.Index(i) == starlark.None {
			allErrored = false
		}
	}
	if !allErrored {
		t.Errorf("expected all branches to error on cancelled context, errors=%v", errors)
	}
}

// ─── Cap Enforcement ────────────────────────

func TestBatchExec_CapEnforcement(t *testing.T) {
	_, err := execWithBatch(t, `
fns = [lambda: i for i in range(9)]
batch.exec(fns)
`)
	if err == nil {
		t.Fatal("expected error for 9 callables, got nil")
	}
	if !strings.Contains(err.Error(), "max 8") {
		t.Errorf("error = %q, want to contain 'max 8'", err.Error())
	}
}

// ─── Empty List ─────────────────────────────

func TestBatchExec_EmptyList(t *testing.T) {
	_, err := execWithBatch(t, `batch.exec([])`)
	if err == nil {
		t.Fatal("expected error for empty list, got nil")
	}
	if !strings.Contains(err.Error(), "empty") {
		t.Errorf("error = %q, want to contain 'empty'", err.Error())
	}
}

// ─── Non-Callable Element ───────────────────

func TestBatchExec_NonCallableElement(t *testing.T) {
	_, err := execWithBatch(t, `batch.exec([1, 2])`)
	if err == nil {
		t.Fatal("expected error for non-callable, got nil")
	}
	if !strings.Contains(err.Error(), "not callable") {
		t.Errorf("error = %q, want to contain 'not callable'", err.Error())
	}
}

// ─── Permission Gating ──────────────────────

func TestBatchExec_PermissionGating(t *testing.T) {
	sb := New(DefaultConfig())
	_, err := sb.Exec(context.Background(), "test.star",
		`batch.exec([lambda: 1])`,
		map[string]starlark.Value{}, // no batch module
	)
	if err == nil {
		t.Fatal("expected error when batch module not available")
	}
	if !strings.Contains(err.Error(), "batch") {
		t.Errorf("error = %q, want to mention 'batch'", err.Error())
	}
}

// ─── Frozen Struct Sharing ──────────────────

func TestBatchExec_FrozenStructSharing(t *testing.T) {
	// Two branches read the same frozen tuple. Should not race.
	result, err := execWithBatch(t, `
data = (1, 2, 3)  # tuples are frozen/immutable

results, errors = batch.exec([
    lambda: data[0] + data[1],
    lambda: data[1] + data[2],
])
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	results := result.Globals["results"].(*starlark.List)

	got0, _ := starlark.AsInt32(results.Index(0))
	got1, _ := starlark.AsInt32(results.Index(1))
	if got0 != 3 {
		t.Errorf("results[0] = %d, want 3", got0)
	}
	if got1 != 5 {
		t.Errorf("results[1] = %d, want 5", got1)
	}
}

// ─── Nested batch.exec ──────────────────────

func TestBatchExec_NestedBatchExec(t *testing.T) {
	result, err := execWithBatch(t, `
results, errors = batch.exec([
    lambda: batch.exec([lambda: 1]),
])
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	// The inner batch.exec should fail with nesting error.
	errors := result.Globals["errors"].(*starlark.List)
	errStr, _ := starlark.AsString(errors.Index(0))
	if !strings.Contains(errStr, "nested") {
		t.Errorf("errors[0] = %q, want to contain 'nested'", errStr)
	}
}

// ─── Default Timeout ────────────────────────

func TestBatchExec_DefaultTimeout(t *testing.T) {
	// Call without explicit timeout — should succeed.
	result, err := execWithBatch(t, `
results, errors = batch.exec([lambda: 42])
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	results := result.Globals["results"].(*starlark.List)
	got, _ := starlark.AsInt32(results.Index(0))
	if got != 42 {
		t.Errorf("results[0] = %d, want 42", got)
	}
}

// ─── Timeout Clamping ───────────────────────

func TestBatchExec_TimeoutClamp(t *testing.T) {
	// timeout=0 should be clamped to default, not crash.
	start := time.Now()
	result, err := execWithBatch(t, `
results, errors = batch.exec([lambda: "ok"], timeout=0)
`)
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	results := result.Globals["results"].(*starlark.List)
	if s, _ := starlark.AsString(results.Index(0)); s != "ok" {
		t.Errorf("results[0] = %v, want 'ok'", results.Index(0))
	}

	// Should complete quickly (clamped to default 10s, but lambda is instant).
	if elapsed > 2*time.Second {
		t.Errorf("took %v, expected fast completion", elapsed)
	}
}