Refactor initTestTemplate/deployBundle/destroyBundle to not return errors (#2017)

## Changes
These test helpers were updated to handle the error internally and not
return it. Since they have testing.T object, they can do so directly. On
the caller side, this functions were always followed by
require.NoError(t, err), that was cleaned up.

This approach helps reduce the setup/teardown boilerplate in the test
cases.

## Tests
Existing tests.

bundle/internal/tf/codegen/go.mod

@@ -20,7 +20,7 @@ require (
 	github.com/cloudflare/circl v1.5.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-retryablehttp v0.7.7 // indirect
-	golang.org/x/crypto v0.30.0 // indirect
+	golang.org/x/crypto v0.31.0 // indirect
 	golang.org/x/mod v0.22.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect
 	golang.org/x/text v0.21.0 // indirect

bundle/internal/tf/codegen/go.sum

@@ -60,8 +60,8 @@ github.com/xanzy/ssh-agent v0.3.3 h1:+/15pJfg/RsTxqYcX6fHqOXZwwMP+2VyYWJeWM2qQFM
 github.com/xanzy/ssh-agent v0.3.3/go.mod h1:6dzNDKs0J9rVPHPhaGCukekBHKqfl+L3KghI1Bc68Uw=
 github.com/zclconf/go-cty v1.15.1 h1:RgQYm4j2EvoBRXOPxhUvxPzRrGDo1eCOhHXuGfrj5S0=
 github.com/zclconf/go-cty v1.15.1/go.mod h1:VvMs5i0vgZdhYawQNq5kePSpLAoz8u1xvZgrPIxfnZE=
-golang.org/x/crypto v0.30.0 h1:RwoQn3GkWiMkzlX562cLB7OxWvjH1L8xutO2WoJcRoY=
-golang.org/x/crypto v0.30.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
+golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U=
+golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
 golang.org/x/exp v0.0.0-20241204233417-43b7b7cde48d h1:0olWaB5pg3+oychR51GUVCEsGkeCU/2JxjBgIo4f3M0=
 golang.org/x/exp v0.0.0-20241204233417-43b7b7cde48d/go.mod h1:qj5a5QZpwLU2NLQudwIN5koi3beDhSAlJwa67PuM98c=
 golang.org/x/mod v0.22.0 h1:D4nJWe9zXqHOmWqj4VMOJhvzj7bEZg4wEYa759z1pH4=

bundle/run/job_test.go

@@ -160,7 +160,7 @@ func TestJobRunnerRestart(t *testing.T) {
 		m := mocks.NewMockWorkspaceClient(t)
 		b.SetWorkpaceClient(m.WorkspaceClient)
 		ctx := context.Background()
-		ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", ""))
+		ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", ""))
 		ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
 
 		jobApi := m.GetMockJobsAPI()
@@ -231,7 +231,7 @@ func TestJobRunnerRestartForContinuousUnpausedJobs(t *testing.T) {
 	m := mocks.NewMockWorkspaceClient(t)
 	b.SetWorkpaceClient(m.WorkspaceClient)
 	ctx := context.Background()
-	ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
+	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
 	ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
 
 	jobApi := m.GetMockJobsAPI()

bundle/run/pipeline_test.go

@@ -76,7 +76,7 @@ func TestPipelineRunnerRestart(t *testing.T) {
 	}
 	b.SetWorkpaceClient(m.WorkspaceClient)
 	ctx := context.Background()
-	ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
+	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
 	ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
 
 	mockWait := &pipelines.WaitGetPipelineIdle[struct{}]{

cmd/labs/installed_test.go

@@ -11,7 +11,7 @@ import (
 func TestListsInstalledProjects(t *testing.T) {
 	ctx := context.Background()
 	ctx = env.WithUserHomeDir(ctx, "project/testdata/installed-in-home")
-	r := testcli.NewRunnerWithContext(t, ctx, "labs", "installed")
+	r := testcli.NewRunner(t, ctx, "labs", "installed")
 	r.RunAndExpectOutput(`
 	Name       Description        Version
 	blueprint  Blueprint Project  v0.3.15

cmd/labs/list_test.go

@@ -12,7 +12,7 @@ import (
 func TestListingWorks(t *testing.T) {
 	ctx := context.Background()
 	ctx = env.WithUserHomeDir(ctx, "project/testdata/installed-in-home")
-	c := testcli.NewRunnerWithContext(t, ctx, "labs", "list")
+	c := testcli.NewRunner(t, ctx, "labs", "list")
 	stdout, _, err := c.Run()
 	require.NoError(t, err)
 	require.Contains(t, stdout.String(), "ucx")

cmd/labs/project/command_test.go

@@ -30,7 +30,7 @@ func devEnvContext(t *testing.T) context.Context {
 
 func TestRunningBlueprintEcho(t *testing.T) {
 	ctx := devEnvContext(t)
-	r := testcli.NewRunnerWithContext(t, ctx, "labs", "blueprint", "echo")
+	r := testcli.NewRunner(t, ctx, "labs", "blueprint", "echo")
 	var out echoOut
 	r.RunAndParseJSON(&out)
 	assert.Equal(t, "echo", out.Command)
@@ -41,14 +41,14 @@ func TestRunningBlueprintEcho(t *testing.T) {
 
 func TestRunningBlueprintEchoProfileWrongOverride(t *testing.T) {
 	ctx := devEnvContext(t)
-	r := testcli.NewRunnerWithContext(t, ctx, "labs", "blueprint", "echo", "--profile", "workspace-profile")
+	r := testcli.NewRunner(t, ctx, "labs", "blueprint", "echo", "--profile", "workspace-profile")
 	_, _, err := r.Run()
 	assert.ErrorIs(t, err, databricks.ErrNotAccountClient)
 }
 
 func TestRunningCommand(t *testing.T) {
 	ctx := devEnvContext(t)
-	r := testcli.NewRunnerWithContext(t, ctx, "labs", "blueprint", "foo")
+	r := testcli.NewRunner(t, ctx, "labs", "blueprint", "foo")
 	r.WithStdin()
 	defer r.CloseStdin()
 
@@ -60,7 +60,7 @@ func TestRunningCommand(t *testing.T) {
 
 func TestRenderingTable(t *testing.T) {
 	ctx := devEnvContext(t)
-	r := testcli.NewRunnerWithContext(t, ctx, "labs", "blueprint", "table")
+	r := testcli.NewRunner(t, ctx, "labs", "blueprint", "table")
 	r.RunAndExpectOutput(`
 	Key    Value
 	First  Second

cmd/labs/project/installer_test.go

@@ -236,7 +236,7 @@ func TestInstallerWorksForReleases(t *testing.T) {
 	//     │               │   │       └── site-packages
 	//     │               │   │           ├── ...
 	//     │               │   │           ├── distutils-precedence.pth
-	r := testcli.NewRunnerWithContext(t, ctx, "labs", "install", "blueprint", "--debug")
+	r := testcli.NewRunner(t, ctx, "labs", "install", "blueprint", "--debug")
 	r.RunAndExpectOutput("setting up important infrastructure")
 }
 
@@ -356,7 +356,7 @@ account_id = abc
 	// 					└── databrickslabs-blueprint-releases.json
 
 	// `databricks labs install .` means "verify this installer i'm developing does work"
-	r := testcli.NewRunnerWithContext(t, ctx, "labs", "install", ".")
+	r := testcli.NewRunner(t, ctx, "labs", "install", ".")
 	r.WithStdin()
 	defer r.CloseStdin()
 
@@ -426,7 +426,7 @@ func TestUpgraderWorksForReleases(t *testing.T) {
 	ctx = env.Set(ctx, "DATABRICKS_CLUSTER_ID", "installer-cluster")
 	ctx = env.Set(ctx, "DATABRICKS_WAREHOUSE_ID", "installer-warehouse")
 
-	r := testcli.NewRunnerWithContext(t, ctx, "labs", "upgrade", "blueprint")
+	r := testcli.NewRunner(t, ctx, "labs", "upgrade", "blueprint")
 	r.RunAndExpectOutput("setting up important infrastructure")
 
 	// Check if the stub was called with the 'python -m pip install' command

cmd/root/io.go

@@ -45,8 +45,9 @@ func (f *outputFlag) initializeIO(cmd *cobra.Command) error {
 		headerTemplate = cmd.Annotations["headerTemplate"]
 	}
 
-	cmdIO := cmdio.NewIO(f.output, cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr(), headerTemplate, template)
-	ctx := cmdio.InContext(cmd.Context(), cmdIO)
+	ctx := cmd.Context()
+	cmdIO := cmdio.NewIO(ctx, f.output, cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr(), headerTemplate, template)
+	ctx = cmdio.InContext(ctx, cmdIO)
 	cmd.SetContext(ctx)
 	return nil
 }

go.mod

@@ -28,7 +28,7 @@ require (
 	golang.org/x/oauth2 v0.24.0
 	golang.org/x/sync v0.10.0
 	golang.org/x/term v0.27.0
-	golang.org/x/text v0.20.0
+	golang.org/x/text v0.21.0
 	gopkg.in/ini.v1 v1.67.0 // Apache 2.0
 	gopkg.in/yaml.v3 v3.0.1
 )
@@ -62,7 +62,7 @@ require (
 	go.opentelemetry.io/otel v1.24.0 // indirect
 	go.opentelemetry.io/otel/metric v1.24.0 // indirect
 	go.opentelemetry.io/otel/trace v1.24.0 // indirect
-	golang.org/x/crypto v0.24.0 // indirect
+	golang.org/x/crypto v0.31.0 // indirect
 	golang.org/x/net v0.26.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect
 	golang.org/x/time v0.5.0 // indirect

go.sum

@@ -176,8 +176,8 @@ go.opentelemetry.io/otel/trace v1.24.0 h1:CsKnnL4dUAr/0llH9FKuc698G04IrpWV0MQA/Y
 go.opentelemetry.io/otel/trace v1.24.0/go.mod h1:HPc3Xr/cOApsBI154IU0OI0HJexz+aw5uPdbs3UCjNU=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
-golang.org/x/crypto v0.24.0 h1:mnl8DM0o513X8fdIkmyFE/5hTYxbwYOjDS/+rK6qpRI=
-golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM=
+golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U=
+golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20240222234643-814bf88cf225 h1:LfspQV/FYTatPTr/3HzIcmiUFH7PGP+OQ6mgDYo3yuQ=
 golang.org/x/exp v0.0.0-20240222234643-814bf88cf225/go.mod h1:CxmFvTBINI24O/j8iY7H1xHzx2i4OsyguNBmN/uPtqc=
@@ -218,8 +218,8 @@ golang.org/x/term v0.27.0 h1:WP60Sv1nlK1T6SupCHbXzSaN0b9wUmsPoRS9b61A23Q=
 golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.20.0 h1:gK/Kv2otX8gz+wn7Rmb3vT96ZwuoxnQlY+HlJVj7Qug=
-golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4=
+golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
+golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
 golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
 golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

integration/bundle/artifacts_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testcli"
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/databricks-sdk-go/service/catalog"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
@@ -246,15 +247,14 @@ func TestUploadArtifactFileToVolumeThatDoesNotExist(t *testing.T) {
 		require.NoError(t, err)
 	})
 
-	bundleRoot, err := initTestTemplate(t, ctx, "artifact_path_with_volume", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "artifact_path_with_volume", map[string]any{
 		"unique_id":   uuid.New().String(),
 		"schema_name": schemaName,
 		"volume_name": "doesnotexist",
 	})
-	require.NoError(t, err)
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	stdout, stderr, err := testcli.RequireErrorRun(t, "bundle", "deploy")
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	stdout, stderr, err := testcli.RequireErrorRun(t, ctx, "bundle", "deploy")
 
 	assert.Error(t, err)
 	assert.Equal(t, fmt.Sprintf(`Error: volume /Volumes/main/%s/doesnotexist does not exist: Not Found
@@ -283,15 +283,14 @@ func TestUploadArtifactToVolumeNotYetDeployed(t *testing.T) {
 		require.NoError(t, err)
 	})
 
-	bundleRoot, err := initTestTemplate(t, ctx, "artifact_path_with_volume", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "artifact_path_with_volume", map[string]any{
 		"unique_id":   uuid.New().String(),
 		"schema_name": schemaName,
 		"volume_name": "my_volume",
 	})
-	require.NoError(t, err)
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	stdout, stderr, err := testcli.RequireErrorRun(t, "bundle", "deploy")
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	stdout, stderr, err := testcli.RequireErrorRun(t, ctx, "bundle", "deploy")
 
 	assert.Error(t, err)
 	assert.Equal(t, fmt.Sprintf(`Error: volume /Volumes/main/%s/my_volume does not exist: Not Found

integration/bundle/basic_test.go

@@ -16,27 +16,22 @@ func TestBasicBundleDeployWithFailOnActiveRuns(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	root, err := initTestTemplate(t, ctx, "basic", map[string]any{
+	root := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
 		"node_type_id":  nodeTypeId,
 		"spark_version": defaultSparkVersion,
 	})
-	require.NoError(t, err)
 
 	t.Cleanup(func() {
-		err = destroyBundle(t, ctx, root)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, root)
 	})
 
 	// deploy empty bundle
-	err = deployBundleWithFlags(t, ctx, root, []string{"--fail-on-active-runs"})
-	require.NoError(t, err)
+	deployBundleWithFlags(t, ctx, root, []string{"--fail-on-active-runs"})
 
 	// Remove .databricks directory to simulate a fresh deployment
-	err = os.RemoveAll(filepath.Join(root, ".databricks"))
-	require.NoError(t, err)
+	require.NoError(t, os.RemoveAll(filepath.Join(root, ".databricks")))
 
 	// deploy empty bundle again
-	err = deployBundleWithFlags(t, ctx, root, []string{"--fail-on-active-runs"})
-	require.NoError(t, err)
+	deployBundleWithFlags(t, ctx, root, []string{"--fail-on-active-runs"})
 }

integration/bundle/bind_resource_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testcli"
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/google/uuid"
@@ -22,30 +23,27 @@ func TestBindJobToExistingJob(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "basic", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
 		"spark_version": "13.3.x-scala2.12",
 		"node_type_id":  nodeTypeId,
 	})
-	require.NoError(t, err)
 
 	jobId := gt.createTestJob(ctx)
 	t.Cleanup(func() {
 		gt.destroyJob(ctx, jobId)
-		require.NoError(t, err)
 	})
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	c := testcli.NewRunner(t, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId), "--auto-approve")
-	_, _, err = c.Run()
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	c := testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId), "--auto-approve")
+	_, _, err := c.Run()
 	require.NoError(t, err)
 
 	// Remove .databricks directory to simulate a fresh deployment
 	err = os.RemoveAll(filepath.Join(bundleRoot, ".databricks"))
 	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	w, err := databricks.NewWorkspaceClient()
 	require.NoError(t, err)
@@ -58,7 +56,7 @@ func TestBindJobToExistingJob(t *testing.T) {
 	require.Equal(t, job.Settings.Name, fmt.Sprintf("test-job-basic-%s", uniqueId))
 	require.Contains(t, job.Settings.Tasks[0].SparkPythonTask.PythonFile, "hello_world.py")
 
-	c = testcli.NewRunner(t, "bundle", "deployment", "unbind", "foo")
+	c = testcli.NewRunner(t, ctx, "bundle", "deployment", "unbind", "foo")
 	_, _, err = c.Run()
 	require.NoError(t, err)
 
@@ -66,8 +64,7 @@ func TestBindJobToExistingJob(t *testing.T) {
 	err = os.RemoveAll(filepath.Join(bundleRoot, ".databricks"))
 	require.NoError(t, err)
 
-	err = destroyBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	destroyBundle(t, ctx, bundleRoot)
 
 	// Check that job is unbound and exists after bundle is destroyed
 	job, err = w.Jobs.Get(ctx, jobs.GetJobRequest{
@@ -84,32 +81,29 @@ func TestAbortBind(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "basic", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
 		"spark_version": "13.3.x-scala2.12",
 		"node_type_id":  nodeTypeId,
 	})
-	require.NoError(t, err)
 
 	jobId := gt.createTestJob(ctx)
 	t.Cleanup(func() {
 		gt.destroyJob(ctx, jobId)
-		err := destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	// Bind should fail because prompting is not possible.
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	t.Setenv("TERM", "dumb")
-	c := testcli.NewRunner(t, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId))
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "TERM", "dumb")
+	c := testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId))
 
 	// Expect error suggesting to use --auto-approve
-	_, _, err = c.Run()
+	_, _, err := c.Run()
 	assert.ErrorContains(t, err, "failed to bind the resource")
 	assert.ErrorContains(t, err, "This bind operation requires user confirmation, but the current console does not support prompting. Please specify --auto-approve if you would like to skip prompts and proceed")
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	w, err := databricks.NewWorkspaceClient()
 	require.NoError(t, err)
@@ -129,10 +123,9 @@ func TestGenerateAndBind(t *testing.T) {
 	gt := &generateJobTest{T: wt, w: wt.W}
 
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "with_includes", map[string]any{
 		"unique_id": uniqueId,
 	})
-	require.NoError(t, err)
 
 	w, err := databricks.NewWorkspaceClient()
 	require.NoError(t, err)
@@ -147,8 +140,8 @@ func TestGenerateAndBind(t *testing.T) {
 		}
 	})
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "generate", "job",
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	c := testcli.NewRunner(t, ctx, "bundle", "generate", "job",
 		"--key", "test_job_key",
 		"--existing-job-id", fmt.Sprint(jobId),
 		"--config-dir", filepath.Join(bundleRoot, "resources"),
@@ -164,15 +157,13 @@ func TestGenerateAndBind(t *testing.T) {
 
 	require.Len(t, matches, 1)
 
-	c = testcli.NewRunner(t, "bundle", "deployment", "bind", "test_job_key", fmt.Sprint(jobId), "--auto-approve")
+	c = testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "test_job_key", fmt.Sprint(jobId), "--auto-approve")
 	_, _, err = c.Run()
 	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
-	err = destroyBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	destroyBundle(t, ctx, bundleRoot)
 
 	// Check that job is bound and does not extsts after bundle is destroyed
 	_, err = w.Jobs.Get(ctx, jobs.GetJobRequest{

integration/bundle/clusters_test.go

@@ -20,16 +20,14 @@ func TestDeployBundleWithCluster(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	root, err := initTestTemplate(t, ctx, "clusters", map[string]any{
+	root := initTestTemplate(t, ctx, "clusters", map[string]any{
 		"unique_id":     uniqueId,
 		"node_type_id":  nodeTypeId,
 		"spark_version": defaultSparkVersion,
 	})
-	require.NoError(t, err)
 
 	t.Cleanup(func() {
-		err = destroyBundle(t, ctx, root)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, root)
 
 		cluster, err := wt.W.Clusters.GetByClusterName(ctx, fmt.Sprintf("test-cluster-%s", uniqueId))
 		if err != nil {
@@ -39,8 +37,7 @@ func TestDeployBundleWithCluster(t *testing.T) {
 		}
 	})
 
-	err = deployBundle(t, ctx, root)
-	require.NoError(t, err)
+	deployBundle(t, ctx, root)
 
 	// Cluster should exists after bundle deployment
 	cluster, err := wt.W.Clusters.GetByClusterName(ctx, fmt.Sprintf("test-cluster-%s", uniqueId))

integration/bundle/dashboards_test.go

@@ -18,19 +18,16 @@ func TestDashboards(t *testing.T) {
 
 	warehouseID := testutil.GetEnvOrSkipTest(t, "TEST_DEFAULT_WAREHOUSE_ID")
 	uniqueID := uuid.New().String()
-	root, err := initTestTemplate(t, ctx, "dashboards", map[string]any{
+	root := initTestTemplate(t, ctx, "dashboards", map[string]any{
 		"unique_id":    uniqueID,
 		"warehouse_id": warehouseID,
 	})
-	require.NoError(t, err)
 
 	t.Cleanup(func() {
-		err = destroyBundle(t, ctx, root)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, root)
 	})
 
-	err = deployBundle(t, ctx, root)
-	require.NoError(t, err)
+	deployBundle(t, ctx, root)
 
 	// Load bundle configuration by running the validate command.
 	b := unmarshalConfig(t, mustValidateBundle(t, ctx, root))
@@ -55,12 +52,11 @@ func TestDashboards(t *testing.T) {
 	require.NoError(t, err)
 
 	// Try to redeploy the bundle and confirm that the out of band modification is detected.
-	stdout, _, err := deployBundleWithArgs(t, ctx, root)
+	stdout, _, err := deployBundleWithArgsErr(t, ctx, root)
 	require.Error(t, err)
 	assert.Contains(t, stdout, `Error: dashboard "file_reference" has been modified remotely`+"\n")
 
 	// Redeploy the bundle with the --force flag and confirm that the out of band modification is ignored.
-	_, stderr, err := deployBundleWithArgs(t, ctx, root, "--force")
-	require.NoError(t, err)
+	_, stderr := deployBundleWithArgs(t, ctx, root, "--force")
 	assert.Contains(t, stderr, `Deployment complete!`+"\n")
 }

integration/bundle/deploy_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testcli"
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/service/catalog"
@@ -24,17 +25,14 @@ import (
 )
 
 func setupUcSchemaBundle(t *testing.T, ctx context.Context, w *databricks.WorkspaceClient, uniqueId string) string {
-	bundleRoot, err := initTestTemplate(t, ctx, "uc_schema", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "uc_schema", map[string]any{
 		"unique_id": uniqueId,
 	})
-	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	t.Cleanup(func() {
-		err := destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	// Assert the schema is created
@@ -96,8 +94,7 @@ func TestBundleDeployUcSchema(t *testing.T) {
 	require.NoError(t, err)
 
 	// Redeploy the bundle
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	// Assert the schema is deleted
 	_, err = w.Schemas.GetByFullName(ctx, strings.Join([]string{catalogName, schemaName}, "."))
@@ -118,9 +115,9 @@ func TestBundleDeployUcSchemaFailsWithoutAutoApprove(t *testing.T) {
 	require.NoError(t, err)
 
 	// Redeploy the bundle
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	t.Setenv("TERM", "dumb")
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock")
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "TERM", "dumb")
+	c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock")
 	stdout, stderr, err := c.Run()
 
 	assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -134,16 +131,14 @@ func TestBundlePipelineDeleteWithoutAutoApprove(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
 		"unique_id":     uniqueId,
 		"node_type_id":  nodeTypeId,
 		"spark_version": defaultSparkVersion,
 	})
-	require.NoError(t, err)
 
 	// deploy pipeline
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	// assert pipeline is created
 	pipelineName := "test-bundle-pipeline-" + uniqueId
@@ -162,9 +157,9 @@ func TestBundlePipelineDeleteWithoutAutoApprove(t *testing.T) {
 	require.NoError(t, err)
 
 	// Redeploy the bundle. Expect it to fail because deleting the pipeline requires --auto-approve.
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	t.Setenv("TERM", "dumb")
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock")
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "TERM", "dumb")
+	c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock")
 	stdout, stderr, err := c.Run()
 
 	assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -181,17 +176,14 @@ func TestBundlePipelineRecreateWithoutAutoApprove(t *testing.T) {
 	w := wt.W
 	uniqueId := uuid.New().String()
 
-	bundleRoot, err := initTestTemplate(t, ctx, "recreate_pipeline", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "recreate_pipeline", map[string]any{
 		"unique_id": uniqueId,
 	})
-	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	t.Cleanup(func() {
-		err := destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	// Assert the pipeline is created
@@ -201,9 +193,9 @@ func TestBundlePipelineRecreateWithoutAutoApprove(t *testing.T) {
 	require.Equal(t, pipelineName, pipeline.Name)
 
 	// Redeploy the bundle, pointing the DLT pipeline to a different UC catalog.
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	t.Setenv("TERM", "dumb")
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock", "--var=\"catalog=whatever\"")
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "TERM", "dumb")
+	c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock", "--var=\"catalog=whatever\"")
 	stdout, stderr, err := c.Run()
 
 	assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -220,22 +212,20 @@ func TestDeployBasicBundleLogs(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	root, err := initTestTemplate(t, ctx, "basic", map[string]any{
+	root := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
 		"node_type_id":  nodeTypeId,
 		"spark_version": defaultSparkVersion,
 	})
-	require.NoError(t, err)
 
 	t.Cleanup(func() {
-		err = destroyBundle(t, ctx, root)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, root)
 	})
 
 	currentUser, err := wt.W.CurrentUser.Me(ctx)
 	require.NoError(t, err)
 
-	stdout, stderr := blackBoxRun(t, root, "bundle", "deploy")
+	stdout, stderr := blackBoxRun(t, ctx, root, "bundle", "deploy")
 	assert.Equal(t, strings.Join([]string{
 		fmt.Sprintf("Uploading bundle files to /Workspace/Users/%s/.bundle/%s/files...", currentUser.UserName, uniqueId),
 		"Deploying resources...",
@@ -250,17 +240,14 @@ func TestDeployUcVolume(t *testing.T) {
 	w := wt.W
 
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "volume", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "volume", map[string]any{
 		"unique_id": uniqueId,
 	})
-	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	t.Cleanup(func() {
-		err := destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	// Assert the volume is created successfully
@@ -282,9 +269,9 @@ func TestDeployUcVolume(t *testing.T) {
 	assert.Equal(t, []catalog.Privilege{catalog.PrivilegeWriteVolume}, grants.PrivilegeAssignments[0].Privileges)
 
 	// Recreation of the volume without --auto-approve should fail since prompting is not possible
-	t.Setenv("TERM", "dumb")
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	stdout, stderr, err := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}").Run()
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "TERM", "dumb")
+	stdout, stderr, err := testcli.NewRunner(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}").Run()
 	assert.Error(t, err)
 	assert.Contains(t, stderr.String(), `This action will result in the deletion or recreation of the following volumes.
 For managed volumes, the files stored in the volume are also deleted from your
@@ -294,9 +281,9 @@ is removed from the catalog, but the underlying files are not deleted:
 	assert.Contains(t, stdout.String(), "the deployment requires destructive actions, but current console does not support prompting. Please specify --auto-approve if you would like to skip prompts and proceed")
 
 	// Successfully recreate the volume with --auto-approve
-	t.Setenv("TERM", "dumb")
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	_, _, err = testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}", "--auto-approve").Run()
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "TERM", "dumb")
+	_, _, err = testcli.NewRunner(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}", "--auto-approve").Run()
 	assert.NoError(t, err)
 
 	// Assert the volume is updated successfully

integration/bundle/deploy_then_remove_resources_test.go

@@ -18,16 +18,14 @@ func TestBundleDeployThenRemoveResources(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
 		"unique_id":     uniqueId,
 		"node_type_id":  nodeTypeId,
 		"spark_version": defaultSparkVersion,
 	})
-	require.NoError(t, err)
 
 	// deploy pipeline
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	// assert pipeline is created
 	pipelineName := "test-bundle-pipeline-" + uniqueId
@@ -46,8 +44,7 @@ func TestBundleDeployThenRemoveResources(t *testing.T) {
 	require.NoError(t, err)
 
 	// deploy again
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	// assert pipeline is deleted
 	_, err = w.Pipelines.GetByName(ctx, pipelineName)
@@ -58,7 +55,6 @@ func TestBundleDeployThenRemoveResources(t *testing.T) {
 	assert.ErrorContains(t, err, "does not exist")
 
 	t.Cleanup(func() {
-		err = destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 }

integration/bundle/deploy_to_shared_test.go

@@ -19,19 +19,16 @@ func TestDeployBasicToSharedWorkspacePath(t *testing.T) {
 	currentUser, err := wt.W.CurrentUser.Me(ctx)
 	require.NoError(t, err)
 
-	bundleRoot, err := initTestTemplate(t, ctx, "basic", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
 		"node_type_id":  nodeTypeId,
 		"spark_version": defaultSparkVersion,
 		"root_path":     fmt.Sprintf("/Shared/%s", currentUser.UserName),
 	})
-	require.NoError(t, err)
 
 	t.Cleanup(func() {
-		err = destroyBundle(wt, ctx, bundleRoot)
-		require.NoError(wt, err)
+		destroyBundle(wt, ctx, bundleRoot)
 	})
 
-	err = deployBundle(wt, ctx, bundleRoot)
-	require.NoError(wt, err)
+	deployBundle(wt, ctx, bundleRoot)
 }

integration/bundle/deployment_state_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/cli/bundle/deploy"
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/env"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 )
@@ -19,17 +20,16 @@ func TestFilesAreSyncedCorrectlyWhenNoSnapshot(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "basic", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "basic", map[string]any{
 		"unique_id":     uniqueId,
 		"spark_version": "13.3.x-scala2.12",
 		"node_type_id":  nodeTypeId,
 	})
-	require.NoError(t, err)
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
 
 	// Add some test file to the bundle
-	err = os.WriteFile(filepath.Join(bundleRoot, "test.py"), []byte("print('Hello, World!')"), 0o644)
+	err := os.WriteFile(filepath.Join(bundleRoot, "test.py"), []byte("print('Hello, World!')"), 0o644)
 	require.NoError(t, err)
 
 	err = os.WriteFile(filepath.Join(bundleRoot, "test_to_modify.py"), []byte("print('Hello, World!')"), 0o644)
@@ -39,11 +39,10 @@ func TestFilesAreSyncedCorrectlyWhenNoSnapshot(t *testing.T) {
 	err = os.WriteFile(filepath.Join(bundleRoot, "notebook.py"), []byte("# Databricks notebook source\nHello, World!"), 0o644)
 	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	t.Cleanup(func() {
-		require.NoError(t, destroyBundle(t, ctx, bundleRoot))
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	remoteRoot := getBundleRemoteRootPath(w, t, uniqueId)
@@ -79,8 +78,7 @@ func TestFilesAreSyncedCorrectlyWhenNoSnapshot(t *testing.T) {
 	err = os.WriteFile(filepath.Join(bundleRoot, "test_to_modify.py"), []byte("print('Modified!')"), 0o644)
 	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	// Check that removed file is not in workspace anymore
 	_, err = w.Workspace.GetStatusByPath(ctx, path.Join(remoteRoot, "files", "test.py"))

integration/bundle/destroy_test.go

@@ -20,22 +20,20 @@ func TestBundleDestroy(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
 		"unique_id":     uniqueId,
 		"node_type_id":  nodeTypeId,
 		"spark_version": defaultSparkVersion,
 	})
-	require.NoError(t, err)
 
 	snapshotsDir := filepath.Join(bundleRoot, ".databricks", "bundle", "default", "sync-snapshots")
 
 	// Assert the snapshot file does not exist
-	_, err = os.ReadDir(snapshotsDir)
+	_, err := os.ReadDir(snapshotsDir)
 	assert.ErrorIs(t, err, os.ErrNotExist)
 
 	// deploy resources
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	// Assert the snapshot file exists
 	entries, err := os.ReadDir(snapshotsDir)
@@ -60,8 +58,7 @@ func TestBundleDestroy(t *testing.T) {
 	assert.Equal(t, job.Settings.Name, jobName)
 
 	// destroy bundle
-	err = destroyBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	destroyBundle(t, ctx, bundleRoot)
 
 	// assert pipeline is deleted
 	_, err = w.Pipelines.GetByName(ctx, pipelineName)

integration/bundle/empty_bundle_test.go

@@ -26,11 +26,9 @@ func TestEmptyBundleDeploy(t *testing.T) {
 	f.Close()
 
 	// deploy empty bundle
-	err = deployBundle(t, ctx, tmpDir)
-	require.NoError(t, err)
+	deployBundle(t, ctx, tmpDir)
 
 	t.Cleanup(func() {
-		err = destroyBundle(t, ctx, tmpDir)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, tmpDir)
 	})
 }

integration/bundle/environments_test.go

@@ -13,17 +13,14 @@ func TestPythonWheelTaskWithEnvironmentsDeployAndRun(t *testing.T) {
 
 	ctx, _ := acc.WorkspaceTest(t)
 
-	bundleRoot, err := initTestTemplate(t, ctx, "python_wheel_task_with_environments", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "python_wheel_task_with_environments", map[string]any{
 		"unique_id": uuid.New().String(),
 	})
-	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	t.Cleanup(func() {
-		err := destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	out, err := runResource(t, ctx, bundleRoot, "some_other_job")

integration/bundle/generate_job_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testcli"
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/compute"
@@ -25,22 +26,21 @@ func TestGenerateFromExistingJobAndDeploy(t *testing.T) {
 	gt := &generateJobTest{T: wt, w: wt.W}
 
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "with_includes", map[string]any{
 		"unique_id": uniqueId,
 	})
-	require.NoError(t, err)
 
 	jobId := gt.createTestJob(ctx)
 	t.Cleanup(func() {
 		gt.destroyJob(ctx, jobId)
 	})
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "generate", "job",
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	c := testcli.NewRunner(t, ctx, "bundle", "generate", "job",
 		"--existing-job-id", fmt.Sprint(jobId),
 		"--config-dir", filepath.Join(bundleRoot, "resources"),
 		"--source-dir", filepath.Join(bundleRoot, "src"))
-	_, _, err = c.Run()
+	_, _, err := c.Run()
 	require.NoError(t, err)
 
 	_, err = os.Stat(filepath.Join(bundleRoot, "src", "test.py"))
@@ -61,11 +61,9 @@ func TestGenerateFromExistingJobAndDeploy(t *testing.T) {
 	require.Contains(t, generatedYaml, "spark_version: 13.3.x-scala2.12")
 	require.Contains(t, generatedYaml, "num_workers: 1")
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
-	err = destroyBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	destroyBundle(t, ctx, bundleRoot)
 }
 
 type generateJobTest struct {

integration/bundle/generate_pipeline_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/databricks/cli/internal/acc"
 	"github.com/databricks/cli/internal/testcli"
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/env"
 	"github.com/databricks/cli/libs/filer"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/pipelines"
@@ -24,22 +25,21 @@ func TestGenerateFromExistingPipelineAndDeploy(t *testing.T) {
 	gt := &generatePipelineTest{T: wt, w: wt.W}
 
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "with_includes", map[string]any{
 		"unique_id": uniqueId,
 	})
-	require.NoError(t, err)
 
 	pipelineId, name := gt.createTestPipeline(ctx)
 	t.Cleanup(func() {
 		gt.destroyPipeline(ctx, pipelineId)
 	})
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "generate", "pipeline",
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
+	c := testcli.NewRunner(t, ctx, "bundle", "generate", "pipeline",
 		"--existing-pipeline-id", fmt.Sprint(pipelineId),
 		"--config-dir", filepath.Join(bundleRoot, "resources"),
 		"--source-dir", filepath.Join(bundleRoot, "src"))
-	_, _, err = c.Run()
+	_, _, err := c.Run()
 	require.NoError(t, err)
 
 	_, err = os.Stat(filepath.Join(bundleRoot, "src", "notebook.py"))
@@ -69,11 +69,9 @@ func TestGenerateFromExistingPipelineAndDeploy(t *testing.T) {
 	require.Contains(t, generatedYaml, "- file:")
 	require.Contains(t, generatedYaml, fmt.Sprintf("path: %s", filepath.Join("..", "src", "test.py")))
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
-	err = destroyBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	destroyBundle(t, ctx, bundleRoot)
 }
 
 type generatePipelineTest struct {

integration/bundle/helpers_test.go

@@ -26,46 +26,43 @@ import (
 
 const defaultSparkVersion = "13.3.x-snapshot-scala2.12"
 
-func initTestTemplate(t testutil.TestingT, ctx context.Context, templateName string, config map[string]any) (string, error) {
+func initTestTemplate(t testutil.TestingT, ctx context.Context, templateName string, config map[string]any) string {
 	bundleRoot := t.TempDir()
 	return initTestTemplateWithBundleRoot(t, ctx, templateName, config, bundleRoot)
 }
 
-func initTestTemplateWithBundleRoot(t testutil.TestingT, ctx context.Context, templateName string, config map[string]any, bundleRoot string) (string, error) {
+func initTestTemplateWithBundleRoot(t testutil.TestingT, ctx context.Context, templateName string, config map[string]any, bundleRoot string) string {
 	templateRoot := filepath.Join("bundles", templateName)
 
-	configFilePath, err := writeConfigFile(t, config)
-	if err != nil {
-		return "", err
-	}
+	configFilePath := writeConfigFile(t, config)
 
 	ctx = root.SetWorkspaceClient(ctx, nil)
-	cmd := cmdio.NewIO(flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "", "bundles")
+	cmd := cmdio.NewIO(ctx, flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "", "bundles")
 	ctx = cmdio.InContext(ctx, cmd)
 
 	out, err := filer.NewLocalClient(bundleRoot)
 	require.NoError(t, err)
 	err = template.Materialize(ctx, configFilePath, os.DirFS(templateRoot), out)
-	return bundleRoot, err
+	require.NoError(t, err)
+	return bundleRoot
 }
 
-func writeConfigFile(t testutil.TestingT, config map[string]any) (string, error) {
+func writeConfigFile(t testutil.TestingT, config map[string]any) string {
 	bytes, err := json.Marshal(config)
-	if err != nil {
-		return "", err
-	}
+	require.NoError(t, err)
 
 	dir := t.TempDir()
 	filepath := filepath.Join(dir, "config.json")
 	t.Log("Configuration for template: ", string(bytes))
 
 	err = os.WriteFile(filepath, bytes, 0o644)
-	return filepath, err
+	require.NoError(t, err)
+	return filepath
 }
 
 func validateBundle(t testutil.TestingT, ctx context.Context, path string) ([]byte, error) {
 	ctx = env.Set(ctx, "BUNDLE_ROOT", path)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "validate", "--output", "json")
+	c := testcli.NewRunner(t, ctx, "bundle", "validate", "--output", "json")
 	stdout, _, err := c.Run()
 	return stdout.Bytes(), err
 }
@@ -83,35 +80,41 @@ func unmarshalConfig(t testutil.TestingT, data []byte) *bundle.Bundle {
 	return bundle
 }
 
-func deployBundle(t testutil.TestingT, ctx context.Context, path string) error {
+func deployBundle(t testutil.TestingT, ctx context.Context, path string) {
 	ctx = env.Set(ctx, "BUNDLE_ROOT", path)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock", "--auto-approve")
+	c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock", "--auto-approve")
 	_, _, err := c.Run()
-	return err
+	require.NoError(t, err)
 }
 
-func deployBundleWithArgs(t testutil.TestingT, ctx context.Context, path string, args ...string) (string, string, error) {
+func deployBundleWithArgsErr(t testutil.TestingT, ctx context.Context, path string, args ...string) (string, string, error) {
 	ctx = env.Set(ctx, "BUNDLE_ROOT", path)
 	args = append([]string{"bundle", "deploy"}, args...)
-	c := testcli.NewRunnerWithContext(t, ctx, args...)
+	c := testcli.NewRunner(t, ctx, args...)
 	stdout, stderr, err := c.Run()
 	return stdout.String(), stderr.String(), err
 }
 
-func deployBundleWithFlags(t testutil.TestingT, ctx context.Context, path string, flags []string) error {
+func deployBundleWithArgs(t testutil.TestingT, ctx context.Context, path string, args ...string) (string, string) {
+	stdout, stderr, err := deployBundleWithArgsErr(t, ctx, path, args...)
+	require.NoError(t, err)
+	return stdout, stderr
+}
+
+func deployBundleWithFlags(t testutil.TestingT, ctx context.Context, path string, flags []string) {
 	ctx = env.Set(ctx, "BUNDLE_ROOT", path)
 	args := []string{"bundle", "deploy", "--force-lock"}
 	args = append(args, flags...)
-	c := testcli.NewRunnerWithContext(t, ctx, args...)
+	c := testcli.NewRunner(t, ctx, args...)
 	_, _, err := c.Run()
-	return err
+	require.NoError(t, err)
 }
 
 func runResource(t testutil.TestingT, ctx context.Context, path, key string) (string, error) {
 	ctx = env.Set(ctx, "BUNDLE_ROOT", path)
 	ctx = cmdio.NewContext(ctx, cmdio.Default())
 
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "run", key)
+	c := testcli.NewRunner(t, ctx, "bundle", "run", key)
 	stdout, _, err := c.Run()
 	return stdout.String(), err
 }
@@ -123,16 +126,16 @@ func runResourceWithParams(t testutil.TestingT, ctx context.Context, path, key s
 	args := make([]string, 0)
 	args = append(args, "bundle", "run", key)
 	args = append(args, params...)
-	c := testcli.NewRunnerWithContext(t, ctx, args...)
+	c := testcli.NewRunner(t, ctx, args...)
 	stdout, _, err := c.Run()
 	return stdout.String(), err
 }
 
-func destroyBundle(t testutil.TestingT, ctx context.Context, path string) error {
+func destroyBundle(t testutil.TestingT, ctx context.Context, path string) {
 	ctx = env.Set(ctx, "BUNDLE_ROOT", path)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "destroy", "--auto-approve")
+	c := testcli.NewRunner(t, ctx, "bundle", "destroy", "--auto-approve")
 	_, _, err := c.Run()
-	return err
+	require.NoError(t, err)
 }
 
 func getBundleRemoteRootPath(w *databricks.WorkspaceClient, t testutil.TestingT, uniqueId string) string {
@@ -143,16 +146,20 @@ func getBundleRemoteRootPath(w *databricks.WorkspaceClient, t testutil.TestingT,
 	return root
 }
 
-func blackBoxRun(t testutil.TestingT, root string, args ...string) (stdout, stderr string) {
+func blackBoxRun(t testutil.TestingT, ctx context.Context, root string, args ...string) (stdout, stderr string) {
 	gitRoot, err := folders.FindDirWithLeaf(".", ".git")
 	require.NoError(t, err)
 
-	t.Setenv("BUNDLE_ROOT", root)
-
 	// Create the command
 	cmd := exec.Command("go", append([]string{"run", "main.go"}, args...)...)
 	cmd.Dir = gitRoot
 
+	// Configure the environment
+	ctx = env.Set(ctx, "BUNDLE_ROOT", root)
+	for key, value := range env.All(ctx) {
+		cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", key, value))
+	}
+
 	// Create buffers to capture output
 	var outBuffer, errBuffer bytes.Buffer
 	cmd.Stdout = &outBuffer

integration/bundle/init_test.go

@@ -20,8 +20,9 @@ import (
 )
 
 func TestBundleInitErrorOnUnknownFields(t *testing.T) {
+	ctx := context.Background()
 	tmpDir := t.TempDir()
-	_, _, err := testcli.RequireErrorRun(t, "bundle", "init", "./testdata/init/field-does-not-exist", "--output-dir", tmpDir)
+	_, _, err := testcli.RequireErrorRun(t, ctx, "bundle", "init", "./testdata/init/field-does-not-exist", "--output-dir", tmpDir)
 	assert.EqualError(t, err, "failed to compute file content for bar.tmpl. variable \"does_not_exist\" not defined")
 }
 
@@ -38,7 +39,7 @@ func TestBundleInitErrorOnUnknownFields(t *testing.T) {
 //     make changes that can break the MLOps Stacks DAB. In which case we should
 //     skip this test until the MLOps Stacks DAB is updated to work again.
 func TestBundleInitOnMlopsStacks(t *testing.T) {
-	_, wt := acc.WorkspaceTest(t)
+	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W
 
 	tmpDir1 := t.TempDir()
@@ -61,7 +62,7 @@ func TestBundleInitOnMlopsStacks(t *testing.T) {
 
 	// Run bundle init
 	assert.NoFileExists(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md"))
-	testcli.RequireSuccessfulRun(t, "bundle", "init", "mlops-stacks", "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir1, "config.json"))
+	testcli.RequireSuccessfulRun(t, ctx, "bundle", "init", "mlops-stacks", "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir1, "config.json"))
 
 	// Assert that the README.md file was created
 	contents := testutil.ReadFile(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md"))
@@ -69,17 +70,17 @@ func TestBundleInitOnMlopsStacks(t *testing.T) {
 
 	// Validate the stack
 	testutil.Chdir(t, filepath.Join(tmpDir2, "repo_name", projectName))
-	testcli.RequireSuccessfulRun(t, "bundle", "validate")
+	testcli.RequireSuccessfulRun(t, ctx, "bundle", "validate")
 
 	// Deploy the stack
-	testcli.RequireSuccessfulRun(t, "bundle", "deploy")
+	testcli.RequireSuccessfulRun(t, ctx, "bundle", "deploy")
 	t.Cleanup(func() {
 		// Delete the stack
-		testcli.RequireSuccessfulRun(t, "bundle", "destroy", "--auto-approve")
+		testcli.RequireSuccessfulRun(t, ctx, "bundle", "destroy", "--auto-approve")
 	})
 
 	// Get summary of the bundle deployment
-	stdout, _ := testcli.RequireSuccessfulRun(t, "bundle", "summary", "--output", "json")
+	stdout, _ := testcli.RequireSuccessfulRun(t, ctx, "bundle", "summary", "--output", "json")
 	summary := &config.Root{}
 	err = json.Unmarshal(stdout.Bytes(), summary)
 	require.NoError(t, err)
@@ -156,7 +157,7 @@ func TestBundleInitHelpers(t *testing.T) {
 		require.NoError(t, err)
 
 		// Run bundle init.
-		testcli.RequireSuccessfulRun(t, "bundle", "init", tmpDir, "--output-dir", tmpDir2)
+		testcli.RequireSuccessfulRun(t, ctx, "bundle", "init", tmpDir, "--output-dir", tmpDir2)
 
 		// Assert that the helper function was correctly computed.
 		contents := testutil.ReadFile(t, filepath.Join(tmpDir2, "foo.txt"))

integration/bundle/job_metadata_test.go

@@ -24,21 +24,18 @@ func TestJobsMetadataFile(t *testing.T) {
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
-	bundleRoot, err := initTestTemplate(t, ctx, "job_metadata", map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, "job_metadata", map[string]any{
 		"unique_id":     uniqueId,
 		"node_type_id":  nodeTypeId,
 		"spark_version": defaultSparkVersion,
 	})
-	require.NoError(t, err)
 
 	// deploy bundle
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	// Cleanup the deployed bundle
 	t.Cleanup(func() {
-		err = destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	// assert job 1 is created

integration/bundle/local_state_staleness_test.go

@@ -27,16 +27,14 @@ func TestLocalStateStaleness(t *testing.T) {
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	uniqueId := uuid.New().String()
 	initialize := func() string {
-		root, err := initTestTemplate(t, ctx, "basic", map[string]any{
+		root := initTestTemplate(t, ctx, "basic", map[string]any{
 			"unique_id":     uniqueId,
 			"node_type_id":  nodeTypeId,
 			"spark_version": defaultSparkVersion,
 		})
-		require.NoError(t, err)
 
 		t.Cleanup(func() {
-			err = destroyBundle(t, ctx, root)
-			require.NoError(t, err)
+			destroyBundle(t, ctx, root)
 		})
 
 		return root
@@ -48,16 +46,13 @@ func TestLocalStateStaleness(t *testing.T) {
 	bundleB := initialize()
 
 	// 1) Deploy bundle A
-	err = deployBundle(t, ctx, bundleA)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleA)
 
 	// 2) Deploy bundle B
-	err = deployBundle(t, ctx, bundleB)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleB)
 
 	// 3) Deploy bundle A again
-	err = deployBundle(t, ctx, bundleA)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleA)
 
 	// Assert that there is only a single job in the workspace corresponding to this bundle.
 	iter := w.Jobs.List(context.Background(), jobs.ListJobsRequest{

integration/bundle/python_wheel_test.go

@@ -15,21 +15,18 @@ func runPythonWheelTest(t *testing.T, templateName, sparkVersion string, pythonW
 
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	instancePoolId := env.Get(ctx, "TEST_INSTANCE_POOL_ID")
-	bundleRoot, err := initTestTemplate(t, ctx, templateName, map[string]any{
+	bundleRoot := initTestTemplate(t, ctx, templateName, map[string]any{
 		"node_type_id":         nodeTypeId,
 		"unique_id":            uuid.New().String(),
 		"spark_version":        sparkVersion,
 		"python_wheel_wrapper": pythonWheelWrapper,
 		"instance_pool_id":     instancePoolId,
 	})
-	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	t.Cleanup(func() {
-		err := destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	out, err := runResource(t, ctx, bundleRoot, "some_other_job")

integration/bundle/spark_jar_test.go

@@ -15,7 +15,7 @@ func runSparkJarTestCommon(t *testing.T, ctx context.Context, sparkVersion, arti
 	nodeTypeId := testutil.GetCloud(t).NodeTypeID()
 	tmpDir := t.TempDir()
 	instancePoolId := env.Get(ctx, "TEST_INSTANCE_POOL_ID")
-	bundleRoot, err := initTestTemplateWithBundleRoot(t, ctx, "spark_jar_task", map[string]any{
+	bundleRoot := initTestTemplateWithBundleRoot(t, ctx, "spark_jar_task", map[string]any{
 		"node_type_id":     nodeTypeId,
 		"unique_id":        uuid.New().String(),
 		"spark_version":    sparkVersion,
@@ -23,14 +23,11 @@ func runSparkJarTestCommon(t *testing.T, ctx context.Context, sparkVersion, arti
 		"artifact_path":    artifactPath,
 		"instance_pool_id": instancePoolId,
 	}, tmpDir)
-	require.NoError(t, err)
 
-	err = deployBundle(t, ctx, bundleRoot)
-	require.NoError(t, err)
+	deployBundle(t, ctx, bundleRoot)
 
 	t.Cleanup(func() {
-		err := destroyBundle(t, ctx, bundleRoot)
-		require.NoError(t, err)
+		destroyBundle(t, ctx, bundleRoot)
 	})
 
 	out, err := runResource(t, ctx, bundleRoot, "jar_job")

integration/cmd/alerts/alerts_test.go

@@ -1,6 +1,7 @@
 package alerts_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/databricks/cli/internal/testcli"
@@ -8,6 +9,7 @@ import (
 )
 
 func TestAlertsCreateErrWhenNoArguments(t *testing.T) {
-	_, _, err := testcli.RequireErrorRun(t, "alerts-legacy", "create")
+	ctx := context.Background()
+	_, _, err := testcli.RequireErrorRun(t, ctx, "alerts-legacy", "create")
 	assert.Equal(t, "please provide command input in JSON format by specifying the --json flag", err.Error())
 }

integration/cmd/api/api_test.go

@@ -1,6 +1,7 @@
 package api_test
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"path"
@@ -16,7 +17,9 @@ import (
 )
 
 func TestApiGet(t *testing.T) {
-	stdout, _ := testcli.RequireSuccessfulRun(t, "api", "get", "/api/2.0/preview/scim/v2/Me")
+	ctx := context.Background()
+
+	stdout, _ := testcli.RequireSuccessfulRun(t, ctx, "api", "get", "/api/2.0/preview/scim/v2/Me")
 
 	// Deserialize SCIM API response.
 	var out map[string]any
@@ -29,6 +32,8 @@ func TestApiGet(t *testing.T) {
 }
 
 func TestApiPost(t *testing.T) {
+	ctx := context.Background()
+
 	if testutil.GetCloud(t) == testutil.GCP {
 		t.Skip("DBFS REST API is disabled on gcp")
 	}
@@ -41,11 +46,11 @@ func TestApiPost(t *testing.T) {
 
 	// Post to mkdir
 	{
-		testcli.RequireSuccessfulRun(t, "api", "post", "--json=@"+requestPath, "/api/2.0/dbfs/mkdirs")
+		testcli.RequireSuccessfulRun(t, ctx, "api", "post", "--json=@"+requestPath, "/api/2.0/dbfs/mkdirs")
 	}
 
 	// Post to delete
 	{
-		testcli.RequireSuccessfulRun(t, "api", "post", "--json=@"+requestPath, "/api/2.0/dbfs/delete")
+		testcli.RequireSuccessfulRun(t, ctx, "api", "post", "--json=@"+requestPath, "/api/2.0/dbfs/delete")
 	}
 }

integration/cmd/auth/describe_test.go

@@ -13,7 +13,8 @@ import (
 func TestAuthDescribeSuccess(t *testing.T) {
 	t.Skipf("Skipping because of https://github.com/databricks/cli/issues/2010")
 
-	stdout, _ := testcli.RequireSuccessfulRun(t, "auth", "describe")
+	ctx := context.Background()
+	stdout, _ := testcli.RequireSuccessfulRun(t, ctx, "auth", "describe")
 	outStr := stdout.String()
 
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{})
@@ -34,7 +35,8 @@ func TestAuthDescribeSuccess(t *testing.T) {
 func TestAuthDescribeFailure(t *testing.T) {
 	t.Skipf("Skipping because of https://github.com/databricks/cli/issues/2010")
 
-	stdout, _ := testcli.RequireSuccessfulRun(t, "auth", "describe", "--profile", "nonexistent")
+	ctx := context.Background()
+	stdout, _ := testcli.RequireSuccessfulRun(t, ctx, "auth", "describe", "--profile", "nonexistent")
 	outStr := stdout.String()
 
 	require.NotEmpty(t, outStr)

integration/cmd/clusters/clusters_test.go

@@ -1,6 +1,7 @@
 package clusters_test
 
 import (
+	"context"
 	"fmt"
 	"regexp"
 	"testing"
@@ -14,7 +15,8 @@ import (
 )
 
 func TestClustersList(t *testing.T) {
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "clusters", "list")
+	ctx := context.Background()
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "clusters", "list")
 	outStr := stdout.String()
 	assert.Contains(t, outStr, "ID")
 	assert.Contains(t, outStr, "Name")
@@ -27,15 +29,17 @@ func TestClustersList(t *testing.T) {
 }
 
 func TestClustersGet(t *testing.T) {
+	ctx := context.Background()
 	clusterId := findValidClusterID(t)
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "clusters", "get", clusterId)
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "clusters", "get", clusterId)
 	outStr := stdout.String()
 	assert.Contains(t, outStr, fmt.Sprintf(`"cluster_id":"%s"`, clusterId))
 	assert.Equal(t, "", stderr.String())
 }
 
 func TestClusterCreateErrorWhenNoArguments(t *testing.T) {
-	_, _, err := testcli.RequireErrorRun(t, "clusters", "create")
+	ctx := context.Background()
+	_, _, err := testcli.RequireErrorRun(t, ctx, "clusters", "create")
 	assert.Contains(t, err.Error(), "accepts 1 arg(s), received 0")
 }
 

integration/cmd/fs/cat_test.go

@@ -23,11 +23,13 @@ func TestFsCat(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
+
 			err := f.Write(context.Background(), "hello.txt", strings.NewReader("abcd"), filer.CreateParentDirectories)
 			require.NoError(t, err)
 
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "cat", path.Join(tmpDir, "hello.txt"))
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "cat", path.Join(tmpDir, "hello.txt"))
 			assert.Equal(t, "", stderr.String())
 			assert.Equal(t, "abcd", stdout.String())
 		})
@@ -43,11 +45,13 @@ func TestFsCatOnADir(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
+
 			err := f.Mkdir(context.Background(), "dir1")
 			require.NoError(t, err)
 
-			_, _, err = testcli.RequireErrorRun(t, "fs", "cat", path.Join(tmpDir, "dir1"))
+			_, _, err = testcli.RequireErrorRun(t, ctx, "fs", "cat", path.Join(tmpDir, "dir1"))
 			assert.ErrorAs(t, err, &filer.NotAFile{})
 		})
 	}
@@ -62,16 +66,18 @@ func TestFsCatOnNonExistentFile(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			_, tmpDir := tc.setupFiler(t)
 
-			_, _, err := testcli.RequireErrorRun(t, "fs", "cat", path.Join(tmpDir, "non-existent-file"))
+			_, _, err := testcli.RequireErrorRun(t, ctx, "fs", "cat", path.Join(tmpDir, "non-existent-file"))
 			assert.ErrorIs(t, err, fs.ErrNotExist)
 		})
 	}
 }
 
 func TestFsCatForDbfsInvalidScheme(t *testing.T) {
-	_, _, err := testcli.RequireErrorRun(t, "fs", "cat", "dab:/non-existent-file")
+	ctx := context.Background()
+	_, _, err := testcli.RequireErrorRun(t, ctx, "fs", "cat", "dab:/non-existent-file")
 	assert.ErrorContains(t, err, "invalid scheme: dab")
 }
 
@@ -86,6 +92,6 @@ func TestFsCatDoesNotSupportOutputModeJson(t *testing.T) {
 	err = f.Write(ctx, "hello.txt", strings.NewReader("abc"))
 	require.NoError(t, err)
 
-	_, _, err = testcli.RequireErrorRun(t, "fs", "cat", "dbfs:"+path.Join(tmpDir, "hello.txt"), "--output=json")
+	_, _, err = testcli.RequireErrorRun(t, ctx, "fs", "cat", "dbfs:"+path.Join(tmpDir, "hello.txt"), "--output=json")
 	assert.ErrorContains(t, err, "json output not supported")
 }

integration/cmd/fs/completion_test.go

@@ -19,10 +19,11 @@ func setupCompletionFile(t *testing.T, f filer.Filer) {
 }
 
 func TestFsCompletion(t *testing.T) {
+	ctx := context.Background()
 	f, tmpDir := setupDbfsFiler(t)
 	setupCompletionFile(t, f)
 
-	stdout, _ := testcli.RequireSuccessfulRun(t, "__complete", "fs", "ls", tmpDir+"/")
+	stdout, _ := testcli.RequireSuccessfulRun(t, ctx, "__complete", "fs", "ls", tmpDir+"/")
 	expectedOutput := fmt.Sprintf("%s/dir1/\n:2\n", tmpDir)
 	assert.Equal(t, expectedOutput, stdout.String())
 }

integration/cmd/fs/cp_test.go

@@ -131,11 +131,12 @@ func TestFsCpDir(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceDir(t, context.Background(), sourceFiler)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", sourceDir, targetDir, "--recursive")
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", sourceDir, targetDir, "--recursive")
 
 			assertTargetDir(t, context.Background(), targetFiler)
 		})
@@ -151,11 +152,12 @@ func TestFsCpFileToFile(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceFile(t, context.Background(), sourceFiler)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", path.Join(sourceDir, "foo.txt"), path.Join(targetDir, "bar.txt"))
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", path.Join(sourceDir, "foo.txt"), path.Join(targetDir, "bar.txt"))
 
 			assertTargetFile(t, context.Background(), targetFiler, "bar.txt")
 		})
@@ -171,11 +173,12 @@ func TestFsCpFileToDir(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceFile(t, context.Background(), sourceFiler)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", path.Join(sourceDir, "foo.txt"), targetDir)
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", path.Join(sourceDir, "foo.txt"), targetDir)
 
 			assertTargetFile(t, context.Background(), targetFiler, "foo.txt")
 		})
@@ -194,7 +197,7 @@ func TestFsCpFileToDirForWindowsPaths(t *testing.T) {
 
 	windowsPath := filepath.Join(filepath.FromSlash(sourceDir), "foo.txt")
 
-	testcli.RequireSuccessfulRun(t, "fs", "cp", windowsPath, targetDir)
+	testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", windowsPath, targetDir)
 	assertTargetFile(t, ctx, targetFiler, "foo.txt")
 }
 
@@ -207,6 +210,7 @@ func TestFsCpDirToDirFileNotOverwritten(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceDir(t, context.Background(), sourceFiler)
@@ -215,7 +219,7 @@ func TestFsCpDirToDirFileNotOverwritten(t *testing.T) {
 			err := targetFiler.Write(context.Background(), "a/b/c/hello.txt", strings.NewReader("this should not be overwritten"), filer.CreateParentDirectories)
 			require.NoError(t, err)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", sourceDir, targetDir, "--recursive")
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", sourceDir, targetDir, "--recursive")
 			assertFileContent(t, context.Background(), targetFiler, "a/b/c/hello.txt", "this should not be overwritten")
 			assertFileContent(t, context.Background(), targetFiler, "query.sql", "SELECT 1")
 			assertFileContent(t, context.Background(), targetFiler, "pyNb.py", "# Databricks notebook source\nprint(123)")
@@ -232,6 +236,7 @@ func TestFsCpFileToDirFileNotOverwritten(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceDir(t, context.Background(), sourceFiler)
@@ -240,7 +245,7 @@ func TestFsCpFileToDirFileNotOverwritten(t *testing.T) {
 			err := targetFiler.Write(context.Background(), "a/b/c/hello.txt", strings.NewReader("this should not be overwritten"), filer.CreateParentDirectories)
 			require.NoError(t, err)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", path.Join(sourceDir, "a/b/c/hello.txt"), path.Join(targetDir, "a/b/c"))
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", path.Join(sourceDir, "a/b/c/hello.txt"), path.Join(targetDir, "a/b/c"))
 			assertFileContent(t, context.Background(), targetFiler, "a/b/c/hello.txt", "this should not be overwritten")
 		})
 	}
@@ -255,6 +260,7 @@ func TestFsCpFileToFileFileNotOverwritten(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceDir(t, context.Background(), sourceFiler)
@@ -263,7 +269,7 @@ func TestFsCpFileToFileFileNotOverwritten(t *testing.T) {
 			err := targetFiler.Write(context.Background(), "a/b/c/dontoverwrite.txt", strings.NewReader("this should not be overwritten"), filer.CreateParentDirectories)
 			require.NoError(t, err)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", path.Join(sourceDir, "a/b/c/hello.txt"), path.Join(targetDir, "a/b/c/dontoverwrite.txt"))
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", path.Join(sourceDir, "a/b/c/hello.txt"), path.Join(targetDir, "a/b/c/dontoverwrite.txt"))
 			assertFileContent(t, context.Background(), targetFiler, "a/b/c/dontoverwrite.txt", "this should not be overwritten")
 		})
 	}
@@ -278,6 +284,7 @@ func TestFsCpDirToDirWithOverwriteFlag(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceDir(t, context.Background(), sourceFiler)
@@ -286,7 +293,7 @@ func TestFsCpDirToDirWithOverwriteFlag(t *testing.T) {
 			err := targetFiler.Write(context.Background(), "a/b/c/hello.txt", strings.NewReader("this should be overwritten"), filer.CreateParentDirectories)
 			require.NoError(t, err)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", sourceDir, targetDir, "--recursive", "--overwrite")
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", sourceDir, targetDir, "--recursive", "--overwrite")
 			assertTargetDir(t, context.Background(), targetFiler)
 		})
 	}
@@ -301,6 +308,7 @@ func TestFsCpFileToFileWithOverwriteFlag(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceDir(t, context.Background(), sourceFiler)
@@ -309,7 +317,7 @@ func TestFsCpFileToFileWithOverwriteFlag(t *testing.T) {
 			err := targetFiler.Write(context.Background(), "a/b/c/overwritten.txt", strings.NewReader("this should be overwritten"), filer.CreateParentDirectories)
 			require.NoError(t, err)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", path.Join(sourceDir, "a/b/c/hello.txt"), path.Join(targetDir, "a/b/c/overwritten.txt"), "--overwrite")
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", path.Join(sourceDir, "a/b/c/hello.txt"), path.Join(targetDir, "a/b/c/overwritten.txt"), "--overwrite")
 			assertFileContent(t, context.Background(), targetFiler, "a/b/c/overwritten.txt", "hello, world\n")
 		})
 	}
@@ -324,6 +332,7 @@ func TestFsCpFileToDirWithOverwriteFlag(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceDir(t, context.Background(), sourceFiler)
@@ -332,7 +341,7 @@ func TestFsCpFileToDirWithOverwriteFlag(t *testing.T) {
 			err := targetFiler.Write(context.Background(), "a/b/c/hello.txt", strings.NewReader("this should be overwritten"), filer.CreateParentDirectories)
 			require.NoError(t, err)
 
-			testcli.RequireSuccessfulRun(t, "fs", "cp", path.Join(sourceDir, "a/b/c/hello.txt"), path.Join(targetDir, "a/b/c"), "--overwrite")
+			testcli.RequireSuccessfulRun(t, ctx, "fs", "cp", path.Join(sourceDir, "a/b/c/hello.txt"), path.Join(targetDir, "a/b/c"), "--overwrite")
 			assertFileContent(t, context.Background(), targetFiler, "a/b/c/hello.txt", "hello, world\n")
 		})
 	}
@@ -347,9 +356,10 @@ func TestFsCpErrorsWhenSourceIsDirWithoutRecursiveFlag(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			_, tmpDir := tc.setupFiler(t)
 
-			_, _, err := testcli.RequireErrorRun(t, "fs", "cp", path.Join(tmpDir), path.Join(tmpDir, "foobar"))
+			_, _, err := testcli.RequireErrorRun(t, ctx, "fs", "cp", path.Join(tmpDir), path.Join(tmpDir, "foobar"))
 			r := regexp.MustCompile("source path .* is a directory. Please specify the --recursive flag")
 			assert.Regexp(t, r, err.Error())
 		})
@@ -357,7 +367,8 @@ func TestFsCpErrorsWhenSourceIsDirWithoutRecursiveFlag(t *testing.T) {
 }
 
 func TestFsCpErrorsOnInvalidScheme(t *testing.T) {
-	_, _, err := testcli.RequireErrorRun(t, "fs", "cp", "dbfs:/a", "https:/b")
+	ctx := context.Background()
+	_, _, err := testcli.RequireErrorRun(t, ctx, "fs", "cp", "dbfs:/a", "https:/b")
 	assert.Equal(t, "invalid scheme: https", err.Error())
 }
 
@@ -370,6 +381,7 @@ func TestFsCpSourceIsDirectoryButTargetIsFile(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			sourceFiler, sourceDir := tc.setupSource(t)
 			targetFiler, targetDir := tc.setupTarget(t)
 			setupSourceDir(t, context.Background(), sourceFiler)
@@ -378,7 +390,7 @@ func TestFsCpSourceIsDirectoryButTargetIsFile(t *testing.T) {
 			err := targetFiler.Write(context.Background(), "my_target", strings.NewReader("I'll block any attempts to recursively copy"), filer.CreateParentDirectories)
 			require.NoError(t, err)
 
-			_, _, err = testcli.RequireErrorRun(t, "fs", "cp", sourceDir, path.Join(targetDir, "my_target"), "--recursive")
+			_, _, err = testcli.RequireErrorRun(t, ctx, "fs", "cp", sourceDir, path.Join(targetDir, "my_target"), "--recursive")
 			assert.Error(t, err)
 		})
 	}

integration/cmd/fs/ls_test.go

@@ -49,10 +49,11 @@ func TestFsLs(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 			setupLsFiles(t, f)
 
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "ls", tmpDir, "--output=json")
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "ls", tmpDir, "--output=json")
 			assert.Equal(t, "", stderr.String())
 
 			var parsedStdout []map[string]any
@@ -82,10 +83,11 @@ func TestFsLsWithAbsolutePaths(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 			setupLsFiles(t, f)
 
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "ls", tmpDir, "--output=json", "--absolute")
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "ls", tmpDir, "--output=json", "--absolute")
 			assert.Equal(t, "", stderr.String())
 
 			var parsedStdout []map[string]any
@@ -114,10 +116,12 @@ func TestFsLsOnFile(t *testing.T) {
 
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
+
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 			setupLsFiles(t, f)
 
-			_, _, err := testcli.RequireErrorRun(t, "fs", "ls", path.Join(tmpDir, "a", "hello.txt"), "--output=json")
+			_, _, err := testcli.RequireErrorRun(t, ctx, "fs", "ls", path.Join(tmpDir, "a", "hello.txt"), "--output=json")
 			assert.Regexp(t, regexp.MustCompile("not a directory: .*/a/hello.txt"), err.Error())
 			assert.ErrorAs(t, err, &filer.NotADirectory{})
 		})
@@ -133,9 +137,10 @@ func TestFsLsOnEmptyDir(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			_, tmpDir := tc.setupFiler(t)
 
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "ls", tmpDir, "--output=json")
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "ls", tmpDir, "--output=json")
 			assert.Equal(t, "", stderr.String())
 			var parsedStdout []map[string]any
 			err := json.Unmarshal(stdout.Bytes(), &parsedStdout)
@@ -156,9 +161,10 @@ func TestFsLsForNonexistingDir(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			_, tmpDir := tc.setupFiler(t)
 
-			_, _, err := testcli.RequireErrorRun(t, "fs", "ls", path.Join(tmpDir, "nonexistent"), "--output=json")
+			_, _, err := testcli.RequireErrorRun(t, ctx, "fs", "ls", path.Join(tmpDir, "nonexistent"), "--output=json")
 			assert.ErrorIs(t, err, fs.ErrNotExist)
 			assert.Regexp(t, regexp.MustCompile("no such directory: .*/nonexistent"), err.Error())
 		})
@@ -168,6 +174,7 @@ func TestFsLsForNonexistingDir(t *testing.T) {
 func TestFsLsWithoutScheme(t *testing.T) {
 	t.Parallel()
 
-	_, _, err := testcli.RequireErrorRun(t, "fs", "ls", "/path-without-a-dbfs-scheme", "--output=json")
+	ctx := context.Background()
+	_, _, err := testcli.RequireErrorRun(t, ctx, "fs", "ls", "/path-without-a-dbfs-scheme", "--output=json")
 	assert.ErrorIs(t, err, fs.ErrNotExist)
 }

integration/cmd/fs/mkdir_test.go

@@ -22,10 +22,11 @@ func TestFsMkdir(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 
 			// create directory "a"
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "mkdir", path.Join(tmpDir, "a"))
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "mkdir", path.Join(tmpDir, "a"))
 			assert.Equal(t, "", stderr.String())
 			assert.Equal(t, "", stdout.String())
 
@@ -47,10 +48,11 @@ func TestFsMkdirCreatesIntermediateDirectories(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 
 			// create directory "a/b/c"
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "mkdir", path.Join(tmpDir, "a", "b", "c"))
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "mkdir", path.Join(tmpDir, "a", "b", "c"))
 			assert.Equal(t, "", stderr.String())
 			assert.Equal(t, "", stdout.String())
 
@@ -84,6 +86,7 @@ func TestFsMkdirWhenDirectoryAlreadyExists(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 
 			// create directory "a"
@@ -91,7 +94,7 @@ func TestFsMkdirWhenDirectoryAlreadyExists(t *testing.T) {
 			require.NoError(t, err)
 
 			// assert run is successful without any errors
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "mkdir", path.Join(tmpDir, "a"))
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "mkdir", path.Join(tmpDir, "a"))
 			assert.Equal(t, "", stderr.String())
 			assert.Equal(t, "", stdout.String())
 		})
@@ -104,6 +107,7 @@ func TestFsMkdirWhenFileExistsAtPath(t *testing.T) {
 	t.Run("dbfs", func(t *testing.T) {
 		t.Parallel()
 
+		ctx := context.Background()
 		f, tmpDir := setupDbfsFiler(t)
 
 		// create file "hello"
@@ -111,7 +115,7 @@ func TestFsMkdirWhenFileExistsAtPath(t *testing.T) {
 		require.NoError(t, err)
 
 		// assert mkdir fails
-		_, _, err = testcli.RequireErrorRun(t, "fs", "mkdir", path.Join(tmpDir, "hello"))
+		_, _, err = testcli.RequireErrorRun(t, ctx, "fs", "mkdir", path.Join(tmpDir, "hello"))
 
 		// Different cloud providers or cloud configurations return different errors.
 		regex := regexp.MustCompile(`(^|: )Path is a file: .*$|(^|: )Cannot create directory .* because .* is an existing file\.$|(^|: )mkdirs\(hadoopPath: .*, permission: rwxrwxrwx\): failed$|(^|: )"The specified path already exists.".*$`)
@@ -121,6 +125,7 @@ func TestFsMkdirWhenFileExistsAtPath(t *testing.T) {
 	t.Run("uc-volumes", func(t *testing.T) {
 		t.Parallel()
 
+		ctx := context.Background()
 		f, tmpDir := setupUcVolumesFiler(t)
 
 		// create file "hello"
@@ -128,7 +133,7 @@ func TestFsMkdirWhenFileExistsAtPath(t *testing.T) {
 		require.NoError(t, err)
 
 		// assert mkdir fails
-		_, _, err = testcli.RequireErrorRun(t, "fs", "mkdir", path.Join(tmpDir, "hello"))
+		_, _, err = testcli.RequireErrorRun(t, ctx, "fs", "mkdir", path.Join(tmpDir, "hello"))
 
 		assert.ErrorAs(t, err, &filer.FileAlreadyExistsError{})
 	})

integration/cmd/fs/rm_test.go

@@ -23,6 +23,7 @@ func TestFsRmFile(t *testing.T) {
 			t.Parallel()
 
 			// Create a file
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 			err := f.Write(context.Background(), "hello.txt", strings.NewReader("abcd"), filer.CreateParentDirectories)
 			require.NoError(t, err)
@@ -32,7 +33,7 @@ func TestFsRmFile(t *testing.T) {
 			assert.NoError(t, err)
 
 			// Run rm command
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "rm", path.Join(tmpDir, "hello.txt"))
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "rm", path.Join(tmpDir, "hello.txt"))
 			assert.Equal(t, "", stderr.String())
 			assert.Equal(t, "", stdout.String())
 
@@ -53,6 +54,7 @@ func TestFsRmEmptyDir(t *testing.T) {
 			t.Parallel()
 
 			// Create a directory
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 			err := f.Mkdir(context.Background(), "a")
 			require.NoError(t, err)
@@ -62,7 +64,7 @@ func TestFsRmEmptyDir(t *testing.T) {
 			assert.NoError(t, err)
 
 			// Run rm command
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "rm", path.Join(tmpDir, "a"))
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "rm", path.Join(tmpDir, "a"))
 			assert.Equal(t, "", stderr.String())
 			assert.Equal(t, "", stdout.String())
 
@@ -83,6 +85,7 @@ func TestFsRmNonEmptyDirectory(t *testing.T) {
 			t.Parallel()
 
 			// Create a directory
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 			err := f.Mkdir(context.Background(), "a")
 			require.NoError(t, err)
@@ -96,7 +99,7 @@ func TestFsRmNonEmptyDirectory(t *testing.T) {
 			assert.NoError(t, err)
 
 			// Run rm command
-			_, _, err = testcli.RequireErrorRun(t, "fs", "rm", path.Join(tmpDir, "a"))
+			_, _, err = testcli.RequireErrorRun(t, ctx, "fs", "rm", path.Join(tmpDir, "a"))
 			assert.ErrorIs(t, err, fs.ErrInvalid)
 			assert.ErrorAs(t, err, &filer.DirectoryNotEmptyError{})
 		})
@@ -112,10 +115,11 @@ func TestFsRmForNonExistentFile(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			_, tmpDir := tc.setupFiler(t)
 
 			// Expect error if file does not exist
-			_, _, err := testcli.RequireErrorRun(t, "fs", "rm", path.Join(tmpDir, "does-not-exist"))
+			_, _, err := testcli.RequireErrorRun(t, ctx, "fs", "rm", path.Join(tmpDir, "does-not-exist"))
 			assert.ErrorIs(t, err, fs.ErrNotExist)
 		})
 	}
@@ -130,6 +134,7 @@ func TestFsRmDirRecursively(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
+			ctx := context.Background()
 			f, tmpDir := tc.setupFiler(t)
 
 			// Create a directory
@@ -145,7 +150,7 @@ func TestFsRmDirRecursively(t *testing.T) {
 			assert.NoError(t, err)
 
 			// Run rm command
-			stdout, stderr := testcli.RequireSuccessfulRun(t, "fs", "rm", path.Join(tmpDir, "a"), "--recursive")
+			stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "fs", "rm", path.Join(tmpDir, "a"), "--recursive")
 			assert.Equal(t, "", stderr.String())
 			assert.Equal(t, "", stdout.String())
 

integration/cmd/jobs/jobs_test.go

@@ -1,6 +1,7 @@
 package jobs_test
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"testing"
@@ -13,10 +14,11 @@ import (
 
 func TestCreateJob(t *testing.T) {
 	testutil.Require(t, testutil.Azure)
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "jobs", "create", "--json", "@testdata/create_job_without_workers.json", "--log-level=debug")
+	ctx := context.Background()
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "jobs", "create", "--json", "@testdata/create_job_without_workers.json", "--log-level=debug")
 	assert.Empty(t, stderr.String())
 	var output map[string]int
 	err := json.Unmarshal(stdout.Bytes(), &output)
 	require.NoError(t, err)
-	testcli.RequireSuccessfulRun(t, "jobs", "delete", fmt.Sprint(output["job_id"]), "--log-level=debug")
+	testcli.RequireSuccessfulRun(t, ctx, "jobs", "delete", fmt.Sprint(output["job_id"]), "--log-level=debug")
 }

integration/cmd/repos/repos_test.go

@@ -53,7 +53,7 @@ func TestReposCreateWithProvider(t *testing.T) {
 	w := wt.W
 	repoPath := synthesizeTemporaryRepoPath(t, w, ctx)
 
-	_, stderr := testcli.RequireSuccessfulRun(t, "repos", "create", repoUrl, "gitHub", "--path", repoPath)
+	_, stderr := testcli.RequireSuccessfulRun(t, ctx, "repos", "create", repoUrl, "gitHub", "--path", repoPath)
 	assert.Equal(t, "", stderr.String())
 
 	// Confirm the repo was created.
@@ -67,7 +67,7 @@ func TestReposCreateWithoutProvider(t *testing.T) {
 	w := wt.W
 	repoPath := synthesizeTemporaryRepoPath(t, w, ctx)
 
-	_, stderr := testcli.RequireSuccessfulRun(t, "repos", "create", repoUrl, "--path", repoPath)
+	_, stderr := testcli.RequireSuccessfulRun(t, ctx, "repos", "create", repoUrl, "--path", repoPath)
 	assert.Equal(t, "", stderr.String())
 
 	// Confirm the repo was created.
@@ -83,22 +83,22 @@ func TestReposGet(t *testing.T) {
 	repoId, repoPath := createTemporaryRepo(t, w, ctx)
 
 	// Get by ID
-	byIdOutput, stderr := testcli.RequireSuccessfulRun(t, "repos", "get", strconv.FormatInt(repoId, 10), "--output=json")
+	byIdOutput, stderr := testcli.RequireSuccessfulRun(t, ctx, "repos", "get", strconv.FormatInt(repoId, 10), "--output=json")
 	assert.Equal(t, "", stderr.String())
 
 	// Get by path
-	byPathOutput, stderr := testcli.RequireSuccessfulRun(t, "repos", "get", repoPath, "--output=json")
+	byPathOutput, stderr := testcli.RequireSuccessfulRun(t, ctx, "repos", "get", repoPath, "--output=json")
 	assert.Equal(t, "", stderr.String())
 
 	// Output should be the same
 	assert.Equal(t, byIdOutput.String(), byPathOutput.String())
 
 	// Get by path fails
-	_, stderr, err := testcli.RequireErrorRun(t, "repos", "get", repoPath+"-doesntexist", "--output=json")
+	_, stderr, err := testcli.RequireErrorRun(t, ctx, "repos", "get", repoPath+"-doesntexist", "--output=json")
 	assert.ErrorContains(t, err, "failed to look up repo")
 
 	// Get by path resolves to something other than a repo
-	_, stderr, err = testcli.RequireErrorRun(t, "repos", "get", "/Repos", "--output=json")
+	_, stderr, err = testcli.RequireErrorRun(t, ctx, "repos", "get", "/Repos", "--output=json")
 	assert.ErrorContains(t, err, "is not a repo")
 }
 
@@ -109,11 +109,11 @@ func TestReposUpdate(t *testing.T) {
 	repoId, repoPath := createTemporaryRepo(t, w, ctx)
 
 	// Update by ID
-	byIdOutput, stderr := testcli.RequireSuccessfulRun(t, "repos", "update", strconv.FormatInt(repoId, 10), "--branch", "ide")
+	byIdOutput, stderr := testcli.RequireSuccessfulRun(t, ctx, "repos", "update", strconv.FormatInt(repoId, 10), "--branch", "ide")
 	assert.Equal(t, "", stderr.String())
 
 	// Update by path
-	byPathOutput, stderr := testcli.RequireSuccessfulRun(t, "repos", "update", repoPath, "--branch", "ide")
+	byPathOutput, stderr := testcli.RequireSuccessfulRun(t, ctx, "repos", "update", repoPath, "--branch", "ide")
 	assert.Equal(t, "", stderr.String())
 
 	// Output should be the same
@@ -127,7 +127,7 @@ func TestReposDeleteByID(t *testing.T) {
 	repoId, _ := createTemporaryRepo(t, w, ctx)
 
 	// Delete by ID
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "repos", "delete", strconv.FormatInt(repoId, 10))
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "repos", "delete", strconv.FormatInt(repoId, 10))
 	assert.Equal(t, "", stdout.String())
 	assert.Equal(t, "", stderr.String())
 
@@ -143,7 +143,7 @@ func TestReposDeleteByPath(t *testing.T) {
 	repoId, repoPath := createTemporaryRepo(t, w, ctx)
 
 	// Delete by path
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "repos", "delete", repoPath)
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "repos", "delete", repoPath)
 	assert.Equal(t, "", stdout.String())
 	assert.Equal(t, "", stderr.String())
 

integration/cmd/secrets/secrets_test.go

@@ -15,7 +15,8 @@ import (
 )
 
 func TestSecretsCreateScopeErrWhenNoArguments(t *testing.T) {
-	_, _, err := testcli.RequireErrorRun(t, "secrets", "create-scope")
+	ctx := context.Background()
+	_, _, err := testcli.RequireErrorRun(t, ctx, "secrets", "create-scope")
 	assert.Contains(t, err.Error(), "accepts 1 arg(s), received 0")
 }
 
@@ -69,7 +70,7 @@ func TestSecretsPutSecretStringValue(tt *testing.T) {
 	key := "test-key"
 	value := "test-value\nwith-newlines\n"
 
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "secrets", "put-secret", scope, key, "--string-value", value)
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "secrets", "put-secret", scope, key, "--string-value", value)
 	assert.Empty(t, stdout)
 	assert.Empty(t, stderr)
 
@@ -83,7 +84,7 @@ func TestSecretsPutSecretBytesValue(tt *testing.T) {
 	key := "test-key"
 	value := []byte{0x00, 0x01, 0x02, 0x03}
 
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "secrets", "put-secret", scope, key, "--bytes-value", string(value))
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "secrets", "put-secret", scope, key, "--bytes-value", string(value))
 	assert.Empty(t, stdout)
 	assert.Empty(t, stderr)
 

integration/cmd/storage_credentials/storage_credentials_test.go

@@ -10,12 +10,12 @@ import (
 )
 
 func TestStorageCredentialsListRendersResponse(t *testing.T) {
-	_, _ = acc.WorkspaceTest(t)
+	ctx, _ := acc.WorkspaceTest(t)
 
 	// Check if metastore is assigned for the workspace, otherwise test will fail
 	t.Log(testutil.GetEnvOrSkipTest(t, "TEST_METASTORE_ID"))
 
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "storage-credentials", "list")
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "storage-credentials", "list")
 	assert.NotEmpty(t, stdout)
 	assert.Empty(t, stderr)
 }

integration/cmd/sync/sync_test.go

@@ -72,8 +72,8 @@ type syncTest struct {
 	remoteRoot string
 }
 
-func setupSyncTest(t *testing.T, args ...string) *syncTest {
-	_, wt := acc.WorkspaceTest(t)
+func setupSyncTest(t *testing.T, args ...string) (context.Context, *syncTest) {
+	ctx, wt := acc.WorkspaceTest(t)
 	w := wt.W
 
 	localRoot := t.TempDir()
@@ -90,10 +90,10 @@ func setupSyncTest(t *testing.T, args ...string) *syncTest {
 		"json",
 	}, args...)
 
-	c := testcli.NewRunner(t, args...)
+	c := testcli.NewRunner(t, ctx, args...)
 	c.RunBackground()
 
-	return &syncTest{
+	return ctx, &syncTest{
 		t:          t,
 		c:          c,
 		w:          w,
@@ -231,8 +231,7 @@ func (a *syncTest) snapshotContains(files []string) {
 }
 
 func TestSyncFullFileSync(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--full", "--watch")
+	ctx, assertSync := setupSyncTest(t, "--full", "--watch")
 
 	// .gitignore is created by the sync process to enforce .databricks is not synced
 	assertSync.waitForCompletionMarker()
@@ -263,8 +262,7 @@ func TestSyncFullFileSync(t *testing.T) {
 }
 
 func TestSyncIncrementalFileSync(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--watch")
+	ctx, assertSync := setupSyncTest(t, "--watch")
 
 	// .gitignore is created by the sync process to enforce .databricks is not synced
 	assertSync.waitForCompletionMarker()
@@ -297,8 +295,7 @@ func TestSyncIncrementalFileSync(t *testing.T) {
 }
 
 func TestSyncNestedFolderSync(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--watch")
+	ctx, assertSync := setupSyncTest(t, "--watch")
 
 	// .gitignore is created by the sync process to enforce .databricks is not synced
 	assertSync.waitForCompletionMarker()
@@ -325,8 +322,7 @@ func TestSyncNestedFolderSync(t *testing.T) {
 }
 
 func TestSyncNestedFolderDoesntFailOnNonEmptyDirectory(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--watch")
+	ctx, assertSync := setupSyncTest(t, "--watch")
 
 	// .gitignore is created by the sync process to enforce .databricks is not synced
 	assertSync.waitForCompletionMarker()
@@ -358,8 +354,7 @@ func TestSyncNestedFolderDoesntFailOnNonEmptyDirectory(t *testing.T) {
 }
 
 func TestSyncNestedSpacePlusAndHashAreEscapedSync(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--watch")
+	ctx, assertSync := setupSyncTest(t, "--watch")
 
 	// .gitignore is created by the sync process to enforce .databricks is not synced
 	assertSync.waitForCompletionMarker()
@@ -394,8 +389,7 @@ func TestSyncNestedSpacePlusAndHashAreEscapedSync(t *testing.T) {
 // In the above scenario sync should delete the empty folder and add foo to the remote
 // file system
 func TestSyncIncrementalFileOverwritesFolder(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--watch")
+	ctx, assertSync := setupSyncTest(t, "--watch")
 
 	// create foo/bar.txt
 	localFilePath := filepath.Join(assertSync.localRoot, "foo/bar.txt")
@@ -424,8 +418,7 @@ func TestSyncIncrementalFileOverwritesFolder(t *testing.T) {
 }
 
 func TestSyncIncrementalSyncPythonNotebookToFile(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--watch")
+	ctx, assertSync := setupSyncTest(t, "--watch")
 
 	// create python notebook
 	localFilePath := filepath.Join(assertSync.localRoot, "foo.py")
@@ -455,8 +448,7 @@ func TestSyncIncrementalSyncPythonNotebookToFile(t *testing.T) {
 }
 
 func TestSyncIncrementalSyncFileToPythonNotebook(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--watch")
+	ctx, assertSync := setupSyncTest(t, "--watch")
 
 	// create vanilla python file
 	localFilePath := filepath.Join(assertSync.localRoot, "foo.py")
@@ -479,8 +471,7 @@ func TestSyncIncrementalSyncFileToPythonNotebook(t *testing.T) {
 }
 
 func TestSyncIncrementalSyncPythonNotebookDelete(t *testing.T) {
-	ctx := context.Background()
-	assertSync := setupSyncTest(t, "--watch")
+	ctx, assertSync := setupSyncTest(t, "--watch")
 
 	// create python notebook
 	localFilePath := filepath.Join(assertSync.localRoot, "foo.py")

integration/cmd/unknown_command_test.go

@@ -1,6 +1,7 @@
 package cmd_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/databricks/cli/internal/testcli"
@@ -8,7 +9,8 @@ import (
 )
 
 func TestUnknownCommand(t *testing.T) {
-	stdout, stderr, err := testcli.RequireErrorRun(t, "unknown-command")
+	ctx := context.Background()
+	stdout, stderr, err := testcli.RequireErrorRun(t, ctx, "unknown-command")
 
 	assert.Error(t, err, "unknown command", `unknown command "unknown-command" for "databricks"`)
 	assert.Equal(t, "", stdout.String())

integration/cmd/version/version_test.go

@@ -1,6 +1,7 @@
 package version_test
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"testing"
@@ -13,25 +14,29 @@ import (
 var expectedVersion = fmt.Sprintf("Databricks CLI v%s\n", build.GetInfo().Version)
 
 func TestVersionFlagShort(t *testing.T) {
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "-v")
+	ctx := context.Background()
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "-v")
 	assert.Equal(t, expectedVersion, stdout.String())
 	assert.Equal(t, "", stderr.String())
 }
 
 func TestVersionFlagLong(t *testing.T) {
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "--version")
+	ctx := context.Background()
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "--version")
 	assert.Equal(t, expectedVersion, stdout.String())
 	assert.Equal(t, "", stderr.String())
 }
 
 func TestVersionCommand(t *testing.T) {
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "version")
+	ctx := context.Background()
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "version")
 	assert.Equal(t, expectedVersion, stdout.String())
 	assert.Equal(t, "", stderr.String())
 }
 
 func TestVersionCommandWithJSONOutput(t *testing.T) {
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "version", "--output", "json")
+	ctx := context.Background()
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "version", "--output", "json")
 	assert.NotEmpty(t, stdout.String())
 	assert.Equal(t, "", stderr.String())
 

integration/cmd/workspace/workspace_test.go

@@ -20,7 +20,8 @@ import (
 )
 
 func TestWorkspaceList(t *testing.T) {
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "workspace", "list", "/")
+	ctx := context.Background()
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "workspace", "list", "/")
 	outStr := stdout.String()
 	assert.Contains(t, outStr, "ID")
 	assert.Contains(t, outStr, "Type")
@@ -30,12 +31,14 @@ func TestWorkspaceList(t *testing.T) {
 }
 
 func TestWorkpaceListErrorWhenNoArguments(t *testing.T) {
-	_, _, err := testcli.RequireErrorRun(t, "workspace", "list")
+	ctx := context.Background()
+	_, _, err := testcli.RequireErrorRun(t, ctx, "workspace", "list")
 	assert.Contains(t, err.Error(), "accepts 1 arg(s), received 0")
 }
 
 func TestWorkpaceGetStatusErrorWhenNoArguments(t *testing.T) {
-	_, _, err := testcli.RequireErrorRun(t, "workspace", "get-status")
+	ctx := context.Background()
+	_, _, err := testcli.RequireErrorRun(t, ctx, "workspace", "get-status")
 	assert.Contains(t, err.Error(), "accepts 1 arg(s), received 0")
 }
 
@@ -53,7 +56,7 @@ func TestWorkpaceExportPrintsContents(t *testing.T) {
 	require.NoError(t, err)
 
 	// Run export
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "workspace", "export", path.Join(tmpdir, "file-a"))
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "workspace", "export", path.Join(tmpdir, "file-a"))
 	assert.Equal(t, contents, stdout.String())
 	assert.Equal(t, "", stderr.String())
 }
@@ -122,7 +125,7 @@ func TestExportDir(t *testing.T) {
 	}, "\n")
 
 	// Run Export
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "workspace", "export-dir", sourceDir, targetDir)
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "workspace", "export-dir", sourceDir, targetDir)
 	assert.Equal(t, expectedLogs, stdout.String())
 	assert.Equal(t, "", stderr.String())
 
@@ -150,7 +153,7 @@ func TestExportDirDoesNotOverwrite(t *testing.T) {
 	require.NoError(t, err)
 
 	// Run Export
-	testcli.RequireSuccessfulRun(t, "workspace", "export-dir", sourceDir, targetDir)
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "export-dir", sourceDir, targetDir)
 
 	// Assert file is not overwritten
 	assertLocalFileContents(t, filepath.Join(targetDir, "file-a"), "local content")
@@ -171,7 +174,7 @@ func TestExportDirWithOverwriteFlag(t *testing.T) {
 	require.NoError(t, err)
 
 	// Run Export
-	testcli.RequireSuccessfulRun(t, "workspace", "export-dir", sourceDir, targetDir, "--overwrite")
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "export-dir", sourceDir, targetDir, "--overwrite")
 
 	// Assert file has been overwritten
 	assertLocalFileContents(t, filepath.Join(targetDir, "file-a"), "content from workspace")
@@ -179,7 +182,7 @@ func TestExportDirWithOverwriteFlag(t *testing.T) {
 
 func TestImportDir(t *testing.T) {
 	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "workspace", "import-dir", "./testdata/import_dir", targetDir, "--log-level=debug")
+	stdout, stderr := testcli.RequireSuccessfulRun(t, ctx, "workspace", "import-dir", "./testdata/import_dir", targetDir, "--log-level=debug")
 
 	expectedLogs := strings.Join([]string{
 		fmt.Sprintf("Importing files from %s", "./testdata/import_dir"),
@@ -220,7 +223,7 @@ func TestImportDirDoesNotOverwrite(t *testing.T) {
 	assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "old file")
 	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"old notebook\")")
 
-	testcli.RequireSuccessfulRun(t, "workspace", "import-dir", "./testdata/import_dir", targetDir)
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import-dir", "./testdata/import_dir", targetDir)
 
 	// Assert files are imported
 	assertFilerFileContents(t, ctx, workspaceFiler, "a/b/c/file-b", "file-in-dir")
@@ -248,7 +251,7 @@ func TestImportDirWithOverwriteFlag(t *testing.T) {
 	assertFilerFileContents(t, ctx, workspaceFiler, "file-a", "old file")
 	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"old notebook\")")
 
-	testcli.RequireSuccessfulRun(t, "workspace", "import-dir", "./testdata/import_dir", targetDir, "--overwrite")
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import-dir", "./testdata/import_dir", targetDir, "--overwrite")
 
 	// Assert files are imported
 	assertFilerFileContents(t, ctx, workspaceFiler, "a/b/c/file-b", "file-in-dir")
@@ -270,7 +273,7 @@ func TestExport(t *testing.T) {
 	// Export vanilla file
 	err = f.Write(ctx, "file-a", strings.NewReader("abc"))
 	require.NoError(t, err)
-	stdout, _ := testcli.RequireSuccessfulRun(t, "workspace", "export", path.Join(sourceDir, "file-a"))
+	stdout, _ := testcli.RequireSuccessfulRun(t, ctx, "workspace", "export", path.Join(sourceDir, "file-a"))
 	b, err := io.ReadAll(&stdout)
 	require.NoError(t, err)
 	assert.Equal(t, "abc", string(b))
@@ -278,13 +281,13 @@ func TestExport(t *testing.T) {
 	// Export python notebook
 	err = f.Write(ctx, "pyNotebook.py", strings.NewReader("# Databricks notebook source"))
 	require.NoError(t, err)
-	stdout, _ = testcli.RequireSuccessfulRun(t, "workspace", "export", path.Join(sourceDir, "pyNotebook"))
+	stdout, _ = testcli.RequireSuccessfulRun(t, ctx, "workspace", "export", path.Join(sourceDir, "pyNotebook"))
 	b, err = io.ReadAll(&stdout)
 	require.NoError(t, err)
 	assert.Equal(t, "# Databricks notebook source\n", string(b))
 
 	// Export python notebook as jupyter
-	stdout, _ = testcli.RequireSuccessfulRun(t, "workspace", "export", path.Join(sourceDir, "pyNotebook"), "--format", "JUPYTER")
+	stdout, _ = testcli.RequireSuccessfulRun(t, ctx, "workspace", "export", path.Join(sourceDir, "pyNotebook"), "--format", "JUPYTER")
 	b, err = io.ReadAll(&stdout)
 	require.NoError(t, err)
 	assert.Contains(t, string(b), `"cells":`, "jupyter notebooks contain the cells field")
@@ -300,7 +303,7 @@ func TestExportWithFileFlag(t *testing.T) {
 	// Export vanilla file
 	err = f.Write(ctx, "file-a", strings.NewReader("abc"))
 	require.NoError(t, err)
-	stdout, _ := testcli.RequireSuccessfulRun(t, "workspace", "export", path.Join(sourceDir, "file-a"), "--file", filepath.Join(localTmpDir, "file.txt"))
+	stdout, _ := testcli.RequireSuccessfulRun(t, ctx, "workspace", "export", path.Join(sourceDir, "file-a"), "--file", filepath.Join(localTmpDir, "file.txt"))
 	b, err := io.ReadAll(&stdout)
 	require.NoError(t, err)
 	// Expect nothing to be printed to stdout
@@ -310,14 +313,14 @@ func TestExportWithFileFlag(t *testing.T) {
 	// Export python notebook
 	err = f.Write(ctx, "pyNotebook.py", strings.NewReader("# Databricks notebook source"))
 	require.NoError(t, err)
-	stdout, _ = testcli.RequireSuccessfulRun(t, "workspace", "export", path.Join(sourceDir, "pyNotebook"), "--file", filepath.Join(localTmpDir, "pyNb.py"))
+	stdout, _ = testcli.RequireSuccessfulRun(t, ctx, "workspace", "export", path.Join(sourceDir, "pyNotebook"), "--file", filepath.Join(localTmpDir, "pyNb.py"))
 	b, err = io.ReadAll(&stdout)
 	require.NoError(t, err)
 	assert.Equal(t, "", string(b))
 	assertLocalFileContents(t, filepath.Join(localTmpDir, "pyNb.py"), "# Databricks notebook source\n")
 
 	// Export python notebook as jupyter
-	stdout, _ = testcli.RequireSuccessfulRun(t, "workspace", "export", path.Join(sourceDir, "pyNotebook"), "--format", "JUPYTER", "--file", filepath.Join(localTmpDir, "jupyterNb.ipynb"))
+	stdout, _ = testcli.RequireSuccessfulRun(t, ctx, "workspace", "export", path.Join(sourceDir, "pyNotebook"), "--format", "JUPYTER", "--file", filepath.Join(localTmpDir, "jupyterNb.ipynb"))
 	b, err = io.ReadAll(&stdout)
 	require.NoError(t, err)
 	assert.Equal(t, "", string(b))
@@ -329,13 +332,13 @@ func TestImportFileUsingContentFormatSource(t *testing.T) {
 	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
 
 	//  Content = `print(1)`. Uploaded as a notebook by default
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "pyScript"),
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "pyScript"),
 		"--content", base64.StdEncoding.EncodeToString([]byte("print(1)")), "--language=PYTHON")
 	assertFilerFileContents(t, ctx, workspaceFiler, "pyScript", "print(1)")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "pyScript", workspace.ObjectTypeNotebook)
 
 	// Import with content = `# Databricks notebook source\nprint(1)`. Uploaded as a notebook with the content just being print(1)
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "pyNb"),
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "pyNb"),
 		"--content", base64.StdEncoding.EncodeToString([]byte("`# Databricks notebook source\nprint(1)")),
 		"--language=PYTHON")
 	assertFilerFileContents(t, ctx, workspaceFiler, "pyNb", "print(1)")
@@ -346,19 +349,19 @@ func TestImportFileUsingContentFormatAuto(t *testing.T) {
 	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
 
 	//  Content = `# Databricks notebook source\nprint(1)`. Upload as file if path has no extension.
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "py-nb-as-file"),
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "py-nb-as-file"),
 		"--content", base64.StdEncoding.EncodeToString([]byte("`# Databricks notebook source\nprint(1)")), "--format=AUTO")
 	assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-file", "# Databricks notebook source\nprint(1)")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "py-nb-as-file", workspace.ObjectTypeFile)
 
 	// Content = `# Databricks notebook source\nprint(1)`. Upload as notebook if path has py extension
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "py-nb-as-notebook.py"),
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "py-nb-as-notebook.py"),
 		"--content", base64.StdEncoding.EncodeToString([]byte("`# Databricks notebook source\nprint(1)")), "--format=AUTO")
 	assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-notebook", "# Databricks notebook source\nprint(1)")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "py-nb-as-notebook", workspace.ObjectTypeNotebook)
 
 	// Content = `print(1)`. Upload as file if content is not notebook (even if path has .py extension)
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "not-a-notebook.py"), "--content",
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "not-a-notebook.py"), "--content",
 		base64.StdEncoding.EncodeToString([]byte("print(1)")), "--format=AUTO")
 	assertFilerFileContents(t, ctx, workspaceFiler, "not-a-notebook.py", "print(1)")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "not-a-notebook.py", workspace.ObjectTypeFile)
@@ -366,15 +369,15 @@ func TestImportFileUsingContentFormatAuto(t *testing.T) {
 
 func TestImportFileFormatSource(t *testing.T) {
 	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "pyNotebook"), "--file", "./testdata/import_dir/pyNotebook.py", "--language=PYTHON")
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "pyNotebook"), "--file", "./testdata/import_dir/pyNotebook.py", "--language=PYTHON")
 	assertFilerFileContents(t, ctx, workspaceFiler, "pyNotebook", "# Databricks notebook source\nprint(\"python\")")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "pyNotebook", workspace.ObjectTypeNotebook)
 
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "scalaNotebook"), "--file", "./testdata/import_dir/scalaNotebook.scala", "--language=SCALA")
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "scalaNotebook"), "--file", "./testdata/import_dir/scalaNotebook.scala", "--language=SCALA")
 	assertFilerFileContents(t, ctx, workspaceFiler, "scalaNotebook", "// Databricks notebook source\nprintln(\"scala\")")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "scalaNotebook", workspace.ObjectTypeNotebook)
 
-	_, _, err := testcli.RequireErrorRun(t, "workspace", "import", path.Join(targetDir, "scalaNotebook"), "--file", "./testdata/import_dir/scalaNotebook.scala")
+	_, _, err := testcli.RequireErrorRun(t, ctx, "workspace", "import", path.Join(targetDir, "scalaNotebook"), "--file", "./testdata/import_dir/scalaNotebook.scala")
 	assert.ErrorContains(t, err, "The zip file may not be valid or may be an unsupported version. Hint: Objects imported using format=SOURCE are expected to be zip encoded databricks source notebook(s) by default. Please specify a language using the --language flag if you are trying to import a single uncompressed notebook")
 }
 
