From 495040e4cd2d8fbbbcc09aff6cf3b88cb4daee78 Mon Sep 17 00:00:00 2001
From: shreyas-goenka <88374338+shreyas-goenka@users.noreply.github.com>
Date: Wed, 25 Sep 2024 21:43:48 +0530
Subject: [PATCH] Modify SetLocation test utility to take full locations as
 argument (#1788)

I plan to use this in https://github.com/databricks/cli/pull/1780, to
set the line and column numbers as well for the locations.

gopatch file used:
```
@@
var x expression
var y expression
var z expression
@@
-bundletest.SetLocation(x, y, z)
+bundletest.SetLocation(x, y, []dyn.Location{{File: z}})
```
---
 bundle/artifacts/expand_globs_test.go         |  7 ++--
 .../expand_pipeline_glob_paths_test.go        |  5 +--
 .../config/mutator/rewrite_sync_paths_test.go | 25 ++++++++-------
 bundle/config/mutator/sync_infer_root_test.go |  3 +-
 bundle/config/mutator/translate_paths_test.go | 32 +++++++++----------
 bundle/deploy/metadata/compute_test.go        |  7 ++--
 bundle/internal/bundletest/location.go        |  6 ++--
 .../libraries/expand_glob_references_test.go  |  7 ++--
 8 files changed, 48 insertions(+), 44 deletions(-)

diff --git a/bundle/artifacts/expand_globs_test.go b/bundle/artifacts/expand_globs_test.go
index c9c47844..1665a480 100644
--- a/bundle/artifacts/expand_globs_test.go
+++ b/bundle/artifacts/expand_globs_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/internal/bundletest"
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/dyn"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -36,7 +37,7 @@ func TestExpandGlobs_Nominal(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "artifacts", filepath.Join(tmpDir, "databricks.yml"))
+	bundletest.SetLocation(b, "artifacts", []dyn.Location{{File: filepath.Join(tmpDir, "databricks.yml")}})
 
 	ctx := context.Background()
 	diags := bundle.Apply(ctx, b, bundle.Seq(
@@ -77,7 +78,7 @@ func TestExpandGlobs_InvalidPattern(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "artifacts", filepath.Join(tmpDir, "databricks.yml"))
+	bundletest.SetLocation(b, "artifacts", []dyn.Location{{File: filepath.Join(tmpDir, "databricks.yml")}})
 
 	ctx := context.Background()
 	diags := bundle.Apply(ctx, b, bundle.Seq(
@@ -125,7 +126,7 @@ func TestExpandGlobs_NoMatches(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "artifacts", filepath.Join(tmpDir, "databricks.yml"))
+	bundletest.SetLocation(b, "artifacts", []dyn.Location{{File: filepath.Join(tmpDir, "databricks.yml")}})
 
 	ctx := context.Background()
 	diags := bundle.Apply(ctx, b, bundle.Seq(
diff --git a/bundle/config/mutator/expand_pipeline_glob_paths_test.go b/bundle/config/mutator/expand_pipeline_glob_paths_test.go
index d1671c25..07dd2021 100644
--- a/bundle/config/mutator/expand_pipeline_glob_paths_test.go
+++ b/bundle/config/mutator/expand_pipeline_glob_paths_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/bundle/internal/bundletest"
+	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/pipelines"
 	"github.com/stretchr/testify/require"
@@ -105,8 +106,8 @@ func TestExpandGlobPathsInPipelines(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
-	bundletest.SetLocation(b, "resources.pipelines.pipeline.libraries[3]", filepath.Join(dir, "relative", "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
+	bundletest.SetLocation(b, "resources.pipelines.pipeline.libraries[3]", []dyn.Location{{File: filepath.Join(dir, "relative", "resource.yml")}})
 
 	m := ExpandPipelineGlobPaths()
 	diags := bundle.Apply(context.Background(), b, m)
diff --git a/bundle/config/mutator/rewrite_sync_paths_test.go b/bundle/config/mutator/rewrite_sync_paths_test.go
index fa7f124b..a66f2763 100644
--- a/bundle/config/mutator/rewrite_sync_paths_test.go
+++ b/bundle/config/mutator/rewrite_sync_paths_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/mutator"
 	"github.com/databricks/cli/bundle/internal/bundletest"
+	"github.com/databricks/cli/libs/dyn"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -33,12 +34,12 @@ func TestRewriteSyncPathsRelative(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "sync.paths[0]", "./databricks.yml")
-	bundletest.SetLocation(b, "sync.paths[1]", "./databricks.yml")
-	bundletest.SetLocation(b, "sync.include[0]", "./file.yml")
-	bundletest.SetLocation(b, "sync.include[1]", "./a/file.yml")
-	bundletest.SetLocation(b, "sync.exclude[0]", "./a/b/file.yml")
-	bundletest.SetLocation(b, "sync.exclude[1]", "./a/b/c/file.yml")
+	bundletest.SetLocation(b, "sync.paths[0]", []dyn.Location{{File: "./databricks.yml"}})
+	bundletest.SetLocation(b, "sync.paths[1]", []dyn.Location{{File: "./databricks.yml"}})
+	bundletest.SetLocation(b, "sync.include[0]", []dyn.Location{{File: "./file.yml"}})
+	bundletest.SetLocation(b, "sync.include[1]", []dyn.Location{{File: "./a/file.yml"}})
+	bundletest.SetLocation(b, "sync.exclude[0]", []dyn.Location{{File: "./a/b/file.yml"}})
+	bundletest.SetLocation(b, "sync.exclude[1]", []dyn.Location{{File: "./a/b/c/file.yml"}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.RewriteSyncPaths())
 	assert.NoError(t, diags.Error())
@@ -72,12 +73,12 @@ func TestRewriteSyncPathsAbsolute(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "sync.paths[0]", "/tmp/dir/databricks.yml")
-	bundletest.SetLocation(b, "sync.paths[1]", "/tmp/dir/databricks.yml")
-	bundletest.SetLocation(b, "sync.include[0]", "/tmp/dir/file.yml")
-	bundletest.SetLocation(b, "sync.include[1]", "/tmp/dir/a/file.yml")
-	bundletest.SetLocation(b, "sync.exclude[0]", "/tmp/dir/a/b/file.yml")
-	bundletest.SetLocation(b, "sync.exclude[1]", "/tmp/dir/a/b/c/file.yml")
+	bundletest.SetLocation(b, "sync.paths[0]", []dyn.Location{{File: "/tmp/dir/databricks.yml"}})
+	bundletest.SetLocation(b, "sync.paths[1]", []dyn.Location{{File: "/tmp/dir/databricks.yml"}})
+	bundletest.SetLocation(b, "sync.include[0]", []dyn.Location{{File: "/tmp/dir/file.yml"}})
+	bundletest.SetLocation(b, "sync.include[1]", []dyn.Location{{File: "/tmp/dir/a/file.yml"}})
+	bundletest.SetLocation(b, "sync.exclude[0]", []dyn.Location{{File: "/tmp/dir/a/b/file.yml"}})
+	bundletest.SetLocation(b, "sync.exclude[1]", []dyn.Location{{File: "/tmp/dir/a/b/c/file.yml"}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.RewriteSyncPaths())
 	assert.NoError(t, diags.Error())
diff --git a/bundle/config/mutator/sync_infer_root_test.go b/bundle/config/mutator/sync_infer_root_test.go
index 383e5676..85e40adc 100644
--- a/bundle/config/mutator/sync_infer_root_test.go
+++ b/bundle/config/mutator/sync_infer_root_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/cli/bundle/config"
 	"github.com/databricks/cli/bundle/config/mutator"
 	"github.com/databricks/cli/bundle/internal/bundletest"
+	"github.com/databricks/cli/libs/dyn"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -184,7 +185,7 @@ func TestSyncInferRoot_Error(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "sync.paths", "databricks.yml")
+	bundletest.SetLocation(b, "sync.paths", []dyn.Location{{File: "databricks.yml"}})
 
 	ctx := context.Background()
 	diags := bundle.Apply(ctx, b, mutator.SyncInferRoot())
diff --git a/bundle/config/mutator/translate_paths_test.go b/bundle/config/mutator/translate_paths_test.go
index 50fcd3b0..c03cee73 100644
--- a/bundle/config/mutator/translate_paths_test.go
+++ b/bundle/config/mutator/translate_paths_test.go
@@ -82,7 +82,7 @@ func TestTranslatePathsSkippedWithGitSource(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	require.NoError(t, diags.Error())
@@ -210,7 +210,7 @@ func TestTranslatePaths(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	require.NoError(t, diags.Error())
@@ -346,8 +346,8 @@ func TestTranslatePathsInSubdirectories(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "resources.jobs", filepath.Join(dir, "job/resource.yml"))
-	bundletest.SetLocation(b, "resources.pipelines", filepath.Join(dir, "pipeline/resource.yml"))
+	bundletest.SetLocation(b, "resources.jobs", []dyn.Location{{File: filepath.Join(dir, "job/resource.yml")}})
+	bundletest.SetLocation(b, "resources.pipelines", []dyn.Location{{File: filepath.Join(dir, "pipeline/resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	require.NoError(t, diags.Error())
@@ -408,7 +408,7 @@ func TestTranslatePathsOutsideSyncRoot(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "../resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "../resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.ErrorContains(t, diags.Error(), "is not contained in sync root path")
@@ -439,7 +439,7 @@ func TestJobNotebookDoesNotExistError(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "fake.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "fake.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.EqualError(t, diags.Error(), "notebook ./doesnt_exist.py not found")
@@ -470,7 +470,7 @@ func TestJobFileDoesNotExistError(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "fake.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "fake.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.EqualError(t, diags.Error(), "file ./doesnt_exist.py not found")
@@ -501,7 +501,7 @@ func TestPipelineNotebookDoesNotExistError(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "fake.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "fake.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.EqualError(t, diags.Error(), "notebook ./doesnt_exist.py not found")
@@ -532,7 +532,7 @@ func TestPipelineFileDoesNotExistError(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "fake.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "fake.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.EqualError(t, diags.Error(), "file ./doesnt_exist.py not found")
@@ -567,7 +567,7 @@ func TestJobSparkPythonTaskWithNotebookSourceError(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.ErrorContains(t, diags.Error(), `expected a file for "resources.jobs.job.tasks[0].spark_python_task.python_file" but got a notebook`)
@@ -602,7 +602,7 @@ func TestJobNotebookTaskWithFileSourceError(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.ErrorContains(t, diags.Error(), `expected a notebook for "resources.jobs.job.tasks[0].notebook_task.notebook_path" but got a file`)
@@ -637,7 +637,7 @@ func TestPipelineNotebookLibraryWithFileSourceError(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.ErrorContains(t, diags.Error(), `expected a notebook for "resources.pipelines.pipeline.libraries[0].notebook.path" but got a file`)
@@ -672,7 +672,7 @@ func TestPipelineFileLibraryWithNotebookSourceError(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	assert.ErrorContains(t, diags.Error(), `expected a file for "resources.pipelines.pipeline.libraries[0].file.path" but got a notebook`)
@@ -710,7 +710,7 @@ func TestTranslatePathJobEnvironments(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "resources.jobs", filepath.Join(dir, "job/resource.yml"))
+	bundletest.SetLocation(b, "resources.jobs", []dyn.Location{{File: filepath.Join(dir, "job/resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, mutator.TranslatePaths())
 	require.NoError(t, diags.Error())
@@ -753,8 +753,8 @@ func TestTranslatePathWithComplexVariables(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "variables", filepath.Join(dir, "variables/variables.yml"))
-	bundletest.SetLocation(b, "resources.jobs", filepath.Join(dir, "job/resource.yml"))
+	bundletest.SetLocation(b, "variables", []dyn.Location{{File: filepath.Join(dir, "variables/variables.yml")}})
+	bundletest.SetLocation(b, "resources.jobs", []dyn.Location{{File: filepath.Join(dir, "job/resource.yml")}})
 
 	ctx := context.Background()
 	// Assign the variables to the dynamic configuration.
diff --git a/bundle/deploy/metadata/compute_test.go b/bundle/deploy/metadata/compute_test.go
index 6d43f845..2c2c7237 100644
--- a/bundle/deploy/metadata/compute_test.go
+++ b/bundle/deploy/metadata/compute_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/bundle/internal/bundletest"
 	"github.com/databricks/cli/bundle/metadata"
+	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -55,9 +56,9 @@ func TestComputeMetadataMutator(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, "resources.jobs.my-job-1", "a/b/c")
-	bundletest.SetLocation(b, "resources.jobs.my-job-2", "d/e/f")
-	bundletest.SetLocation(b, "resources.pipelines.my-pipeline", "abc")
+	bundletest.SetLocation(b, "resources.jobs.my-job-1", []dyn.Location{{File: "a/b/c"}})
+	bundletest.SetLocation(b, "resources.jobs.my-job-2", []dyn.Location{{File: "d/e/f"}})
+	bundletest.SetLocation(b, "resources.pipelines.my-pipeline", []dyn.Location{{File: "abc"}})
 
 	expectedMetadata := metadata.Metadata{
 		Version: metadata.Version,
diff --git a/bundle/internal/bundletest/location.go b/bundle/internal/bundletest/location.go
index 380d6e17..2ffd621b 100644
--- a/bundle/internal/bundletest/location.go
+++ b/bundle/internal/bundletest/location.go
@@ -8,15 +8,13 @@ import (
 // SetLocation sets the location of all values in the bundle to the given path.
 // This is useful for testing where we need to associate configuration
 // with the path it is loaded from.
-func SetLocation(b *bundle.Bundle, prefix string, filePath string) {
+func SetLocation(b *bundle.Bundle, prefix string, locations []dyn.Location) {
 	start := dyn.MustPathFromString(prefix)
 	b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) {
 		return dyn.Walk(root, func(p dyn.Path, v dyn.Value) (dyn.Value, error) {
 			// If the path has the given prefix, set the location.
 			if p.HasPrefix(start) {
-				return v.WithLocations([]dyn.Location{{
-					File: filePath,
-				}}), nil
+				return v.WithLocations(locations), nil
 			}
 
 			// The path is not nested under the given prefix.
diff --git a/bundle/libraries/expand_glob_references_test.go b/bundle/libraries/expand_glob_references_test.go
index e7f2e169..2dfbddb7 100644
--- a/bundle/libraries/expand_glob_references_test.go
+++ b/bundle/libraries/expand_glob_references_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/databricks/cli/bundle/config/resources"
 	"github.com/databricks/cli/bundle/internal/bundletest"
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/databricks-sdk-go/service/compute"
 	"github.com/databricks/databricks-sdk-go/service/jobs"
 	"github.com/stretchr/testify/require"
@@ -61,7 +62,7 @@ func TestGlobReferencesExpandedForTaskLibraries(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, ExpandGlobReferences())
 	require.Empty(t, diags)
@@ -146,7 +147,7 @@ func TestGlobReferencesExpandedForForeachTaskLibraries(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, ExpandGlobReferences())
 	require.Empty(t, diags)
@@ -221,7 +222,7 @@ func TestGlobReferencesExpandedForEnvironmentsDeps(t *testing.T) {
 		},
 	}
 
-	bundletest.SetLocation(b, ".", filepath.Join(dir, "resource.yml"))
+	bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "resource.yml")}})
 
 	diags := bundle.Apply(context.Background(), b, ExpandGlobReferences())
 	require.Empty(t, diags)