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()) } }