core/server/sandbox/modules_test.go

package sandbox

import (
	"context"
	"testing"

	"go.starlark.net/starlark"
)

func TestBuildSecretsModule_Get(t *testing.T) {
	mod := MakeModule("secrets", starlark.StringDict{
		"get": starlark.NewBuiltin("secrets.get", func(
			thread *starlark.Thread, b *starlark.Builtin,
			args starlark.Tuple, kwargs []starlark.Tuple,
		) (starlark.Value, error) {
			var key string
			if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &key); err != nil {
				return nil, err
			}
			if key == "api_key" {
				return starlark.String("sk-test123"), nil
			}
			return starlark.None, nil
		}),
	})

	sb := New(DefaultConfig())
	res, err := sb.Exec(context.Background(), "test.star", `
val = secrets.get("api_key")
missing = secrets.get("nope")
`, map[string]starlark.Value{"secrets": mod})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	val := res.Globals["val"]
	if val.(starlark.String).GoString() != "sk-test123" {
		t.Fatalf("expected 'sk-test123', got %s", val.String())
	}
	missing := res.Globals["missing"]
	if missing != starlark.None {
		t.Fatalf("expected None, got %s", missing.String())
	}
}

func TestBuildNotificationsModule_Send(t *testing.T) {
	var sentTitle string
	var sentUserID string

	sendFn := starlark.NewBuiltin("notifications.send", func(
		thread *starlark.Thread, b *starlark.Builtin,
		args starlark.Tuple, kwargs []starlark.Tuple,
	) (starlark.Value, error) {
		var userID, title string
		if err := starlark.UnpackArgs(b.Name(), args, kwargs,
			"user_id", &userID, "title", &title); err != nil {
			return nil, err
		}
		sentUserID = userID
		sentTitle = title
		return starlark.True, nil
	})

	mod := MakeModule("notifications", starlark.StringDict{
		"send": sendFn,
	})

	sb := New(DefaultConfig())
	_, err := sb.Exec(context.Background(), "test.star", `
result = notifications.send(user_id="u-1", title="Hello World")
`, map[string]starlark.Value{"notifications": mod})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if sentUserID != "u-1" {
		t.Fatalf("expected user_id 'u-1', got %q", sentUserID)
	}
	if sentTitle != "Hello World" {
		t.Fatalf("expected title 'Hello World', got %q", sentTitle)
	}
}

func TestModuleIsolation(t *testing.T) {
	// Script should not be able to access modules it wasn't given
	sb := New(DefaultConfig())
	_, err := sb.Exec(context.Background(), "test.star", `
val = secrets.get("key")
`, nil) // no modules injected

	if err == nil {
		t.Fatal("expected error when accessing undefined module")
	}
}

func TestMakeModuleName(t *testing.T) {
	mod := MakeModule("test_mod", starlark.StringDict{
		"version": starlark.String("1.0"),
	})
	if mod.Name != "test_mod" {
		t.Fatalf("expected 'test_mod', got %q", mod.Name)
	}
}

func TestEntryPointCall(t *testing.T) {
	sb := New(DefaultConfig())

	// Simulate a package script with on_pre_completion entry point
	res, err := sb.Exec(context.Background(), "ext.star", `
def on_pre_completion(ctx):
    channel = ctx["channel_id"]
    return [{"role": "system", "content": "context for " + channel}]
`, nil)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	fn := res.Globals["on_pre_completion"].(starlark.Callable)

	ctxDict := starlark.NewDict(1)
	ctxDict.SetKey(starlark.String("channel_id"), starlark.String("ch-abc"))

	val, _, err := sb.Call(context.Background(), fn, starlark.Tuple{ctxDict}, nil)
	if err != nil {
		t.Fatalf("call error: %v", err)
	}

	list, ok := val.(*starlark.List)
	if !ok {
		t.Fatalf("expected list, got %s", val.Type())
	}
	if list.Len() != 1 {
		t.Fatalf("expected 1 message, got %d", list.Len())
	}

	item := list.Index(0)
	dict, ok := item.(*starlark.Dict)
	if !ok {
		t.Fatalf("expected dict, got %s", item.Type())
	}

	contentVal, _, _ := dict.Get(starlark.String("content"))
	content, _ := starlark.AsString(contentVal)
	if content != "context for ch-abc" {
		t.Fatalf("expected 'context for ch-abc', got %q", content)
	}
}

func TestEntryPointNoneReturn(t *testing.T) {
	sb := New(DefaultConfig())

	res, err := sb.Exec(context.Background(), "ext.star", `
def on_pre_completion(ctx):
    return None
`, nil)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	fn := res.Globals["on_pre_completion"].(starlark.Callable)
	ctxDict := starlark.NewDict(0)

	val, _, err := sb.Call(context.Background(), fn, starlark.Tuple{ctxDict}, nil)
	if err != nil {
		t.Fatalf("call error: %v", err)
	}

	if val != starlark.None {
		t.Fatalf("expected None, got %s", val.Type())
	}
}