Moonshark/core/runner/luarunner_test.go

package runner

import (
	"context"
	"testing"
	"time"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// Helper function to create bytecode for testing
func createTestBytecode(t *testing.T, code string) []byte {
	state := luajit.New()
	if state == nil {
		t.Fatal("Failed to create Lua state")
	}
	defer state.Close()

	bytecode, err := state.CompileBytecode(code, "test")
	if err != nil {
		t.Fatalf("Failed to compile test bytecode: %v", err)
	}

	return bytecode
}

func TestRunnerBasic(t *testing.T) {
	runner, err := NewRunner()
	if err != nil {
		t.Fatalf("Failed to create runner: %v", err)
	}
	defer runner.Close()

	bytecode := createTestBytecode(t, "return 42")

	result, err := runner.Run(bytecode, nil, "")
	if err != nil {
		t.Fatalf("Failed to run script: %v", err)
	}

	num, ok := result.(float64)
	if !ok {
		t.Fatalf("Expected float64 result, got %T", result)
	}

	if num != 42 {
		t.Errorf("Expected 42, got %f", num)
	}
}

func TestRunnerWithContext(t *testing.T) {
	runner, err := NewRunner()
	if err != nil {
		t.Fatalf("Failed to create runner: %v", err)
	}
	defer runner.Close()

	bytecode := createTestBytecode(t, `
		return {
			num = ctx.number,
			str = ctx.text,
			flag = ctx.enabled,
			list = {ctx.table[1], ctx.table[2], ctx.table[3]},
		}
	`)

	execCtx := NewContext()
	execCtx.Set("number", 42.5)
	execCtx.Set("text", "hello")
	execCtx.Set("enabled", true)
	execCtx.Set("table", []float64{10, 20, 30})

	result, err := runner.Run(bytecode, execCtx, "")
	if err != nil {
		t.Fatalf("Failed to run job: %v", err)
	}

	// Result should be a map
	resultMap, ok := result.(map[string]any)
	if !ok {
		t.Fatalf("Expected map result, got %T", result)
	}

	// Check values
	if resultMap["num"] != 42.5 {
		t.Errorf("Expected num=42.5, got %v", resultMap["num"])
	}
	if resultMap["str"] != "hello" {
		t.Errorf("Expected str=hello, got %v", resultMap["str"])
	}
	if resultMap["flag"] != true {
		t.Errorf("Expected flag=true, got %v", resultMap["flag"])
	}

	arr, ok := resultMap["list"].([]float64)
	if !ok {
		t.Fatalf("Expected []float64, got %T", resultMap["list"])
	}

	expected := []float64{10, 20, 30}
	for i, v := range expected {
		if arr[i] != v {
			t.Errorf("Expected list[%d]=%f, got %f", i, v, arr[i])
		}
	}
}

func TestRunnerWithTimeout(t *testing.T) {
	runner, err := NewRunner()
	if err != nil {
		t.Fatalf("Failed to create runner: %v", err)
	}
	defer runner.Close()

	// Create bytecode that sleeps
	bytecode := createTestBytecode(t, `
		-- Sleep for 500ms
		local start = os.time()
		while os.difftime(os.time(), start) < 0.5 do end
		return "done"
	`)

	// Test with timeout that should succeed
	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
	defer cancel()

	result, err := runner.RunWithContext(ctx, bytecode, nil, "")
	if err != nil {
		t.Fatalf("Unexpected error with sufficient timeout: %v", err)
	}
	if result != "done" {
		t.Errorf("Expected 'done', got %v", result)
	}

	// Test with timeout that should fail
	ctx, cancel = context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()

	_, err = runner.RunWithContext(ctx, bytecode, nil, "")
	if err == nil {
		t.Errorf("Expected timeout error, got nil")
	}
}

func TestSandboxIsolation(t *testing.T) {
	runner, err := NewRunner()
	if err != nil {
		t.Fatalf("Failed to create runner: %v", err)
	}
	defer runner.Close()

	// Create a script that tries to modify a global variable
	bytecode1 := createTestBytecode(t, `
		-- Set a "global" variable
		my_global = "test value"
		return true
	`)

	_, err = runner.Run(bytecode1, nil, "")
	if err != nil {
		t.Fatalf("Failed to execute first script: %v", err)
	}

	// Now try to access that variable from another script
	bytecode2 := createTestBytecode(t, `
		-- Try to access the previously set global
		return my_global ~= nil
	`)

	result, err := runner.Run(bytecode2, nil, "")
	if err != nil {
		t.Fatalf("Failed to execute second script: %v", err)
	}

	// The variable should not be accessible (sandbox isolation)
	if result.(bool) {
		t.Errorf("Expected sandbox isolation, but global variable was accessible")
	}
}

func TestRunnerWithInit(t *testing.T) {
	// Define an init function that registers a simple "math" module
	mathInit := func(state *luajit.State) error {
		// Register the "add" function
		err := state.RegisterGoFunction("add", func(s *luajit.State) int {
			a := s.ToNumber(1)
			b := s.ToNumber(2)
			s.PushNumber(a + b)
			return 1 // Return one result
		})

		if err != nil {
			return err
		}

		// Register a whole module
		mathFuncs := map[string]luajit.GoFunction{
			"multiply": func(s *luajit.State) int {
				a := s.ToNumber(1)
				b := s.ToNumber(2)
				s.PushNumber(a * b)
				return 1
			},
			"subtract": func(s *luajit.State) int {
				a := s.ToNumber(1)
				b := s.ToNumber(2)
				s.PushNumber(a - b)
				return 1
			},
		}

		return RegisterModule(state, "math2", mathFuncs)
	}

	// Create a runner with our init function
	runner, err := NewRunner(WithInitFunc(mathInit))
	if err != nil {
		t.Fatalf("Failed to create runner: %v", err)
	}
	defer runner.Close()

	// Test the add function
	bytecode1 := createTestBytecode(t, "return add(5, 7)")
	result1, err := runner.Run(bytecode1, nil, "")
	if err != nil {
		t.Fatalf("Failed to call add function: %v", err)
	}

	num1, ok := result1.(float64)
	if !ok || num1 != 12 {
		t.Errorf("Expected add(5, 7) = 12, got %v", result1)
	}

	// Test the math2 module
	bytecode2 := createTestBytecode(t, "return math2.multiply(6, 8)")
	result2, err := runner.Run(bytecode2, nil, "")
	if err != nil {
		t.Fatalf("Failed to call math2.multiply: %v", err)
	}

	num2, ok := result2.(float64)
	if !ok || num2 != 48 {
		t.Errorf("Expected math2.multiply(6, 8) = 48, got %v", result2)
	}
}

func TestConcurrentExecution(t *testing.T) {
	const jobs = 20

	runner, err := NewRunner(WithBufferSize(20))
	if err != nil {
		t.Fatalf("Failed to create runner: %v", err)
	}
	defer runner.Close()

	// Create bytecode that returns its input
	bytecode := createTestBytecode(t, "return ctx.n")

	// Run multiple jobs concurrently
	results := make(chan int, jobs)
	for i := 0; i < jobs; i++ {
		i := i // Capture loop variable
		go func() {
			execCtx := NewContext()
			execCtx.Set("n", float64(i))

			result, err := runner.Run(bytecode, execCtx, "")
			if err != nil {
				t.Errorf("Job %d failed: %v", i, err)
				results <- -1
				return
			}

			num, ok := result.(float64)
			if !ok {
				t.Errorf("Job %d: expected float64, got %T", i, result)
				results <- -1
				return
			}

			results <- int(num)
		}()
	}

	// Collect results
	seen := make(map[int]bool)
	for i := 0; i < jobs; i++ {
		result := <-results
		if result != -1 {
			seen[result] = true
		}
	}

	// Verify all jobs were processed
	if len(seen) != jobs {
		t.Errorf("Expected %d unique results, got %d", jobs, len(seen))
	}
}

func TestRunnerClose(t *testing.T) {
	runner, err := NewRunner()
	if err != nil {
		t.Fatalf("Failed to create runner: %v", err)
	}

	// Submit a job to verify runner works
	bytecode := createTestBytecode(t, "return 42")
	_, err = runner.Run(bytecode, nil, "")
	if err != nil {
		t.Fatalf("Failed to run job: %v", err)
	}

	// Close
	if err := runner.Close(); err != nil {
		t.Errorf("Close failed: %v", err)
	}

	// Run after close should fail
	_, err = runner.Run(bytecode, nil, "")
	if err != ErrRunnerClosed {
		t.Errorf("Expected ErrRunnerClosed, got %v", err)
	}

	// Second close should return error
	if err := runner.Close(); err != ErrRunnerClosed {
		t.Errorf("Expected ErrRunnerClosed on second close, got %v", err)
	}
}

func TestErrorHandling(t *testing.T) {
	runner, err := NewRunner()
	if err != nil {
		t.Fatalf("Failed to create runner: %v", err)
	}
	defer runner.Close()

	// Test invalid bytecode
	_, err = runner.Run([]byte("not valid bytecode"), nil, "")
	if err == nil {
		t.Errorf("Expected error for invalid bytecode, got nil")
	}

	// Test Lua runtime error
	bytecode := createTestBytecode(t, `
		error("intentional error")
		return true
	`)

	_, err = runner.Run(bytecode, nil, "")
	if err == nil {
		t.Errorf("Expected error from Lua error() call, got nil")
	}

	// Test with nil context
	bytecode = createTestBytecode(t, "return ctx == nil")
	result, err := runner.Run(bytecode, nil, "")
	if err != nil {
		t.Errorf("Unexpected error with nil context: %v", err)
	}
	if result.(bool) != true {
		t.Errorf("Expected ctx to be nil in Lua, but it wasn't")
	}

	// Test invalid context value
	execCtx := NewContext()
	execCtx.Set("param", complex128(1+2i)) // Unsupported type

	bytecode = createTestBytecode(t, "return ctx.param")
	_, err = runner.Run(bytecode, execCtx, "")
	if err == nil {
		t.Errorf("Expected error for unsupported context value type, got nil")
	}
}