Remove calls to t.Setenv from integration tests

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
 }

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)
-	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
@@ -290,8 +291,8 @@ func TestUploadArtifactToVolumeNotYetDeployed(t *testing.T) {
 	})
 	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/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)
-	c := testcli.NewRunner(t, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId), "--auto-approve")
+	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)
 
@@ -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)
-	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()
@@ -147,8 +148,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,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)
-	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())
@@ -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)
-	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())
@@ -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)
-	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())
@@ -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")
-	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 +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")
-	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/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)
-	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"))

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