Use `acc.WorkspaceTest` helper from bundle integration tests (#1181)

## Changes

This helper:
* Constructs a context
* Constructs a `*databricks.WorkspaceClient`
* Ensures required environment variables are present to run an
integration test
* Enables debugging integration tests from VS Code

Debugging integration tests (from VS Code) is made possible by a prelude
in the helper that checks if the calling process is a debug binary, and
if so, sources environment variables from
`~/..databricks/debug-env.json` (if present).

## Tests

Integration tests still pass.

---------

Co-authored-by: Andrew Nester <andrew.nester@databricks.com>

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/cli/internal/acc"
-	"github.com/databricks/databricks-sdk-go"
 	"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")
+	ctx, wt := acc.WorkspaceTest(t)
-	t.Log(env)
+	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)
+	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
-
-	w, err := databricks.NewWorkspaceClient()
 	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)
 	})
 }

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")
+	ctx, _ := acc.WorkspaceTest(t)
-	t.Log(env)
 
 	// 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)
 	})
 }

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")
+	ctx, wt := acc.WorkspaceTest(t)
-	t.Log(env)
+	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)
+		gt.destroyJob(ctx, jobId)
-		require.NoError(t, err)
 	})
 
 	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)
-
 }
 
-func createTestJob(t *testing.T) int64 {
+type generateJobTest struct {
+	T *testing.T
+	w *databricks.WorkspaceClient
+}
+
+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) {
+func (gt *generateJobTest) destroyJob(ctx context.Context, jobId int64) {
-	w, err := databricks.NewWorkspaceClient()
+	err := gt.w.Jobs.Delete(ctx, jobs.DeleteJob{
-	require.NoError(t, err)
-
-	ctx := context.Background()
-	err = w.Jobs.Delete(ctx, jobs.DeleteJob{
 		JobId: jobId,
 	})
-	require.NoError(t, err)
+	require.NoError(gt.T, err)
 }

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")
+	ctx, wt := acc.WorkspaceTest(t)
-	t.Log(env)
+	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)
+		gt.destroyPipeline(ctx, pipelineId)
-		require.NoError(t, err)
 	})
 
 	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 {
+type generatePipelineTest struct {
-	w, err := databricks.NewWorkspaceClient()
+	T *testing.T
-	require.NoError(t, err)
+	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) {
+func (gt *generatePipelineTest) destroyPipeline(ctx context.Context, pipelineId string) {
-	w, err := databricks.NewWorkspaceClient()
+	err := gt.w.Pipelines.Delete(ctx, pipelines.DeletePipelineRequest{
-	require.NoError(t, err)
-
-	ctx := context.Background()
-	err = w.Pipelines.Delete(ctx, pipelines.DeletePipelineRequest{
 		PipelineId: pipelineId,
 	})
-	require.NoError(t, err)
+	require.NoError(gt.T, err)
 }

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) {
+func runResource(t *testing.T, ctx context.Context, path string, key string) (string, error) {
-	ctx := context.Background()
 	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) {
+func runResourceWithParams(t *testing.T, ctx context.Context, path string, key string, params ...string) (string, error) {
-	ctx := context.Background()
 	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
 }

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")
+	ctx, wt := acc.WorkspaceTest(t)
-	t.Log(env)
+	w := wt.W
 
-	w, err := databricks.NewWorkspaceClient()
+	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
-	require.NoError(t, err)
-
-	nodeTypeId := internal.GetNodeTypeId(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)
 	})
 

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")
+	ctx, wt := acc.WorkspaceTest(t)
-	t.Log(env)
+	w := wt.W
-
-	w, err := databricks.NewWorkspaceClient()
-	require.NoError(t, err)
 
 	// 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.

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")
+	ctx, _ := acc.WorkspaceTest(t)
-	t.Log(env)
 
-	var nodeTypeId string
+	nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
-	if env == "gcp" {
+	bundleRoot, err := initTestTemplate(t, ctx, "python_wheel_task", map[string]any{
-		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{
 		"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:")