Changeset 0.29.0 (#195)

server/sandbox/modules_test.go

@@ -0,0 +1,171 @@
+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())
+	}
+}