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)