Use t.Context() instead of context.Background() in tests

AGENTS.md

@@ -134,7 +134,7 @@ func TestApplySomeChangeReturnsDiagnostics(t *testing.T) {
 }
 
 func TestApplySomeChangeFixesThings(t *testing.T) {
-	ctx := context.Background()
+	ctx := t.Context()
 	b, err := ...some operation...
 	require.NoError(t, err)
 	...
@@ -202,7 +202,7 @@ log.Errorf(ctx, "...")
 ```
 
 Note that the 'ctx' variable here is something that should be passed in as
-an argument by the caller. We should not use context.Background() like we do in tests.
+an argument by the caller. In tests, use `t.Context()` (or `b.Context()` for benchmarks) instead of `context.Background()`. This is enforced by a lint rule.
 
 Use cmdio.LogString to print to stdout:
 

acceptance/acceptance_test.go

@@ -567,7 +567,7 @@ func runTest(t *testing.T,
 		timeout = time.Duration(float64(timeout) * config.TimeoutCIMultiplier)
 	}
 
-	ctx, cancelFunc := context.WithTimeout(context.Background(), timeout)
+	ctx, cancelFunc := context.WithTimeout(t.Context(), timeout)
 	defer cancelFunc()
 	args := []string{"bash", "-euo", "pipefail", EntryPointScript}
 	cmd := exec.CommandContext(ctx, args[0], args[1:]...)

acceptance/dbr_test.go

@@ -325,7 +325,7 @@ func runDbrTests(ctx context.Context, t *testing.T, w *databricks.WorkspaceClien
 
 // runDbrAcceptanceTests is the main entry point for running DBR acceptance tests.
 func runDbrAcceptanceTests(t *testing.T, config dbrTestConfig) {
-	ctx := context.Background()
+	ctx := t.Context()
 	uniqueID := uuid.New().String()
 
 	w, f, testDir := setupDbrTestDir(ctx, t, uniqueID)

acceptance/internal/cmd_server.go

@@ -1,7 +1,6 @@
 package internal
 
 import (
-	"context"
 	"encoding/json"
 	"os"
 	"strings"
@@ -27,7 +26,7 @@ func StartCmdServer(t *testing.T) *testserver.Server {
 		// Change current working directory to match the callsite.
 		defer chdir(t, q.Get("cwd"))()
 
-		c := testcli.NewRunner(t, context.Background(), args...)
+		c := testcli.NewRunner(t, t.Context(), args...)
 		c.Verbose = false
 		stdout, stderr, err := c.Run()
 		result := map[string]any{

acceptance/internal/prepare_server.go

@@ -1,7 +1,6 @@
 package internal
 
 import (
-	"context"
 	"encoding/json"
 	"fmt"
 	"net/http"
@@ -86,7 +85,7 @@ func PrepareServerAndClient(t *testing.T, config TestConfig, logRequests bool, o
 		w, err := databricks.NewWorkspaceClient()
 		require.NoError(t, err)
 
-		user, err := w.CurrentUser.Me(context.Background())
+		user, err := w.CurrentUser.Me(t.Context())
 		require.NoError(t, err, "Failed to get current user")
 
 		cfg := w.Config

bundle/apps/validate_test.go

@@ -1,7 +1,6 @@
 package apps
 
 import (
-	"context"
 	"path/filepath"
 	"testing"
 
@@ -50,7 +49,7 @@ func TestAppsValidateSameSourcePath(t *testing.T) {
 
 	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(tmpDir, "databricks.yml")}})
 
-	diags := bundle.ApplySeq(context.Background(), b, mutator.TranslatePaths(), Validate())
+	diags := bundle.ApplySeq(t.Context(), b, mutator.TranslatePaths(), Validate())
 	require.Len(t, diags, 1)
 	require.Equal(t, "Duplicate app source code path", diags[0].Summary)
 	require.Contains(t, diags[0].Detail, "has the same source code path as app resource")
@@ -86,7 +85,7 @@ func TestAppsValidateBothSourceCodePathAndGitSource(t *testing.T) {
 
 	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(tmpDir, "databricks.yml")}})
 
-	diags := bundle.ApplySeq(context.Background(), b, mutator.TranslatePaths(), Validate())
+	diags := bundle.ApplySeq(t.Context(), b, mutator.TranslatePaths(), Validate())
 	require.Len(t, diags, 1)
 	require.Equal(t, "Both source_code_path and git_source fields are set", diags[0].Summary)
 	require.Contains(t, diags[0].Detail, "should have either source_code_path or git_source field, not both")

bundle/bundle_test.go

@@ -1,7 +1,6 @@
 package bundle
 
 import (
-	"context"
 	"io/fs"
 	"os"
 	"path/filepath"
@@ -22,7 +21,7 @@ import (
 
 // mustLoad calls MustLoad and returns the bundle and collected diagnostics
 func mustLoad(t *testing.T) (*Bundle, []diag.Diagnostic) {
-	ctx := logdiag.InitContext(context.Background())
+	ctx := logdiag.InitContext(t.Context())
 	logdiag.SetCollect(ctx, true)
 	b := MustLoad(ctx)
 	diags := logdiag.FlushCollected(ctx)
@@ -31,27 +30,27 @@ func mustLoad(t *testing.T) (*Bundle, []diag.Diagnostic) {
 
 // tryLoad calls TryLoad and returns the bundle and collected diagnostics
 func tryLoad(t *testing.T) (*Bundle, []diag.Diagnostic) {
-	ctx := logdiag.InitContext(context.Background())
+	ctx := logdiag.InitContext(t.Context())
 	logdiag.SetCollect(ctx, true)
 	b := TryLoad(ctx)
 	diags := logdiag.FlushCollected(ctx)
 	return b, diags
 }
 
 func TestLoadNotExists(t *testing.T) {
-	b, err := Load(context.Background(), "/doesntexist")
+	b, err := Load(t.Context(), "/doesntexist")
 	assert.ErrorIs(t, err, fs.ErrNotExist)
 	assert.Nil(t, b)
 }
 
 func TestLoadExists(t *testing.T) {
-	b, err := Load(context.Background(), "./tests/basic")
+	b, err := Load(t.Context(), "./tests/basic")
 	assert.NoError(t, err)
 	assert.NotNil(t, b)
 }
 
 func TestBundleLocalStateDir(t *testing.T) {
-	ctx := context.Background()
+	ctx := t.Context()
 	projectDir := t.TempDir()
 	f1, err := os.Create(filepath.Join(projectDir, "databricks.yml"))
 	require.NoError(t, err)
@@ -75,7 +74,7 @@ func TestBundleLocalStateDir(t *testing.T) {
 }
 
 func TestBundleLocalStateDirOverride(t *testing.T) {
-	ctx := context.Background()
+	ctx := t.Context()
 	projectDir := t.TempDir()
 	bundleTmpDir := t.TempDir()
 	f1, err := os.Create(filepath.Join(projectDir, "databricks.yml"))

bundle/config/loader/entry_point_test.go

@@ -1,7 +1,6 @@
 package loader_test
 
 import (
-	"context"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
@@ -12,15 +11,15 @@ import (
 
 func TestEntryPointNoRootPath(t *testing.T) {
 	b := &bundle.Bundle{}
-	diags := bundle.Apply(context.Background(), b, loader.EntryPoint())
+	diags := bundle.Apply(t.Context(), b, loader.EntryPoint())
 	require.Error(t, diags.Error())
 }
 
 func TestEntryPoint(t *testing.T) {
 	b := &bundle.Bundle{
 		BundleRootPath: "testdata/basic",
 	}
-	diags := bundle.Apply(context.Background(), b, loader.EntryPoint())
+	diags := bundle.Apply(t.Context(), b, loader.EntryPoint())
 	require.NoError(t, diags.Error())
 	assert.Equal(t, "loader_test", b.Config.Bundle.Name)
 }

bundle/config/loader/process_include_test.go

@@ -1,7 +1,6 @@
 package loader_test
 
 import (
-	"context"
 	"path/filepath"
 	"testing"
 
@@ -31,7 +30,7 @@ func TestProcessInclude(t *testing.T) {
 	assert.Equal(t, "foo", b.Config.Workspace.Host)
 
 	// Apply the mutator and assert that the host value has been updated
-	diags := bundle.Apply(context.Background(), b, m)
+	diags := bundle.Apply(t.Context(), b, m)
 	require.NoError(t, diags.Error())
 	assert.Equal(t, "bar", b.Config.Workspace.Host)
 }
@@ -55,7 +54,7 @@ func TestProcessIncludeFormatMatch(t *testing.T) {
 			}
 
 			m := loader.ProcessInclude(filepath.Join(b.BundleRootPath, fileName), fileName)
-			diags := bundle.Apply(context.Background(), b, m)
+			diags := bundle.Apply(t.Context(), b, m)
 			assert.Empty(t, diags)
 		})
 	}
@@ -210,7 +209,7 @@ func TestProcessIncludeFormatNotMatch(t *testing.T) {
 			}
 
 			m := loader.ProcessInclude(filepath.Join(b.BundleRootPath, fileName), fileName)
-			diags := bundle.Apply(context.Background(), b, m)
+			diags := bundle.Apply(t.Context(), b, m)
 			require.Len(t, diags, 1)
 			assert.Equal(t, expectedDiags, diags)
 		})

bundle/config/loader/process_root_includes_test.go

@@ -1,7 +1,6 @@
 package loader_test
 
 import (
-	"context"
 	"runtime"
 	"testing"
 
@@ -18,7 +17,7 @@ func TestProcessRootIncludesEmpty(t *testing.T) {
 	b := &bundle.Bundle{
 		BundleRootPath: ".",
 	}
-	diags := bundle.Apply(context.Background(), b, loader.ProcessRootIncludes())
+	diags := bundle.Apply(t.Context(), b, loader.ProcessRootIncludes())
 	require.NoError(t, diags.Error())
 }
 
@@ -38,7 +37,7 @@ func TestProcessRootIncludesAbs(t *testing.T) {
 			},
 		},
 	}
-	diags := bundle.Apply(context.Background(), b, loader.ProcessRootIncludes())
+	diags := bundle.Apply(t.Context(), b, loader.ProcessRootIncludes())
 	require.True(t, diags.HasError())
 	assert.ErrorContains(t, diags.Error(), "must be relative paths")
 }
@@ -57,7 +56,7 @@ func TestProcessRootIncludesSingleGlob(t *testing.T) {
 	testutil.Touch(t, b.BundleRootPath, "a.yml")
 	testutil.Touch(t, b.BundleRootPath, "b.yml")
 
-	diags := bundle.Apply(context.Background(), b, loader.ProcessRootIncludes())
+	diags := bundle.Apply(t.Context(), b, loader.ProcessRootIncludes())
 	require.NoError(t, diags.Error())
 	assert.Equal(t, []string{"a.yml", "b.yml"}, b.Config.Include)
 }
@@ -76,7 +75,7 @@ func TestProcessRootIncludesMultiGlob(t *testing.T) {
 	testutil.Touch(t, b.BundleRootPath, "a1.yml")
 	testutil.Touch(t, b.BundleRootPath, "b1.yml")
 
-	diags := bundle.Apply(context.Background(), b, loader.ProcessRootIncludes())
+	diags := bundle.Apply(t.Context(), b, loader.ProcessRootIncludes())
 	require.NoError(t, diags.Error())
 	assert.Equal(t, []string{"a1.yml", "b1.yml"}, b.Config.Include)
 }
@@ -94,7 +93,7 @@ func TestProcessRootIncludesRemoveDups(t *testing.T) {
 
 	testutil.Touch(t, b.BundleRootPath, "a.yml")
 
-	diags := bundle.Apply(context.Background(), b, loader.ProcessRootIncludes())
+	diags := bundle.Apply(t.Context(), b, loader.ProcessRootIncludes())
 	require.NoError(t, diags.Error())
 	assert.Equal(t, []string{"a.yml"}, b.Config.Include)
 }
@@ -108,7 +107,7 @@ func TestProcessRootIncludesNotExists(t *testing.T) {
 			},
 		},
 	}
-	diags := bundle.Apply(context.Background(), b, loader.ProcessRootIncludes())
+	diags := bundle.Apply(t.Context(), b, loader.ProcessRootIncludes())
 	require.True(t, diags.HasError())
 	assert.ErrorContains(t, diags.Error(), "notexist.yml defined in 'include' section does not match any files")
 }
@@ -172,7 +171,7 @@ func TestProcessRootIncludesGlobInRootPath(t *testing.T) {
 				BundleRootPath: test.root,
 			}
 
-			diags := bundle.Apply(context.Background(), b, loader.ProcessRootIncludes())
+			diags := bundle.Apply(t.Context(), b, loader.ProcessRootIncludes())
 			require.True(t, diags.HasError())
 			assert.Len(t, diags, 1)
 			assert.Equal(t, test.diag, diags[0])

bundle/config/mutator/apply_source_linked_deployment_preset_test.go

@@ -20,7 +20,7 @@ func TestApplyPresetsSourceLinkedDeployment(t *testing.T) {
 		t.Skip("this test is not applicable on Windows because source-linked mode works only in the Databricks Workspace")
 	}
 
-	testContext := context.Background()
+	testContext := t.Context()
 	enabled := true
 	disabled := false
 	workspacePath := "/Workspace/user.name@company.com"

bundle/config/mutator/compute_id_compate_test.go

@@ -1,7 +1,6 @@
 package mutator_test
 
 import (
-	"context"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
@@ -20,7 +19,7 @@ func TestComputeIdToClusterId(t *testing.T) {
 		},
 	}
 
-	diags := bundle.Apply(context.Background(), b, mutator.ComputeIdToClusterId())
+	diags := bundle.Apply(t.Context(), b, mutator.ComputeIdToClusterId())
 	assert.NoError(t, diags.Error())
 	assert.Equal(t, "compute-id", b.Config.Bundle.ClusterId)
 	assert.Empty(t, b.Config.Bundle.ComputeId)
@@ -41,11 +40,11 @@ func TestComputeIdToClusterIdInTargetOverride(t *testing.T) {
 		},
 	}
 
-	diags := bundle.Apply(context.Background(), b, mutator.ComputeIdToClusterId())
+	diags := bundle.Apply(t.Context(), b, mutator.ComputeIdToClusterId())
 	assert.NoError(t, diags.Error())
 	assert.Empty(t, b.Config.Targets["dev"].ComputeId)
 
-	diags = diags.Extend(bundle.Apply(context.Background(), b, mutator.SelectTarget("dev")))
+	diags = diags.Extend(bundle.Apply(t.Context(), b, mutator.SelectTarget("dev")))
 	assert.NoError(t, diags.Error())
 
 	assert.Equal(t, "compute-id-dev", b.Config.Bundle.ClusterId)

bundle/config/mutator/configure_wsfs_test.go

@@ -1,7 +1,6 @@
 package mutator_test
 
 import (
-	"context"
 	"reflect"
 	"runtime"
 	"testing"
@@ -37,7 +36,7 @@ func TestConfigureWSFS_SkipsIfNotWorkspacePrefix(t *testing.T) {
 	b := mockBundleForConfigureWSFS(t, "/foo")
 	originalSyncRoot := b.SyncRoot
 
-	ctx := context.Background()
+	ctx := t.Context()
 	diags := bundle.Apply(ctx, b, mutator.ConfigureWSFS())
 	assert.Empty(t, diags)
 	assert.Equal(t, originalSyncRoot, b.SyncRoot)
@@ -47,7 +46,7 @@ func TestConfigureWSFS_SkipsIfNotRunningOnRuntime(t *testing.T) {
 	b := mockBundleForConfigureWSFS(t, "/Workspace/foo")
 	originalSyncRoot := b.SyncRoot
 
-	ctx := context.Background()
+	ctx := t.Context()
 	ctx = dbr.MockRuntime(ctx, dbr.Environment{})
 	diags := bundle.Apply(ctx, b, mutator.ConfigureWSFS())
 	assert.Empty(t, diags)
@@ -91,7 +90,7 @@ func TestConfigureWSFS_DBRVersions(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			b := mockBundleForConfigureWSFS(t, "/Workspace/foo")
 
-			ctx := context.Background()
+			ctx := t.Context()
 			ctx = dbr.MockRuntime(ctx, dbr.Environment{IsDbr: true, Version: tt.version})
 			diags := bundle.Apply(ctx, b, mutator.ConfigureWSFS())
 			assert.Empty(t, diags)

bundle/config/mutator/default_target_test.go

@@ -1,7 +1,6 @@
 package mutator_test
 
 import (
-	"context"
 	"testing"
 
 	"github.com/databricks/cli/bundle"
@@ -13,7 +12,7 @@ import (
 
 func TestDefaultTarget(t *testing.T) {
 	b := &bundle.Bundle{}
-	diags := bundle.Apply(context.Background(), b, mutator.DefineDefaultTarget())
+	diags := bundle.Apply(t.Context(), b, mutator.DefineDefaultTarget())
 	require.NoError(t, diags.Error())
 
 	env, ok := b.Config.Targets["default"]
@@ -29,7 +28,7 @@ func TestDefaultTargetAlreadySpecified(t *testing.T) {
 			},
 		},
 	}
-	diags := bundle.Apply(context.Background(), b, mutator.DefineDefaultTarget())
+	diags := bundle.Apply(t.Context(), b, mutator.DefineDefaultTarget())
 	require.NoError(t, diags.Error())
 
 	_, ok := b.Config.Targets["default"]