Use acc.WorkspaceTest helper from bundle integration tests

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)
 	})
 }

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)
 	})
 }

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)
 }

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)
 }

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,24 +46,22 @@ 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)
 	stdout, _, err := c.Run()
 	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
 }

internal/bundle/job_metadata_test.go

@@ -12,36 +12,34 @@ 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",
 	})
 	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)
 	})