feat: Allow builtin manipulation for programmatic use

core/server/api_container/server/startosis_engine/startosis_interpreter.go

@@ -63,11 +63,22 @@ type StartosisInterpreter struct {
 	starlarkValueSerde           *kurtosis_types.StarlarkValueSerde
 	enclaveEnvVars               string
 	interpretationTimeValueStore *interpretation_time_value_store.InterpretationTimeValueStore
+	// This is a function that allows the consumer of the interpreter to adjust the default builtins.
+	// It is useful when external libraries or helpers need to be plugged in to kurtosis,
+	// for example when running unit tests using the starlarktest package
+	processBuiltins StartosisInterpreterBuiltinsProcessor
 }
 
+// StartosisInterpreterBuiltinsProcessor is a builtins transformer function
+type StartosisInterpreterBuiltinsProcessor func(thread *starlark.Thread, predeclared starlark.StringDict) starlark.StringDict
+
 type SerializedInterpretationOutput string
 
 func NewStartosisInterpreter(serviceNetwork service_network.ServiceNetwork, packageContentProvider startosis_packages.PackageContentProvider, runtimeValueStore *runtime_value_store.RuntimeValueStore, starlarkValueSerde *kurtosis_types.StarlarkValueSerde, enclaveVarEnvs string, interpretationTimeValueStore *interpretation_time_value_store.InterpretationTimeValueStore) *StartosisInterpreter {
+	return NewStartosisInterpreterWithBuiltinsProcessor(serviceNetwork, packageContentProvider, runtimeValueStore, starlarkValueSerde, enclaveVarEnvs, interpretationTimeValueStore, nil)
+}
+
+func NewStartosisInterpreterWithBuiltinsProcessor(serviceNetwork service_network.ServiceNetwork, packageContentProvider startosis_packages.PackageContentProvider, runtimeValueStore *runtime_value_store.RuntimeValueStore, starlarkValueSerde *kurtosis_types.StarlarkValueSerde, enclaveVarEnvs string, interpretationTimeValueStore *interpretation_time_value_store.InterpretationTimeValueStore, processBuiltins StartosisInterpreterBuiltinsProcessor) *StartosisInterpreter {
 	return &StartosisInterpreter{
 		mutex:                        &sync.Mutex{},
 		serviceNetwork:               serviceNetwork,
@@ -76,6 +87,7 @@ func NewStartosisInterpreter(serviceNetwork service_network.ServiceNetwork, pack
 		enclaveEnvVars:               enclaveVarEnvs,
 		starlarkValueSerde:           starlarkValueSerde,
 		interpretationTimeValueStore: interpretationTimeValueStore,
+		processBuiltins:              processBuiltins,
 	}
 }
 
@@ -389,6 +401,15 @@ func (interpreter *StartosisInterpreter) buildBindings(
 	for _, kurtosisTypeConstructors := range KurtosisTypeConstructors() {
 		predeclared[kurtosisTypeConstructors.Name()] = kurtosisTypeConstructors
 	}
+
+	// Allow the consumers to adjust the builtins
+	//
+	// This is useful for adding e.g. starlarktest package
+	// for unit testing of kurtosis scripts
+	if interpreter.processBuiltins != nil {
+		predeclared = interpreter.processBuiltins(thread, predeclared)
+	}
+
 	return &predeclared, nil
 }
 

core/server/api_container/server/startosis_engine/startosis_interpreter_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/kurtosis-tech/kurtosis/core/server/api_container/server/startosis_engine/startosis_packages/mock_package_content_provider"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
+	"go.starlark.net/starlark"
 	"net"
 	"strings"
 	"testing"
@@ -1081,6 +1082,57 @@ def run(plan):
 	require.Nil(suite.T(), interpretationError)
 }
 
+func (suite *StartosisInterpreterTestSuite) TestStarlarkInterpreter_WithExtraPredefinedThatFails() {
+	script := `
+def run(plan):
+	yell()
+`
+
+	expectedError := startosis_errors.NewInterpretationError("do not yell() please")
+
+	processBuiltins := func(thread *starlark.Thread, predeclareds starlark.StringDict) starlark.StringDict {
+		// We'll add a new builtin called raise that should raise an evaluation error
+		predeclareds["yell"] = starlark.NewBuiltin("yell", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+			return nil, expectedError
+		})
+
+		return predeclareds
+	}
+
+	// We'll create a new interpreter using the processBuiltins transformer above
+	interpreter := NewStartosisInterpreterWithBuiltinsProcessor(suite.serviceNetwork, suite.packageContentProvider, suite.runtimeValueStore, nil, "", suite.interpretationTimeValueStore, processBuiltins)
+
+	_, _, interpretationError := interpreter.Interpret(context.Background(), startosis_constants.PackageIdPlaceholderForStandaloneScript, useDefaultMainFunctionName, noPackageReplaceOptions, startosis_constants.PlaceHolderMainFileForPlaceStandAloneScript, script, startosis_constants.EmptyInputArgs, defaultNonBlockingMode, emptyEnclaveComponents, emptyInstructionsPlanMask, defaultImageDownloadMode)
+
+	// If everything goes well we should see an error with a correct message & stack trace
+	require.NotNil(suite.T(), interpretationError)
+	require.Equal(suite.T(), fmt.Sprintf("Evaluation error: %v\n\tat [3:6]: run\n\tat [0:0]: yell", expectedError.Error()), interpretationError.GetErrorMessage())
+
+}
+
+func (suite *StartosisInterpreterTestSuite) TestStarlarkInterpreter_WithExtraPredefinedThatReturns() {
+	script := `
+def run(plan):
+	plan.print(chill())
+`
+	processBuiltins := func(thread *starlark.Thread, predeclareds starlark.StringDict) starlark.StringDict {
+		// We create a builtin that simply returns a value
+		predeclareds["chill"] = starlark.NewBuiltin("chill", func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+			return starlark.String("chillin"), nil
+		})
+
+		return predeclareds
+	}
+
+	// We'll create a new interpreter using the processBuiltins transformer above
+	interpreter := NewStartosisInterpreterWithBuiltinsProcessor(suite.serviceNetwork, suite.packageContentProvider, suite.runtimeValueStore, nil, "", suite.interpretationTimeValueStore, processBuiltins)
+
+	// If everything goes well we should see no error & message on the output
+	_, instructionsPlan, interpretationError := interpreter.Interpret(context.Background(), startosis_constants.PackageIdPlaceholderForStandaloneScript, useDefaultMainFunctionName, noPackageReplaceOptions, startosis_constants.PlaceHolderMainFileForPlaceStandAloneScript, script, startosis_constants.EmptyInputArgs, defaultNonBlockingMode, emptyEnclaveComponents, emptyInstructionsPlanMask, defaultImageDownloadMode)
+	require.Nil(suite.T(), interpretationError)
+	validateScriptOutputFromPrintInstructions(suite.T(), instructionsPlan, "chillin\n")
+}
+
 func (suite *StartosisInterpreterTestSuite) TestGetModuleIdPrefix() {
 	githubModuleId := "github.com/some-person/some-pkg/main.star"
 	expectedGithubModuleId := "github.com/some-person/some-pkg"