Merge pull request #217 from vladimirvivien/preload-script-modules

Support for script module preload

exec/executor.go

@@ -4,6 +4,7 @@
 package exec
 
 import (
+	"fmt"
 	"io"
 	"os"
 
@@ -36,3 +37,35 @@ func Execute(name string, source io.Reader, args ArgMap) error {
 func ExecuteFile(file *os.File, args ArgMap) error {
 	return Execute(file.Name(), file, args)
 }
+
+type StarlarkModule struct {
+	Name   string
+	Source io.Reader
+}
+
+func ExecuteWithModules(name string, source io.Reader, args ArgMap, modules ...StarlarkModule) error {
+	star := starlark.New()
+
+	if args != nil {
+		starStruct, err := starlark.NewGoValue(args).ToStarlarkStruct("args")
+		if err != nil {
+			return err
+		}
+
+		star.AddPredeclared("args", starStruct)
+	}
+
+	// load modules
+	for _, mod := range modules {
+		if err := star.Preload(mod.Name, mod.Source); err != nil {
+			return fmt.Errorf("module load: %w", err)
+		}
+	}
+
+	err := star.Exec(name, source)
+	if err != nil {
+		return fmt.Errorf("exec failed: %w", err)
+	}
+
+	return nil
+}

exec/executor_test.go

@@ -17,7 +17,7 @@ var (
 	support *testcrashd.TestSupport
 )
 
-func TestMain(m *testing.M) {
+func setupTestSupport() {
 	test, err := testcrashd.Init()
 	if err != nil {
 		logrus.Fatal(err)
@@ -36,17 +36,17 @@ func TestMain(m *testing.M) {
 	if err != nil {
 		logrus.Fatal(err)
 	}
+}
 
-	result := m.Run()
-
+func teardownTestSupport() {
 	if err := support.TearDown(); err != nil {
 		logrus.Fatal(err)
 	}
-
-	os.Exit(result)
 }
 
-func TestKindScript(t *testing.T) {
+func TestExampleScripts(t *testing.T) {
+	setupTestSupport()
+
 	tests := []struct {
 		name       string
 		scriptPath string
@@ -80,16 +80,6 @@ func TestKindScript(t *testing.T) {
 				"ssh_port":    support.PortValue(),
 			},
 		},
-		//{
-		//	name:       "kube-nodes provider",
-		//	scriptPath: "../examples/kube-nodes-provider.crsh",
-		//	args: ArgMap{
-		//		"kubecfg":  getTestKubeConf(),
-		//		"ssh_port": testSSHPort,
-		//		"username": getUsername(),
-		//		"key_path": getPrivateKey(),
-		//	},
-		//},
 		{
 			name:       "kind-capi-bootstrap",
 			scriptPath: "../examples/kind-capi-bootstrap.crsh",
@@ -109,6 +99,7 @@ func TestKindScript(t *testing.T) {
 			}
 		})
 	}
+	teardownTestSupport()
 }
 
 func TestExecute(t *testing.T) {
@@ -118,14 +109,31 @@ func TestExecute(t *testing.T) {
 		exec   func(t *testing.T, script string)
 	}{
 		{
-			name:   "run_local",
+			name:   "execute single script",
 			script: `result = run_local("echo 'Hello World!'")`,
 			exec: func(t *testing.T, script string) {
 				if err := Execute("run_local", strings.NewReader(script), ArgMap{}); err != nil {
 					t.Fatal(err)
 				}
 			},
 		},
+		{
+			name:   "execute with modules",
+			script: `result = multiply(2, 3)`,
+			exec: func(t *testing.T, script string) {
+				mod := `
+def multiply(x, y):
+    log (msg="{} * {} = {}".format(x,y,x*y))
+`
+				if err := ExecuteWithModules(
+					"multiply",
+					strings.NewReader(script),
+					ArgMap{},
+					StarlarkModule{Name: "lib", Source: strings.NewReader(mod)}); err != nil {
+					t.Fatal(err)
+				}
+			},
+		},
 	}
 
 	for _, test := range tests {

starlark/starlark_exec.go

@@ -34,6 +34,27 @@ func (e *Executor) AddPredeclared(name string, value starlark.Value) {
 	}
 }
 
+// Preload loads parse and load code modules to be used in other scripts.
+// This call should take place prior to calling the Exec call for main script execution.
+// The result of the loaded module are added to the global predeclared
+// values used in Exec call.
+func (e *Executor) Preload(name string, source io.Reader) error {
+	result, err := starlark.ExecFile(e.thread, name, source, e.predecs)
+	if err != nil {
+		if evalErr, ok := err.(*starlark.EvalError); ok {
+			return fmt.Errorf(evalErr.Backtrace())
+		}
+		return err
+	}
+
+	// add result to predeclared
+	for k, v := range result {
+		e.predecs[k] = v
+	}
+
+	return nil
+}
+
 func (e *Executor) Exec(name string, source io.Reader) error {
 	if err := setupLocalDefaults(e.thread); err != nil {
 		return fmt.Errorf("failed to setup defaults: %s", err)