@@ -382,18 +385,18 @@ func TestImportFileFormatAuto(t *testing.T) {
 	ctx, workspaceFiler, targetDir := setupWorkspaceImportExportTest(t)
 
 	// Upload as file if path has no extension
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "py-nb-as-file"), "--file", "./testdata/import_dir/pyNotebook.py", "--format=AUTO")
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "py-nb-as-file"), "--file", "./testdata/import_dir/pyNotebook.py", "--format=AUTO")
 	assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-file", "# Databricks notebook source")
 	assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-file", "print(\"python\")")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "py-nb-as-file", workspace.ObjectTypeFile)
 
 	// Upload as notebook if path has extension
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "py-nb-as-notebook.py"), "--file", "./testdata/import_dir/pyNotebook.py", "--format=AUTO")
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "py-nb-as-notebook.py"), "--file", "./testdata/import_dir/pyNotebook.py", "--format=AUTO")
 	assertFilerFileContents(t, ctx, workspaceFiler, "py-nb-as-notebook", "# Databricks notebook source\nprint(\"python\")")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "py-nb-as-notebook", workspace.ObjectTypeNotebook)
 
 	// Upload as file if content is not notebook (even if path has .py extension)
-	testcli.RequireSuccessfulRun(t, "workspace", "import", path.Join(targetDir, "not-a-notebook.py"), "--file", "./testdata/import_dir/file-a", "--format=AUTO")
+	testcli.RequireSuccessfulRun(t, ctx, "workspace", "import", path.Join(targetDir, "not-a-notebook.py"), "--file", "./testdata/import_dir/file-a", "--format=AUTO")
 	assertFilerFileContents(t, ctx, workspaceFiler, "not-a-notebook.py", "hello, world")
 	assertWorkspaceFileType(t, ctx, workspaceFiler, "not-a-notebook.py", workspace.ObjectTypeFile)
 }

