Remove calls to `t.Setenv` from integration tests (#2018)

## Changes

The `Setenv` helper function configures an environment variable and
resets it to its original value when exiting the test scope. It is
incompatible with running tests in parallel because it modifies
process-wide state. The `libs/env` package defines functions to interact
with the environment but records `Setenv` calls on a `context.Context`.
This enables us to override/specialize the environment scoped to a
context.

Pre-requisites for removing the `t.Setenv` calls:
* Make `cmdio.NewIO` accept a context and use it with `libs/env`
* Make all `internal/testcli` functions use a context

The rest of this change:
* Modifies integration tests to initialize a context to use if there
wasn't already one
* Updates `t.Setenv` calls to use `env.Set`

## Tests

n/a

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 := cmd.Context()
-	ctx := cmdio.InContext(cmd.Context(), cmdIO)
+	cmdIO := cmdio.NewIO(ctx, f.output, cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr(), headerTemplate, template)
+	ctx = cmdio.InContext(ctx, cmdIO)
 	cmd.SetContext(ctx)
 	return nil
 }

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"
@@ -253,8 +254,8 @@ func TestUploadArtifactFileToVolumeThatDoesNotExist(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	stdout, stderr, err := testcli.RequireErrorRun(t, "bundle", "deploy")
+	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
@@ -290,8 +291,8 @@ func TestUploadArtifactToVolumeNotYetDeployed(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	stdout, stderr, err := testcli.RequireErrorRun(t, "bundle", "deploy")
+	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/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"
@@ -35,8 +36,8 @@ func TestBindJobToExistingJob(t *testing.T) {
 		require.NoError(t, err)
 	})
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	c := testcli.NewRunner(t, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId), "--auto-approve")
+	c := testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId), "--auto-approve")
 	_, _, err = c.Run()
 	require.NoError(t, err)
 
@@ -58,7 +59,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)
 
@@ -99,9 +100,9 @@ func TestAbortBind(t *testing.T) {
 	})
 
 	// Bind should fail because prompting is not possible.
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	t.Setenv("TERM", "dumb")
+	ctx = env.Set(ctx, "TERM", "dumb")
-	c := testcli.NewRunner(t, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId))
+	c := testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId))
 
 	// Expect error suggesting to use --auto-approve
 	_, _, err = c.Run()
@@ -147,8 +148,8 @@ func TestGenerateAndBind(t *testing.T) {
 		}
 	})
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "generate", "job",
+	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,7 +165,7 @@ 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)
 

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"
@@ -118,9 +119,9 @@ func TestBundleDeployUcSchemaFailsWithoutAutoApprove(t *testing.T) {
 	require.NoError(t, err)
 
 	// Redeploy the bundle
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	t.Setenv("TERM", "dumb")
+	ctx = env.Set(ctx, "TERM", "dumb")
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock")
+	c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock")
 	stdout, stderr, err := c.Run()
 
 	assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -162,9 +163,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)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	t.Setenv("TERM", "dumb")
+	ctx = env.Set(ctx, "TERM", "dumb")
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock")
+	c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock")
 	stdout, stderr, err := c.Run()
 
 	assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -201,9 +202,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)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	t.Setenv("TERM", "dumb")
+	ctx = env.Set(ctx, "TERM", "dumb")
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock", "--var=\"catalog=whatever\"")
+	c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock", "--var=\"catalog=whatever\"")
 	stdout, stderr, err := c.Run()
 
 	assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -235,7 +236,7 @@ func TestDeployBasicBundleLogs(t *testing.T) {
 	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...",
@@ -282,9 +283,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")
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "TERM", "dumb")
-	stdout, stderr, err := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}").Run()
+	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 +295,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")
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "TERM", "dumb")
-	_, _, err = testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}", "--auto-approve").Run()
+	_, _, 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/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"
 )
@@ -26,7 +27,7 @@ func TestFilesAreSyncedCorrectlyWhenNoSnapshot(t *testing.T) {
 	})
 	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)

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"
@@ -35,8 +36,8 @@ func TestGenerateFromExistingJobAndDeploy(t *testing.T) {
 		gt.destroyJob(ctx, jobId)
 	})
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "generate", "job",
+	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"))

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"
@@ -34,8 +35,8 @@ func TestGenerateFromExistingPipelineAndDeploy(t *testing.T) {
 		gt.destroyPipeline(ctx, pipelineId)
 	})
 
-	t.Setenv("BUNDLE_ROOT", bundleRoot)
+	ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
-	c := testcli.NewRunnerWithContext(t, ctx, "bundle", "generate", "pipeline",
+	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"))

integration/bundle/helpers_test.go

@@ -40,7 +40,7 @@ func initTestTemplateWithBundleRoot(t testutil.TestingT, ctx context.Context, te
 	}
 
 	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)
@@ -65,7 +65,7 @@ func writeConfigFile(t testutil.TestingT, config map[string]any) (string, error)
 
 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
 }
@@ -85,7 +85,7 @@ func unmarshalConfig(t testutil.TestingT, data []byte) *bundle.Bundle {
 
 func deployBundle(t testutil.TestingT, ctx context.Context, path string) error {
 	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
 }
@@ -93,7 +93,7 @@ func deployBundle(t testutil.TestingT, ctx context.Context, path string) error {
 func deployBundleWithArgs(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
 }
@@ -102,7 +102,7 @@ func deployBundleWithFlags(t testutil.TestingT, ctx context.Context, path 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
 }
@@ -111,7 +111,7 @@ func runResource(t testutil.TestingT, ctx context.Context, path, key string) (st
 	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,14 +123,14 @@ 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 {
 	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
 }
@@ -143,16 +143,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/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 {
+func setupSyncTest(t *testing.T, args ...string) (context.Context, *syncTest) {
-	_, wt := acc.WorkspaceTest(t)
+	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()
+	ctx, assertSync := setupSyncTest(t, "--full", "--watch")
-	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()
+	ctx, assertSync := setupSyncTest(t, "--watch")
-	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()
+	ctx, assertSync := setupSyncTest(t, "--watch")
-	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()
+	ctx, assertSync := setupSyncTest(t, "--watch")
-	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()
+	ctx, assertSync := setupSyncTest(t, "--watch")
-	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()
+	ctx, assertSync := setupSyncTest(t, "--watch")
-	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()
+	ctx, assertSync := setupSyncTest(t, "--watch")
-	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()
+	ctx, assertSync := setupSyncTest(t, "--watch")
-	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()
+	ctx, assertSync := setupSyncTest(t, "--watch")
-	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)
+	targetPath := ensureWorkspacePrefix(acc.TemporaryRepo(wt, examplesRepoUrl))
-	require.NoError(t, err)
 
-	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 {
+func NewRunner(t testutil.TestingT, ctx context.Context, args ...string) *Runner {
-	return &Runner{
-		TestingT: t,
-
-		ctx:  context.Background(),
-		args: args,
-	}
-}
-
-func NewRunnerWithContext(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) {
+func RequireErrorRun(t testutil.TestingT, ctx context.Context, args ...string) (bytes.Buffer, bytes.Buffer, error) {
-	r := NewRunner(t, args...)
+	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 := context.Background()
-			ctx := InContext(context.Background(), cmdIO)
+			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{