From e0886ddb7d72802a8ef0d44c0a115f502e47ad4f Mon Sep 17 00:00:00 2001 From: Andrew Nester Date: Wed, 7 Feb 2024 10:49:03 +0100 Subject: [PATCH 1/2] Make sure grouped flags are added to the command flag set --- bundle/run/options.go | 2 ++ libs/cmdgroup/command.go | 28 ++++++++++++++++++++++++++++ libs/cmdgroup/command_test.go | 7 +++++++ libs/cmdgroup/template.go | 4 ++-- 4 files changed, 39 insertions(+), 2 deletions(-) diff --git a/bundle/run/options.go b/bundle/run/options.go index 580612d0ea..1f47712e23 100644 --- a/bundle/run/options.go +++ b/bundle/run/options.go @@ -23,4 +23,6 @@ func (o *Options) Define(cmd *cobra.Command) { pipelineGroup := wrappedCmd.AddFlagGroup("Pipeline") o.Pipeline.Define(pipelineGroup.FlagSet()) + + wrappedCmd.RefreshFlags() } diff --git a/libs/cmdgroup/command.go b/libs/cmdgroup/command.go index 19c9af16af..ad3ba6efa2 100644 --- a/libs/cmdgroup/command.go +++ b/libs/cmdgroup/command.go @@ -15,6 +15,12 @@ type CommandWithGroupFlag struct { flagGroups []*FlagGroup } +func (c *CommandWithGroupFlag) RefreshFlags() { + for _, fg := range c.flagGroups { + c.cmd.Flags().AddFlagSet(fg.flagSet) + } +} + func (c *CommandWithGroupFlag) Command() *cobra.Command { return c.cmd } @@ -23,6 +29,24 @@ func (c *CommandWithGroupFlag) FlagGroups() []*FlagGroup { return c.flagGroups } +func (c *CommandWithGroupFlag) NonGroupedFlags() *pflag.FlagSet { + nonGrouped := pflag.NewFlagSet("non-grouped", pflag.ContinueOnError) + c.cmd.LocalFlags().VisitAll(func(f *pflag.Flag) { + for _, fg := range c.flagGroups { + if fg.Has(f) { + return + } + } + nonGrouped.AddFlag(f) + }) + + return nonGrouped +} + +func (c *CommandWithGroupFlag) HasNonGroupedFlags() bool { + return c.NonGroupedFlags().HasFlags() +} + func NewCommandWithGroupFlag(cmd *cobra.Command) *CommandWithGroupFlag { cmdWithFlagGroups := &CommandWithGroupFlag{cmd: cmd, flagGroups: make([]*FlagGroup, 0)} cmd.SetUsageFunc(func(c *cobra.Command) error { @@ -64,6 +88,10 @@ func (c *FlagGroup) FlagSet() *pflag.FlagSet { return c.flagSet } +func (c *FlagGroup) Has(f *pflag.Flag) bool { + return c.flagSet.Lookup(f.Name) != nil +} + var templateFuncs = template.FuncMap{ "trim": strings.TrimSpace, "trimRightSpace": trimRightSpace, diff --git a/libs/cmdgroup/command_test.go b/libs/cmdgroup/command_test.go index 2eae31d14e..ea33e66650 100644 --- a/libs/cmdgroup/command_test.go +++ b/libs/cmdgroup/command_test.go @@ -29,6 +29,7 @@ func TestCommandFlagGrouping(t *testing.T) { fs.String("pipeline-type", "", "Type of the pipeline") cmd.Flags().BoolP("bool", "b", false, "Bool flag") + wrappedCmd.RefreshFlags() buf := bytes.NewBuffer(nil) cmd.SetOutput(buf) @@ -48,4 +49,10 @@ Pipeline Flags: Flags: -b, --bool Bool flag` require.Equal(t, expected, buf.String()) + + require.NotNil(t, cmd.Flags().Lookup("job-name")) + require.NotNil(t, cmd.Flags().Lookup("job-type")) + require.NotNil(t, cmd.Flags().Lookup("pipeline-name")) + require.NotNil(t, cmd.Flags().Lookup("pipeline-type")) + require.NotNil(t, cmd.Flags().Lookup("bool")) } diff --git a/libs/cmdgroup/template.go b/libs/cmdgroup/template.go index aac967b0ee..5c1be48fb4 100644 --- a/libs/cmdgroup/template.go +++ b/libs/cmdgroup/template.go @@ -7,8 +7,8 @@ const usageTemplate = `Usage:{{if .Command.Runnable}} {{.Description}}{{end}} {{.FlagSet.FlagUsages | trimTrailingWhitespaces}} {{end}} -{{if .Command.HasAvailableLocalFlags}}Flags: -{{.Command.LocalFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}{{if .Command.HasAvailableInheritedFlags}} +{{if .HasNonGroupedFlags}}Flags: +{{.NonGroupedFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}{{if .Command.HasAvailableInheritedFlags}} Global Flags: {{.Command.InheritedFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}` From 8d965afe4cdf1866c8632df1975afd8cde2a32fe Mon Sep 17 00:00:00 2001 From: Pieter Noordhuis Date: Wed, 7 Feb 2024 11:03:49 +0100 Subject: [PATCH 2/2] Use acc.WorkspaceTest helper from bundle integration tests --- .../deploy_then_remove_resources_test.go | 23 ++++------- internal/bundle/empty_bundle_test.go | 9 ++-- internal/bundle/generate_job_test.go | 41 +++++++++---------- internal/bundle/generate_pipeline_test.go | 40 +++++++++--------- internal/bundle/helpers.go | 18 ++++---- internal/bundle/job_metadata_test.go | 18 ++++---- internal/bundle/local_state_staleness_test.go | 24 +++++------ internal/bundle/python_wheel_test.go | 27 +++++------- 8 files changed, 91 insertions(+), 109 deletions(-) diff --git a/internal/bundle/deploy_then_remove_resources_test.go b/internal/bundle/deploy_then_remove_resources_test.go index 73860593c0..72baf798cd 100644 --- a/internal/bundle/deploy_then_remove_resources_test.go +++ b/internal/bundle/deploy_then_remove_resources_test.go @@ -1,38 +1,33 @@ package bundle import ( - "context" "os" "path/filepath" "testing" - "github.com/databricks/cli/internal" - "github.com/databricks/databricks-sdk-go" + "github.com/databricks/cli/internal/acc" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestAccBundleDeployThenRemoveResources(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, wt := acc.WorkspaceTest(t) + w := wt.W uniqueId := uuid.New().String() - bundleRoot, err := initTestTemplate(t, "deploy_then_remove_resources", map[string]any{ + bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{ "unique_id": uniqueId, }) require.NoError(t, err) // deploy pipeline - err = deployBundle(t, bundleRoot) - require.NoError(t, err) - - w, err := databricks.NewWorkspaceClient() + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) // assert pipeline is created pipelineName := "test-bundle-pipeline-" + uniqueId - pipeline, err := w.Pipelines.GetByName(context.Background(), pipelineName) + pipeline, err := w.Pipelines.GetByName(ctx, pipelineName) require.NoError(t, err) assert.Equal(t, pipeline.Name, pipelineName) @@ -41,15 +36,15 @@ func TestAccBundleDeployThenRemoveResources(t *testing.T) { require.NoError(t, err) // deploy again - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) // assert pipeline is deleted - _, err = w.Pipelines.GetByName(context.Background(), pipelineName) + _, err = w.Pipelines.GetByName(ctx, pipelineName) assert.ErrorContains(t, err, "does not exist") t.Cleanup(func() { - err = destroyBundle(t, bundleRoot) + err = destroyBundle(t, ctx, bundleRoot) require.NoError(t, err) }) } diff --git a/internal/bundle/empty_bundle_test.go b/internal/bundle/empty_bundle_test.go index 9b39368f45..36883ae001 100644 --- a/internal/bundle/empty_bundle_test.go +++ b/internal/bundle/empty_bundle_test.go @@ -6,14 +6,13 @@ import ( "path/filepath" "testing" - "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" "github.com/google/uuid" "github.com/stretchr/testify/require" ) func TestAccEmptyBundleDeploy(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, _ := acc.WorkspaceTest(t) // create empty bundle tmpDir := t.TempDir() @@ -27,11 +26,11 @@ func TestAccEmptyBundleDeploy(t *testing.T) { f.Close() // deploy empty bundle - err = deployBundle(t, tmpDir) + err = deployBundle(t, ctx, tmpDir) require.NoError(t, err) t.Cleanup(func() { - err = destroyBundle(t, tmpDir) + err = destroyBundle(t, ctx, tmpDir) require.NoError(t, err) }) } diff --git a/internal/bundle/generate_job_test.go b/internal/bundle/generate_job_test.go index e9445abc55..e6f1578092 100644 --- a/internal/bundle/generate_job_test.go +++ b/internal/bundle/generate_job_test.go @@ -10,6 +10,7 @@ import ( "testing" "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" "github.com/databricks/cli/internal/testutil" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go" @@ -20,23 +21,22 @@ import ( ) func TestAccGenerateFromExistingJobAndDeploy(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, wt := acc.WorkspaceTest(t) + gt := &generateJobTest{T: t, w: wt.W} uniqueId := uuid.New().String() - bundleRoot, err := initTestTemplate(t, "with_includes", map[string]any{ + bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{ "unique_id": uniqueId, }) require.NoError(t, err) - jobId := createTestJob(t) + jobId := gt.createTestJob(ctx) t.Cleanup(func() { - destroyJob(t, jobId) - require.NoError(t, err) + gt.destroyJob(ctx, jobId) }) t.Setenv("BUNDLE_ROOT", bundleRoot) - c := internal.NewCobraTestRunner(t, "bundle", "generate", "job", + c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "generate", "job", "--existing-job-id", fmt.Sprint(jobId), "--config-dir", filepath.Join(bundleRoot, "resources"), "--source-dir", filepath.Join(bundleRoot, "src")) @@ -61,15 +61,22 @@ func TestAccGenerateFromExistingJobAndDeploy(t *testing.T) { require.Contains(t, generatedYaml, "spark_version: 13.3.x-scala2.12") require.Contains(t, generatedYaml, "num_workers: 1") - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) - err = destroyBundle(t, bundleRoot) + err = destroyBundle(t, ctx, bundleRoot) require.NoError(t, err) +} +type generateJobTest struct { + T *testing.T + w *databricks.WorkspaceClient } -func createTestJob(t *testing.T) int64 { +func (gt *generateJobTest) createTestJob(ctx context.Context) int64 { + t := gt.T + w := gt.w + var nodeTypeId string switch testutil.GetCloud(t) { case testutil.AWS: @@ -80,10 +87,6 @@ func createTestJob(t *testing.T) int64 { nodeTypeId = "n1-standard-4" } - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - ctx := context.Background() tmpdir := internal.TemporaryWorkspaceDir(t, w) f, err := filer.NewWorkspaceFilesClient(w, tmpdir) require.NoError(t, err) @@ -112,13 +115,9 @@ func createTestJob(t *testing.T) int64 { return resp.JobId } -func destroyJob(t *testing.T, jobId int64) { - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - ctx := context.Background() - err = w.Jobs.Delete(ctx, jobs.DeleteJob{ +func (gt *generateJobTest) destroyJob(ctx context.Context, jobId int64) { + err := gt.w.Jobs.Delete(ctx, jobs.DeleteJob{ JobId: jobId, }) - require.NoError(t, err) + require.NoError(gt.T, err) } diff --git a/internal/bundle/generate_pipeline_test.go b/internal/bundle/generate_pipeline_test.go index 7b2323e6a9..0005e29fa1 100644 --- a/internal/bundle/generate_pipeline_test.go +++ b/internal/bundle/generate_pipeline_test.go @@ -10,6 +10,7 @@ import ( "testing" "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" "github.com/databricks/cli/libs/filer" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/pipelines" @@ -18,23 +19,22 @@ import ( ) func TestAccGenerateFromExistingPipelineAndDeploy(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, wt := acc.WorkspaceTest(t) + gt := &generatePipelineTest{T: t, w: wt.W} uniqueId := uuid.New().String() - bundleRoot, err := initTestTemplate(t, "with_includes", map[string]any{ + bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{ "unique_id": uniqueId, }) require.NoError(t, err) - pipelineId := createTestPipeline(t) + pipelineId := gt.createTestPipeline(ctx) t.Cleanup(func() { - destroyPipeline(t, pipelineId) - require.NoError(t, err) + gt.destroyPipeline(ctx, pipelineId) }) t.Setenv("BUNDLE_ROOT", bundleRoot) - c := internal.NewCobraTestRunner(t, "bundle", "generate", "pipeline", + c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "generate", "pipeline", "--existing-pipeline-id", fmt.Sprint(pipelineId), "--config-dir", filepath.Join(bundleRoot, "resources"), "--source-dir", filepath.Join(bundleRoot, "src")) @@ -61,18 +61,22 @@ func TestAccGenerateFromExistingPipelineAndDeploy(t *testing.T) { require.Contains(t, generatedYaml, "- file:") require.Contains(t, generatedYaml, fmt.Sprintf("path: %s", filepath.Join("..", "src", "test.py"))) - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) - err = destroyBundle(t, bundleRoot) + err = destroyBundle(t, ctx, bundleRoot) require.NoError(t, err) } -func createTestPipeline(t *testing.T) string { - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) +type generatePipelineTest struct { + T *testing.T + w *databricks.WorkspaceClient +} + +func (gt *generatePipelineTest) createTestPipeline(ctx context.Context) string { + t := gt.T + w := gt.w - ctx := context.Background() tmpdir := internal.TemporaryWorkspaceDir(t, w) f, err := filer.NewWorkspaceFilesClient(w, tmpdir) require.NoError(t, err) @@ -103,13 +107,9 @@ func createTestPipeline(t *testing.T) string { return resp.PipelineId } -func destroyPipeline(t *testing.T, pipelineId string) { - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - ctx := context.Background() - err = w.Pipelines.Delete(ctx, pipelines.DeletePipelineRequest{ +func (gt *generatePipelineTest) destroyPipeline(ctx context.Context, pipelineId string) { + err := gt.w.Pipelines.Delete(ctx, pipelines.DeletePipelineRequest{ PipelineId: pipelineId, }) - require.NoError(t, err) + require.NoError(gt.T, err) } diff --git a/internal/bundle/helpers.go b/internal/bundle/helpers.go index 681edc2d6c..2c2b2dac94 100644 --- a/internal/bundle/helpers.go +++ b/internal/bundle/helpers.go @@ -15,7 +15,7 @@ import ( "github.com/databricks/cli/libs/template" ) -func initTestTemplate(t *testing.T, templateName string, config map[string]any) (string, error) { +func initTestTemplate(t *testing.T, ctx context.Context, templateName string, config map[string]any) (string, error) { templateRoot := filepath.Join("bundles", templateName) bundleRoot := t.TempDir() @@ -24,7 +24,7 @@ func initTestTemplate(t *testing.T, templateName string, config map[string]any) return "", err } - ctx := root.SetWorkspaceClient(context.Background(), nil) + ctx = root.SetWorkspaceClient(ctx, nil) cmd := cmdio.NewIO(flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "bundles") ctx = cmdio.InContext(ctx, cmd) @@ -46,15 +46,14 @@ func writeConfigFile(t *testing.T, config map[string]any) (string, error) { return filepath, err } -func deployBundle(t *testing.T, path string) error { +func deployBundle(t *testing.T, ctx context.Context, path string) error { t.Setenv("BUNDLE_ROOT", path) - c := internal.NewCobraTestRunner(t, "bundle", "deploy", "--force-lock") + c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock") _, _, err := c.Run() return err } -func runResource(t *testing.T, path string, key string) (string, error) { - ctx := context.Background() +func runResource(t *testing.T, ctx context.Context, path string, key string) (string, error) { ctx = cmdio.NewContext(ctx, cmdio.Default()) c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "run", key) @@ -62,8 +61,7 @@ func runResource(t *testing.T, path string, key string) (string, error) { return stdout.String(), err } -func runResourceWithParams(t *testing.T, path string, key string, params ...string) (string, error) { - ctx := context.Background() +func runResourceWithParams(t *testing.T, ctx context.Context, path string, key string, params ...string) (string, error) { ctx = cmdio.NewContext(ctx, cmdio.Default()) args := make([]string, 0) @@ -74,9 +72,9 @@ func runResourceWithParams(t *testing.T, path string, key string, params ...stri return stdout.String(), err } -func destroyBundle(t *testing.T, path string) error { +func destroyBundle(t *testing.T, ctx context.Context, path string) error { t.Setenv("BUNDLE_ROOT", path) - c := internal.NewCobraTestRunner(t, "bundle", "destroy", "--auto-approve") + c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "destroy", "--auto-approve") _, _, err := c.Run() return err } diff --git a/internal/bundle/job_metadata_test.go b/internal/bundle/job_metadata_test.go index 3e2bb7f032..0d8a431e4f 100644 --- a/internal/bundle/job_metadata_test.go +++ b/internal/bundle/job_metadata_test.go @@ -12,23 +12,21 @@ import ( "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/metadata" "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" + "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/filer" - "github.com/databricks/databricks-sdk-go" "github.com/google/uuid" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestAccJobsMetadataFile(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) + ctx, wt := acc.WorkspaceTest(t) + w := wt.W - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) - - nodeTypeId := internal.GetNodeTypeId(env) + nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV")) uniqueId := uuid.New().String() - bundleRoot, err := initTestTemplate(t, "job_metadata", map[string]any{ + bundleRoot, err := initTestTemplate(t, ctx, "job_metadata", map[string]any{ "unique_id": uniqueId, "node_type_id": nodeTypeId, "spark_version": "13.2.x-snapshot-scala2.12", @@ -36,12 +34,12 @@ func TestAccJobsMetadataFile(t *testing.T) { require.NoError(t, err) // deploy bundle - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) // Cleanup the deployed bundle t.Cleanup(func() { - err = destroyBundle(t, bundleRoot) + err = destroyBundle(t, ctx, bundleRoot) require.NoError(t, err) }) diff --git a/internal/bundle/local_state_staleness_test.go b/internal/bundle/local_state_staleness_test.go index 06cfe0e0d9..872ac8a8e8 100644 --- a/internal/bundle/local_state_staleness_test.go +++ b/internal/bundle/local_state_staleness_test.go @@ -5,7 +5,8 @@ import ( "testing" "github.com/databricks/cli/internal" - "github.com/databricks/databricks-sdk-go" + "github.com/databricks/cli/internal/acc" + "github.com/databricks/cli/libs/env" "github.com/databricks/databricks-sdk-go/listing" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/google/uuid" @@ -14,11 +15,8 @@ import ( ) func TestAccLocalStateStaleness(t *testing.T) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) - - w, err := databricks.NewWorkspaceClient() - require.NoError(t, err) + ctx, wt := acc.WorkspaceTest(t) + w := wt.W // The approach for this test is as follows: // 1) First deploy of bundle instance A @@ -27,10 +25,10 @@ func TestAccLocalStateStaleness(t *testing.T) { // Because of deploy (2), the locally cached state of bundle instance A should be stale. // Then for deploy (3), it must use the remote state over the stale local state. - nodeTypeId := internal.GetNodeTypeId(env) + nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV")) uniqueId := uuid.New().String() initialize := func() string { - root, err := initTestTemplate(t, "basic", map[string]any{ + root, err := initTestTemplate(t, ctx, "basic", map[string]any{ "unique_id": uniqueId, "node_type_id": nodeTypeId, "spark_version": "13.2.x-snapshot-scala2.12", @@ -38,26 +36,28 @@ func TestAccLocalStateStaleness(t *testing.T) { require.NoError(t, err) t.Cleanup(func() { - err = destroyBundle(t, root) + err = destroyBundle(t, ctx, root) require.NoError(t, err) }) return root } + var err error + bundleA := initialize() bundleB := initialize() // 1) Deploy bundle A - err = deployBundle(t, bundleA) + err = deployBundle(t, ctx, bundleA) require.NoError(t, err) // 2) Deploy bundle B - err = deployBundle(t, bundleB) + err = deployBundle(t, ctx, bundleB) require.NoError(t, err) // 3) Deploy bundle A again - err = deployBundle(t, bundleA) + err = deployBundle(t, ctx, bundleA) require.NoError(t, err) // Assert that there is only a single job in the workspace corresponding to this bundle. diff --git a/internal/bundle/python_wheel_test.go b/internal/bundle/python_wheel_test.go index c94ed93a3d..fc14fd17bd 100644 --- a/internal/bundle/python_wheel_test.go +++ b/internal/bundle/python_wheel_test.go @@ -4,24 +4,17 @@ import ( "testing" "github.com/databricks/cli/internal" + "github.com/databricks/cli/internal/acc" + "github.com/databricks/cli/libs/env" "github.com/google/uuid" "github.com/stretchr/testify/require" ) func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bool) { - env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV") - t.Log(env) - - var nodeTypeId string - if env == "gcp" { - nodeTypeId = "n1-standard-4" - } else if env == "aws" { - nodeTypeId = "i3.xlarge" - } else { - nodeTypeId = "Standard_DS4_v2" - } - - bundleRoot, err := initTestTemplate(t, "python_wheel_task", map[string]any{ + ctx, _ := acc.WorkspaceTest(t) + + nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV")) + bundleRoot, err := initTestTemplate(t, ctx, "python_wheel_task", map[string]any{ "node_type_id": nodeTypeId, "unique_id": uuid.New().String(), "spark_version": sparkVersion, @@ -29,20 +22,20 @@ func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bo }) require.NoError(t, err) - err = deployBundle(t, bundleRoot) + err = deployBundle(t, ctx, bundleRoot) require.NoError(t, err) t.Cleanup(func() { - destroyBundle(t, bundleRoot) + destroyBundle(t, ctx, bundleRoot) }) - out, err := runResource(t, bundleRoot, "some_other_job") + out, err := runResource(t, ctx, bundleRoot, "some_other_job") require.NoError(t, err) require.Contains(t, out, "Hello from my func") require.Contains(t, out, "Got arguments:") require.Contains(t, out, "['my_test_code', 'one', 'two']") - out, err = runResourceWithParams(t, bundleRoot, "some_other_job", "--python-params=param1,param2") + out, err = runResourceWithParams(t, ctx, bundleRoot, "some_other_job", "--python-params=param1,param2") require.NoError(t, err) require.Contains(t, out, "Hello from my func") require.Contains(t, out, "Got arguments:")