integration/libs/git/git_fetch_test.go

@@ -5,11 +5,10 @@ import (
 	"os/exec"
 	"path"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/databricks/cli/internal/acc"
-	"github.com/databricks/cli/internal/testcli"
-	"github.com/databricks/cli/internal/testutil"
 	"github.com/databricks/cli/libs/dbr"
 	"github.com/databricks/cli/libs/git"
 	"github.com/stretchr/testify/assert"
@@ -39,19 +38,18 @@ func assertSparseGitInfo(t *testing.T, expectedRoot string, info git.RepositoryI
 	assert.Equal(t, expectedRoot, info.WorktreeRoot)
 }
 
+func ensureWorkspacePrefix(root string) string {
+	// The fixture helper doesn't include /Workspace, so include it here.
+	if !strings.HasPrefix(root, "/Workspace/") {
+		return path.Join("/Workspace", root)
+	}
+	return root
+}
+
 func TestFetchRepositoryInfoAPI_FromRepo(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
-	me, err := wt.W.CurrentUser.Me(ctx)
-	require.NoError(t, err)
+	targetPath := ensureWorkspacePrefix(acc.TemporaryRepo(wt, examplesRepoUrl))
 
-	targetPath := testutil.RandomName(path.Join("/Workspace/Users", me.UserName, "/testing-clone-bundle-examples-"))
-	stdout, stderr := testcli.RequireSuccessfulRun(t, "repos", "create", examplesRepoUrl, examplesRepoProvider, "--path", targetPath)
-	t.Cleanup(func() {
-		testcli.RequireSuccessfulRun(t, "repos", "delete", targetPath)
-	})
-
-	assert.Empty(t, stderr.String())
-	assert.NotEmpty(t, stdout.String())
 	ctx = dbr.MockRuntime(ctx, true)
 
 	for _, inputPath := range []string{
@@ -68,16 +66,12 @@ func TestFetchRepositoryInfoAPI_FromRepo(t *testing.T) {
 
 func TestFetchRepositoryInfoAPI_FromNonRepo(t *testing.T) {
 	ctx, wt := acc.WorkspaceTest(t)
-	me, err := wt.W.CurrentUser.Me(ctx)
+	rootPath := ensureWorkspacePrefix(acc.TemporaryWorkspaceDir(wt, "testing-nonrepo-"))
+
+	// Create directory inside this root path (this is cleaned up as part of the root path).
+	err := wt.W.Workspace.MkdirsByPath(ctx, path.Join(rootPath, "a/b/c"))
 	require.NoError(t, err)
 
-	rootPath := testutil.RandomName(path.Join("/Workspace/Users", me.UserName, "testing-nonrepo-"))
-	_, stderr := testcli.RequireSuccessfulRun(t, "workspace", "mkdirs", path.Join(rootPath, "a/b/c"))
-	t.Cleanup(func() {
-		testcli.RequireSuccessfulRun(t, "workspace", "delete", "--recursive", rootPath)
-	})
-
-	assert.Empty(t, stderr.String())
 	ctx = dbr.MockRuntime(ctx, true)
 
 	tests := []struct {

internal/testcli/runner.go

@@ -281,16 +281,7 @@ func (r *Runner) RunAndParseJSON(v any) {
 	require.NoError(r, err)
 }
 
-func NewRunner(t testutil.TestingT, args ...string) *Runner {
-	return &Runner{
-		TestingT: t,
-
-		ctx:  context.Background(),
-		args: args,
-	}
-}
-
-func NewRunnerWithContext(t testutil.TestingT, ctx context.Context, args ...string) *Runner {
+func NewRunner(t testutil.TestingT, ctx context.Context, args ...string) *Runner {
 	return &Runner{
 		TestingT: t,
 
@@ -299,16 +290,16 @@ func NewRunnerWithContext(t testutil.TestingT, ctx context.Context, args ...stri
 	}
 }
 
-func RequireSuccessfulRun(t testutil.TestingT, args ...string) (bytes.Buffer, bytes.Buffer) {
+func RequireSuccessfulRun(t testutil.TestingT, ctx context.Context, args ...string) (bytes.Buffer, bytes.Buffer) {
 	t.Logf("run args: [%s]", strings.Join(args, ", "))
-	r := NewRunner(t, args...)
+	r := NewRunner(t, ctx, args...)
 	stdout, stderr, err := r.Run()
 	require.NoError(t, err)
 	return stdout, stderr
 }
 
-func RequireErrorRun(t testutil.TestingT, args ...string) (bytes.Buffer, bytes.Buffer, error) {
-	r := NewRunner(t, args...)
+func RequireErrorRun(t testutil.TestingT, ctx context.Context, args ...string) (bytes.Buffer, bytes.Buffer, error) {
+	r := NewRunner(t, ctx, args...)
 	stdout, stderr, err := r.Run()
 	require.Error(t, err)
 	return stdout, stderr, err

libs/cmdio/io.go

@@ -31,9 +31,9 @@ type cmdIO struct {
 	err            io.Writer
 }
 
-func NewIO(outputFormat flags.Output, in io.Reader, out, err io.Writer, headerTemplate, template string) *cmdIO {
+func NewIO(ctx context.Context, outputFormat flags.Output, in io.Reader, out, err io.Writer, headerTemplate, template string) *cmdIO {
 	// The check below is similar to color.NoColor but uses the specified err writer.
-	dumb := os.Getenv("NO_COLOR") != "" || os.Getenv("TERM") == "dumb"
+	dumb := env.Get(ctx, "NO_COLOR") != "" || env.Get(ctx, "TERM") == "dumb"
 	if f, ok := err.(*os.File); ok && !dumb {
 		dumb = !isatty.IsTerminal(f.Fd()) && !isatty.IsCygwinTerminal(f.Fd())
 	}

libs/cmdio/render_test.go

@@ -171,8 +171,9 @@ func TestRender(t *testing.T) {
 	for _, c := range testCases {
 		t.Run(c.name, func(t *testing.T) {
 			output := &bytes.Buffer{}
-			cmdIO := NewIO(c.outputFormat, nil, output, output, c.headerTemplate, c.template)
-			ctx := InContext(context.Background(), cmdIO)
+			ctx := context.Background()
+			cmdIO := NewIO(ctx, c.outputFormat, nil, output, output, c.headerTemplate, c.template)
+			ctx = InContext(ctx, cmdIO)
 			var err error
 			if vv, ok := c.v.(listing.Iterator[*provisioning.Workspace]); ok {
 				err = RenderIterator(ctx, vv)

libs/databrickscfg/cfgpickers/clusters_test.go

@@ -115,7 +115,7 @@ func TestFirstCompatibleCluster(t *testing.T) {
 	w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg)))
 
 	ctx := context.Background()
-	ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
+	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
 	clusterID, err := AskForCluster(ctx, w, WithDatabricksConnect("13.1"))
 	require.NoError(t, err)
 	require.Equal(t, "bcd-id", clusterID)
@@ -162,7 +162,7 @@ func TestNoCompatibleClusters(t *testing.T) {
 	w := databricks.Must(databricks.NewWorkspaceClient((*databricks.Config)(cfg)))
 
 	ctx := context.Background()
-	ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
+	ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
 	_, err := AskForCluster(ctx, w, WithDatabricksConnect("13.1"))
 	require.Equal(t, ErrNoCompatibleClusters, err)
 }

libs/template/helpers_test.go

@@ -162,7 +162,7 @@ func TestWorkspaceHost(t *testing.T) {
 
 func TestWorkspaceHostNotConfigured(t *testing.T) {
 	ctx := context.Background()
-	cmd := cmdio.NewIO(flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "", "template")
+	cmd := cmdio.NewIO(ctx, flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "", "template")
 	ctx = cmdio.InContext(ctx, cmd)
 
 	w := &databricks.WorkspaceClient{