From e0886ddb7d72802a8ef0d44c0a115f502e47ad4f Mon Sep 17 00:00:00 2001
From: Andrew Nester <andrew.nester@databricks.com>
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 <pieter.noordhuis@databricks.com>
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:")
