Skip to content

Commit 70b7bbf

Browse files
pieternpietern
authored
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
1 parent d929ea3 commit 70b7bbf

39 files changed

+274
-223
lines changed

‎bundle/run/job_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ func TestJobRunnerRestart(t *testing.T) {
160160
m := mocks.NewMockWorkspaceClient(t)
161161
b.SetWorkpaceClient(m.WorkspaceClient)
162162
ctx := context.Background()
163-
ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", ""))
163+
ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", ""))
164164
ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
165165

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

237237
jobApi := m.GetMockJobsAPI()

‎bundle/run/pipeline_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ func TestPipelineRunnerRestart(t *testing.T) {
7676
}
7777
b.SetWorkpaceClient(m.WorkspaceClient)
7878
ctx := context.Background()
79-
ctx = cmdio.InContext(ctx, cmdio.NewIO(flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
79+
ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputText, &bytes.Buffer{}, &bytes.Buffer{}, &bytes.Buffer{}, "", "..."))
8080
ctx = cmdio.NewContext(ctx, cmdio.NewLogger(flags.ModeAppend))
8181

8282
mockWait := &pipelines.WaitGetPipelineIdle[struct{}]{

‎cmd/labs/installed_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import (
1111
func TestListsInstalledProjects(t *testing.T) {
1212
ctx := context.Background()
1313
ctx = env.WithUserHomeDir(ctx, "project/testdata/installed-in-home")
14-
r := testcli.NewRunnerWithContext(t, ctx, "labs", "installed")
14+
r := testcli.NewRunner(t, ctx, "labs", "installed")
1515
r.RunAndExpectOutput(`
1616
Name Description Version
1717
blueprint Blueprint Project v0.3.15

‎cmd/labs/list_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import (
1212
func TestListingWorks(t *testing.T) {
1313
ctx := context.Background()
1414
ctx = env.WithUserHomeDir(ctx, "project/testdata/installed-in-home")
15-
c := testcli.NewRunnerWithContext(t, ctx, "labs", "list")
15+
c := testcli.NewRunner(t, ctx, "labs", "list")
1616
stdout, _, err := c.Run()
1717
require.NoError(t, err)
1818
require.Contains(t, stdout.String(), "ucx")

‎cmd/labs/project/command_test.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ func devEnvContext(t *testing.T) context.Context {
3030

3131
func TestRunningBlueprintEcho(t *testing.T) {
3232
ctx := devEnvContext(t)
33-
r := testcli.NewRunnerWithContext(t, ctx, "labs", "blueprint", "echo")
33+
r := testcli.NewRunner(t, ctx, "labs", "blueprint", "echo")
3434
var out echoOut
3535
r.RunAndParseJSON(&out)
3636
assert.Equal(t, "echo", out.Command)
@@ -41,14 +41,14 @@ func TestRunningBlueprintEcho(t *testing.T) {
4141

4242
func TestRunningBlueprintEchoProfileWrongOverride(t *testing.T) {
4343
ctx := devEnvContext(t)
44-
r := testcli.NewRunnerWithContext(t, ctx, "labs", "blueprint", "echo", "--profile", "workspace-profile")
44+
r := testcli.NewRunner(t, ctx, "labs", "blueprint", "echo", "--profile", "workspace-profile")
4545
_, _, err := r.Run()
4646
assert.ErrorIs(t, err, databricks.ErrNotAccountClient)
4747
}
4848

4949
func TestRunningCommand(t *testing.T) {
5050
ctx := devEnvContext(t)
51-
r := testcli.NewRunnerWithContext(t, ctx, "labs", "blueprint", "foo")
51+
r := testcli.NewRunner(t, ctx, "labs", "blueprint", "foo")
5252
r.WithStdin()
5353
defer r.CloseStdin()
5454

@@ -60,7 +60,7 @@ func TestRunningCommand(t *testing.T) {
6060

6161
func TestRenderingTable(t *testing.T) {
6262
ctx := devEnvContext(t)
63-
r := testcli.NewRunnerWithContext(t, ctx, "labs", "blueprint", "table")
63+
r := testcli.NewRunner(t, ctx, "labs", "blueprint", "table")
6464
r.RunAndExpectOutput(`
6565
Key Value
6666
First Second

‎cmd/labs/project/installer_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ func TestInstallerWorksForReleases(t *testing.T) {
236236
// │ │ │ └── site-packages
237237
// │ │ │ ├── ...
238238
// │ │ │ ├── distutils-precedence.pth
239-
r := testcli.NewRunnerWithContext(t, ctx, "labs", "install", "blueprint", "--debug")
239+
r := testcli.NewRunner(t, ctx, "labs", "install", "blueprint", "--debug")
240240
r.RunAndExpectOutput("setting up important infrastructure")
241241
}
242242

@@ -356,7 +356,7 @@ account_id = abc
356356
// └── databrickslabs-blueprint-releases.json
357357

358358
// `databricks labs install .` means "verify this installer i'm developing does work"
359-
r := testcli.NewRunnerWithContext(t, ctx, "labs", "install", ".")
359+
r := testcli.NewRunner(t, ctx, "labs", "install", ".")
360360
r.WithStdin()
361361
defer r.CloseStdin()
362362

@@ -426,7 +426,7 @@ func TestUpgraderWorksForReleases(t *testing.T) {
426426
ctx = env.Set(ctx, "DATABRICKS_CLUSTER_ID", "installer-cluster")
427427
ctx = env.Set(ctx, "DATABRICKS_WAREHOUSE_ID", "installer-warehouse")
428428

429-
r := testcli.NewRunnerWithContext(t, ctx, "labs", "upgrade", "blueprint")
429+
r := testcli.NewRunner(t, ctx, "labs", "upgrade", "blueprint")
430430
r.RunAndExpectOutput("setting up important infrastructure")
431431

432432
// Check if the stub was called with the 'python -m pip install' command

‎cmd/root/io.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,9 @@ func (f *outputFlag) initializeIO(cmd *cobra.Command) error {
4545
headerTemplate = cmd.Annotations["headerTemplate"]
4646
}
4747

48-
cmdIO := cmdio.NewIO(f.output, cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr(), headerTemplate, template)
49-
ctx := cmdio.InContext(cmd.Context(), cmdIO)
48+
ctx := cmd.Context()
49+
cmdIO := cmdio.NewIO(ctx, f.output, cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr(), headerTemplate, template)
50+
ctx = cmdio.InContext(ctx, cmdIO)
5051
cmd.SetContext(ctx)
5152
return nil
5253
}

‎integration/bundle/artifacts_test.go‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/databricks/cli/internal/acc"
1616
"github.com/databricks/cli/internal/testcli"
1717
"github.com/databricks/cli/internal/testutil"
18+
"github.com/databricks/cli/libs/env"
1819
"github.com/databricks/databricks-sdk-go/service/catalog"
1920
"github.com/databricks/databricks-sdk-go/service/compute"
2021
"github.com/databricks/databricks-sdk-go/service/jobs"
@@ -253,8 +254,8 @@ func TestUploadArtifactFileToVolumeThatDoesNotExist(t *testing.T) {
253254
})
254255
require.NoError(t, err)
255256

256-
t.Setenv("BUNDLE_ROOT", bundleRoot)
257-
stdout, stderr, err := testcli.RequireErrorRun(t, "bundle", "deploy")
257+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
258+
stdout, stderr, err := testcli.RequireErrorRun(t, ctx, "bundle", "deploy")
258259

259260
assert.Error(t, err)
260261
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) {
290291
})
291292
require.NoError(t, err)
292293

293-
t.Setenv("BUNDLE_ROOT", bundleRoot)
294-
stdout, stderr, err := testcli.RequireErrorRun(t, "bundle", "deploy")
294+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
295+
stdout, stderr, err := testcli.RequireErrorRun(t, ctx, "bundle", "deploy")
295296

296297
assert.Error(t, err)
297298
assert.Equal(t, fmt.Sprintf(`Error: volume /Volumes/main/%s/my_volume does not exist: Not Found

‎integration/bundle/bind_resource_test.go‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"github.com/databricks/cli/internal/acc"
1010
"github.com/databricks/cli/internal/testcli"
1111
"github.com/databricks/cli/internal/testutil"
12+
"github.com/databricks/cli/libs/env"
1213
"github.com/databricks/databricks-sdk-go"
1314
"github.com/databricks/databricks-sdk-go/service/jobs"
1415
"github.com/google/uuid"
@@ -35,8 +36,8 @@ func TestBindJobToExistingJob(t *testing.T) {
3536
require.NoError(t, err)
3637
})
3738

38-
t.Setenv("BUNDLE_ROOT", bundleRoot)
39-
c := testcli.NewRunner(t, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId), "--auto-approve")
39+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
40+
c := testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId), "--auto-approve")
4041
_, _, err = c.Run()
4142
require.NoError(t, err)
4243

@@ -58,7 +59,7 @@ func TestBindJobToExistingJob(t *testing.T) {
5859
require.Equal(t, job.Settings.Name, fmt.Sprintf("test-job-basic-%s", uniqueId))
5960
require.Contains(t, job.Settings.Tasks[0].SparkPythonTask.PythonFile, "hello_world.py")
6061

61-
c = testcli.NewRunner(t, "bundle", "deployment", "unbind", "foo")
62+
c = testcli.NewRunner(t, ctx, "bundle", "deployment", "unbind", "foo")
6263
_, _, err = c.Run()
6364
require.NoError(t, err)
6465

@@ -99,9 +100,9 @@ func TestAbortBind(t *testing.T) {
99100
})
100101

101102
// Bind should fail because prompting is not possible.
102-
t.Setenv("BUNDLE_ROOT", bundleRoot)
103-
t.Setenv("TERM", "dumb")
104-
c := testcli.NewRunner(t, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId))
103+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
104+
ctx = env.Set(ctx, "TERM", "dumb")
105+
c := testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "foo", fmt.Sprint(jobId))
105106

106107
// Expect error suggesting to use --auto-approve
107108
_, _, err = c.Run()
@@ -147,8 +148,8 @@ func TestGenerateAndBind(t *testing.T) {
147148
}
148149
})
149150

150-
t.Setenv("BUNDLE_ROOT", bundleRoot)
151-
c := testcli.NewRunnerWithContext(t, ctx, "bundle", "generate", "job",
151+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
152+
c := testcli.NewRunner(t, ctx, "bundle", "generate", "job",
152153
"--key", "test_job_key",
153154
"--existing-job-id", fmt.Sprint(jobId),
154155
"--config-dir", filepath.Join(bundleRoot, "resources"),
@@ -164,7 +165,7 @@ func TestGenerateAndBind(t *testing.T) {
164165

165166
require.Len(t, matches, 1)
166167

167-
c = testcli.NewRunner(t, "bundle", "deployment", "bind", "test_job_key", fmt.Sprint(jobId), "--auto-approve")
168+
c = testcli.NewRunner(t, ctx, "bundle", "deployment", "bind", "test_job_key", fmt.Sprint(jobId), "--auto-approve")
168169
_, _, err = c.Run()
169170
require.NoError(t, err)
170171

‎integration/bundle/deploy_test.go‎

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import (
1414
"github.com/databricks/cli/internal/acc"
1515
"github.com/databricks/cli/internal/testcli"
1616
"github.com/databricks/cli/internal/testutil"
17+
"github.com/databricks/cli/libs/env"
1718
"github.com/databricks/databricks-sdk-go"
1819
"github.com/databricks/databricks-sdk-go/apierr"
1920
"github.com/databricks/databricks-sdk-go/service/catalog"
@@ -118,9 +119,9 @@ func TestBundleDeployUcSchemaFailsWithoutAutoApprove(t *testing.T) {
118119
require.NoError(t, err)
119120

120121
// Redeploy the bundle
121-
t.Setenv("BUNDLE_ROOT", bundleRoot)
122-
t.Setenv("TERM", "dumb")
123-
c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock")
122+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
123+
ctx = env.Set(ctx, "TERM", "dumb")
124+
c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock")
124125
stdout, stderr, err := c.Run()
125126

126127
assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -162,9 +163,9 @@ func TestBundlePipelineDeleteWithoutAutoApprove(t *testing.T) {
162163
require.NoError(t, err)
163164

164165
// Redeploy the bundle. Expect it to fail because deleting the pipeline requires --auto-approve.
165-
t.Setenv("BUNDLE_ROOT", bundleRoot)
166-
t.Setenv("TERM", "dumb")
167-
c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock")
166+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
167+
ctx = env.Set(ctx, "TERM", "dumb")
168+
c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock")
168169
stdout, stderr, err := c.Run()
169170

170171
assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -201,9 +202,9 @@ func TestBundlePipelineRecreateWithoutAutoApprove(t *testing.T) {
201202
require.Equal(t, pipelineName, pipeline.Name)
202203

203204
// Redeploy the bundle, pointing the DLT pipeline to a different UC catalog.
204-
t.Setenv("BUNDLE_ROOT", bundleRoot)
205-
t.Setenv("TERM", "dumb")
206-
c := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock", "--var=\"catalog=whatever\"")
205+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
206+
ctx = env.Set(ctx, "TERM", "dumb")
207+
c := testcli.NewRunner(t, ctx, "bundle", "deploy", "--force-lock", "--var=\"catalog=whatever\"")
207208
stdout, stderr, err := c.Run()
208209

209210
assert.EqualError(t, err, root.ErrAlreadyPrinted.Error())
@@ -235,7 +236,7 @@ func TestDeployBasicBundleLogs(t *testing.T) {
235236
currentUser, err := wt.W.CurrentUser.Me(ctx)
236237
require.NoError(t, err)
237238

238-
stdout, stderr := blackBoxRun(t, root, "bundle", "deploy")
239+
stdout, stderr := blackBoxRun(t, ctx, root, "bundle", "deploy")
239240
assert.Equal(t, strings.Join([]string{
240241
fmt.Sprintf("Uploading bundle files to /Workspace/Users/%s/.bundle/%s/files...", currentUser.UserName, uniqueId),
241242
"Deploying resources...",
@@ -282,9 +283,9 @@ func TestDeployUcVolume(t *testing.T) {
282283
assert.Equal(t, []catalog.Privilege{catalog.PrivilegeWriteVolume}, grants.PrivilegeAssignments[0].Privileges)
283284

284285
// Recreation of the volume without --auto-approve should fail since prompting is not possible
285-
t.Setenv("TERM", "dumb")
286-
t.Setenv("BUNDLE_ROOT", bundleRoot)
287-
stdout, stderr, err := testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}").Run()
286+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
287+
ctx = env.Set(ctx, "TERM", "dumb")
288+
stdout, stderr, err := testcli.NewRunner(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}").Run()
288289
assert.Error(t, err)
289290
assert.Contains(t, stderr.String(), `This action will result in the deletion or recreation of the following volumes.
290291
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:
294295
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")
295296

296297
// Successfully recreate the volume with --auto-approve
297-
t.Setenv("TERM", "dumb")
298-
t.Setenv("BUNDLE_ROOT", bundleRoot)
299-
_, _, err = testcli.NewRunnerWithContext(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}", "--auto-approve").Run()
298+
ctx = env.Set(ctx, "BUNDLE_ROOT", bundleRoot)
299+
ctx = env.Set(ctx, "TERM", "dumb")
300+
_, _, err = testcli.NewRunner(t, ctx, "bundle", "deploy", "--var=schema_name=${resources.schemas.schema2.name}", "--auto-approve").Run()
300301
assert.NoError(t, err)
301302

302303
// Assert the volume is updated successfully

0 commit comments

Comments
 (0)