From be5c0231436c0136a3749465256bd0010d12a175 Mon Sep 17 00:00:00 2001 From: Jan Nanista Date: Fri, 31 Jan 2025 16:16:28 -0800 Subject: [PATCH] feat: Allow builtin manipulation for programmatic use --- .../startosis_engine/startosis_interpreter.go | 21 ++++++++ .../startosis_interpreter_test.go | 52 +++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/core/server/api_container/server/startosis_engine/startosis_interpreter.go b/core/server/api_container/server/startosis_engine/startosis_interpreter.go index 99b3e9a5a7..1880b13a48 100644 --- a/core/server/api_container/server/startosis_engine/startosis_interpreter.go +++ b/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 } diff --git a/core/server/api_container/server/startosis_engine/startosis_interpreter_test.go b/core/server/api_container/server/startosis_engine/startosis_interpreter_test.go index 2b95563878..7472ad6078 100644 --- a/core/server/api_container/server/startosis_engine/startosis_interpreter_test.go +++ b/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"