From 3d91691f25e85220646df5b9bd8a672e28b14e7f Mon Sep 17 00:00:00 2001 From: Gleb Kanterov Date: Wed, 22 Jan 2025 16:37:37 +0100 Subject: [PATCH] PythonMutator: propagate source locations (#1783) ## Changes Add a mechanism to load Python source locations in the Python mutator. Previously, locations pointed to generated YAML. Now, they point to Python sources instead. Python process outputs "locations.json" containing locations of bundle paths, examples: ```json {"path": "resources.jobs.job_0", "file": "resources/job_0.py", "line": 3, "column": 5} {"path": "resources.jobs.job_0.tasks[0].task_key", "file": "resources/job_0.py", "line": 10, "column": 5} {"path": "resources.jobs.job_1", "file": "resources/job_1.py", "line": 5, "column": 7} ``` Such locations form a tree, and we assign locations of the closest ancestor to each `dyn.Value` based on its path. For example, `resources.jobs.job_0.tasks[0].task_key` is located at `job_0.py:10:5` and `resources.jobs.job_0.tasks[0].email_notifications` is located at `job_0.py:3:5`, because we use the location of the job as the most precise approximation. This feature is only enabled if `experimental/python` is used. Note: for now, we don't update locations with relative paths, because it has a side effect in changing how these paths are resolved ## Example ``` % databricks bundle validate Warning: job_cluster_key abc is not defined at resources.jobs.examples.tasks[0].job_cluster_key in resources/example.py:10:1 ``` ## Tests Unit tests and manually --- .../mutator/python/python_diagnostics.go | 1 + .../config/mutator/python/python_locations.go | 194 ++++++++++++++++++ .../mutator/python/python_locations_test.go | 179 ++++++++++++++++ .../config/mutator/python/python_mutator.go | 101 +++++++-- .../mutator/python/python_mutator_test.go | 79 +++++-- 5 files changed, 518 insertions(+), 36 deletions(-) create mode 100644 bundle/config/mutator/python/python_locations.go create mode 100644 bundle/config/mutator/python/python_locations_test.go diff --git a/bundle/config/mutator/python/python_diagnostics.go b/bundle/config/mutator/python/python_diagnostics.go index 12822065b..7a1e13b4e 100644 --- a/bundle/config/mutator/python/python_diagnostics.go +++ b/bundle/config/mutator/python/python_diagnostics.go @@ -9,6 +9,7 @@ import ( "github.com/databricks/cli/libs/dyn" ) +// pythonDiagnostic is a single entry in diagnostics.json type pythonDiagnostic struct { Severity pythonSeverity `json:"severity"` Summary string `json:"summary"` diff --git a/bundle/config/mutator/python/python_locations.go b/bundle/config/mutator/python/python_locations.go new file mode 100644 index 000000000..2fa86bea0 --- /dev/null +++ b/bundle/config/mutator/python/python_locations.go @@ -0,0 +1,194 @@ +package python + +import ( + "encoding/json" + "fmt" + "io" + "path/filepath" + + "github.com/databricks/cli/libs/dyn" +) + +// generatedFileName is used as the virtual file name for YAML generated by Python code. +// +// mergePythonLocations replaces dyn.Location with generatedFileName with locations loaded +// from locations.json +const generatedFileName = "__generated_by_python__.yml" + +// pythonLocations is data structure for efficient location lookup for a given path +// +// Locations form a tree, and we assign locations of the closest ancestor to each dyn.Value based on its path. +// We implement it as a trie (prefix tree) where keys are components of the path. With that, lookups are O(n) +// where n is the number of components in the path. +// +// For example, with locations.json: +// +// {"path": "resources.jobs.job_0", "file": "resources/job_0.py", "line": 3, "column": 5} +// {"path": "resources.jobs.job_0.tasks[0].task_key", "file": "resources/job_0.py", "line": 10, "column": 5} +// {"path": "resources.jobs.job_1", "file": "resources/job_1.py", "line": 5, "column": 7} +// +// - resources.jobs.job_0.tasks[0].task_key is located at job_0.py:10:5 +// +// - resources.jobs.job_0.tasks[0].email_notifications is located at job_0.py:3:5, +// because we use the location of the job as the most precise approximation. +// +// See pythonLocationEntry for the structure of a single entry in locations.json +type pythonLocations struct { + // descendants referenced by index, e.g. '.foo' + keys map[string]*pythonLocations + + // descendants referenced by key, e.g. '[0]' + indexes map[int]*pythonLocations + + // location for the current node if it exists + location dyn.Location + + // if true, location is present + exists bool +} + +// pythonLocationEntry is a single entry in locations.json +type pythonLocationEntry struct { + Path string `json:"path"` + File string `json:"file"` + Line int `json:"line"` + Column int `json:"column"` +} + +// mergePythonLocations applies locations from Python mutator into given dyn.Value +// +// The primary use-case is to merge locations.json with output.json, so that any +// validation errors will point to Python source code instead of generated YAML. +func mergePythonLocations(value dyn.Value, locations *pythonLocations) (dyn.Value, error) { + return dyn.Walk(value, func(path dyn.Path, value dyn.Value) (dyn.Value, error) { + newLocation, ok := findPythonLocation(locations, path) + if !ok { + return value, nil + } + + // The first item in the list is the "last" location used for error reporting + // + // Loaded YAML uses virtual file path as location, we remove any of such references, + // because they should use 'newLocation' instead. + // + // We preserve any previous non-virtual locations in case when Python function modified + // resource defined in YAML. + newLocations := append( + []dyn.Location{newLocation}, + removeVirtualLocations(value.Locations())..., + ) + + return value.WithLocations(newLocations), nil + }) +} + +func removeVirtualLocations(locations []dyn.Location) []dyn.Location { + var newLocations []dyn.Location + + for _, location := range locations { + if filepath.Base(location.File) == generatedFileName { + continue + } + + newLocations = append(newLocations, location) + } + + return newLocations +} + +// parsePythonLocations parses locations.json from the Python mutator. +// +// locations file is newline-separated JSON objects with pythonLocationEntry structure. +func parsePythonLocations(input io.Reader) (*pythonLocations, error) { + decoder := json.NewDecoder(input) + locations := newPythonLocations() + + for decoder.More() { + var entry pythonLocationEntry + + err := decoder.Decode(&entry) + if err != nil { + return nil, fmt.Errorf("failed to parse python location: %s", err) + } + + path, err := dyn.NewPathFromString(entry.Path) + if err != nil { + return nil, fmt.Errorf("failed to parse python location: %s", err) + } + + location := dyn.Location{ + File: entry.File, + Line: entry.Line, + Column: entry.Column, + } + + putPythonLocation(locations, path, location) + } + + return locations, nil +} + +// putPythonLocation puts the location to the trie for the given path +func putPythonLocation(trie *pythonLocations, path dyn.Path, location dyn.Location) { + currentNode := trie + + for _, component := range path { + if key := component.Key(); key != "" { + if _, ok := currentNode.keys[key]; !ok { + currentNode.keys[key] = newPythonLocations() + } + + currentNode = currentNode.keys[key] + } else { + index := component.Index() + if _, ok := currentNode.indexes[index]; !ok { + currentNode.indexes[index] = newPythonLocations() + } + + currentNode = currentNode.indexes[index] + } + } + + currentNode.location = location + currentNode.exists = true +} + +// newPythonLocations creates a new trie node +func newPythonLocations() *pythonLocations { + return &pythonLocations{ + keys: make(map[string]*pythonLocations), + indexes: make(map[int]*pythonLocations), + } +} + +// findPythonLocation finds the location or closest ancestor location in the trie for the given path +// if no ancestor or exact location is found, false is returned. +func findPythonLocation(locations *pythonLocations, path dyn.Path) (dyn.Location, bool) { + currentNode := locations + lastLocation := locations.location + exists := locations.exists + + for _, component := range path { + if key := component.Key(); key != "" { + if _, ok := currentNode.keys[key]; !ok { + break + } + + currentNode = currentNode.keys[key] + } else { + index := component.Index() + if _, ok := currentNode.indexes[index]; !ok { + break + } + + currentNode = currentNode.indexes[index] + } + + if currentNode.exists { + lastLocation = currentNode.location + exists = true + } + } + + return lastLocation, exists +} diff --git a/bundle/config/mutator/python/python_locations_test.go b/bundle/config/mutator/python/python_locations_test.go new file mode 100644 index 000000000..32afcc92b --- /dev/null +++ b/bundle/config/mutator/python/python_locations_test.go @@ -0,0 +1,179 @@ +package python + +import ( + "bytes" + "path/filepath" + "testing" + + "github.com/databricks/cli/libs/diag" + "github.com/stretchr/testify/require" + + "github.com/databricks/cli/libs/dyn" + assert "github.com/databricks/cli/libs/dyn/dynassert" +) + +func TestMergeLocations(t *testing.T) { + pythonLocation := dyn.Location{File: "foo.py", Line: 1, Column: 1} + generatedLocation := dyn.Location{File: generatedFileName, Line: 1, Column: 1} + yamlLocation := dyn.Location{File: "foo.yml", Line: 1, Column: 1} + + locations := newPythonLocations() + putPythonLocation(locations, dyn.MustPathFromString("foo"), pythonLocation) + + input := dyn.NewValue( + map[string]dyn.Value{ + "foo": dyn.NewValue( + map[string]dyn.Value{ + "baz": dyn.NewValue("baz", []dyn.Location{yamlLocation}), + "qux": dyn.NewValue("baz", []dyn.Location{generatedLocation, yamlLocation}), + }, + []dyn.Location{}, + ), + "bar": dyn.NewValue("baz", []dyn.Location{generatedLocation}), + }, + []dyn.Location{yamlLocation}, + ) + + expected := dyn.NewValue( + map[string]dyn.Value{ + "foo": dyn.NewValue( + map[string]dyn.Value{ + // pythonLocation is appended to the beginning of the list if absent + "baz": dyn.NewValue("baz", []dyn.Location{pythonLocation, yamlLocation}), + // generatedLocation is replaced by pythonLocation + "qux": dyn.NewValue("baz", []dyn.Location{pythonLocation, yamlLocation}), + }, + []dyn.Location{pythonLocation}, + ), + // if location is unknown, we keep it as-is + "bar": dyn.NewValue("baz", []dyn.Location{generatedLocation}), + }, + []dyn.Location{yamlLocation}, + ) + + actual, err := mergePythonLocations(input, locations) + + assert.NoError(t, err) + assert.Equal(t, expected, actual) +} + +func TestFindLocation(t *testing.T) { + location0 := dyn.Location{File: "foo.py", Line: 1, Column: 1} + location1 := dyn.Location{File: "foo.py", Line: 2, Column: 1} + + locations := newPythonLocations() + putPythonLocation(locations, dyn.MustPathFromString("foo"), location0) + putPythonLocation(locations, dyn.MustPathFromString("foo.bar"), location1) + + actual, exists := findPythonLocation(locations, dyn.MustPathFromString("foo.bar")) + + assert.True(t, exists) + assert.Equal(t, location1, actual) +} + +func TestFindLocation_indexPathComponent(t *testing.T) { + location0 := dyn.Location{File: "foo.py", Line: 1, Column: 1} + location1 := dyn.Location{File: "foo.py", Line: 2, Column: 1} + location2 := dyn.Location{File: "foo.py", Line: 3, Column: 1} + + locations := newPythonLocations() + putPythonLocation(locations, dyn.MustPathFromString("foo"), location0) + putPythonLocation(locations, dyn.MustPathFromString("foo.bar"), location1) + putPythonLocation(locations, dyn.MustPathFromString("foo.bar[0]"), location2) + + actual, exists := findPythonLocation(locations, dyn.MustPathFromString("foo.bar[0]")) + + assert.True(t, exists) + assert.Equal(t, location2, actual) +} + +func TestFindLocation_closestAncestorLocation(t *testing.T) { + location0 := dyn.Location{File: "foo.py", Line: 1, Column: 1} + location1 := dyn.Location{File: "foo.py", Line: 2, Column: 1} + + locations := newPythonLocations() + putPythonLocation(locations, dyn.MustPathFromString("foo"), location0) + putPythonLocation(locations, dyn.MustPathFromString("foo.bar"), location1) + + actual, exists := findPythonLocation(locations, dyn.MustPathFromString("foo.bar.baz")) + + assert.True(t, exists) + assert.Equal(t, location1, actual) +} + +func TestFindLocation_unknownLocation(t *testing.T) { + location0 := dyn.Location{File: "foo.py", Line: 1, Column: 1} + location1 := dyn.Location{File: "foo.py", Line: 2, Column: 1} + + locations := newPythonLocations() + putPythonLocation(locations, dyn.MustPathFromString("foo"), location0) + putPythonLocation(locations, dyn.MustPathFromString("foo.bar"), location1) + + _, exists := findPythonLocation(locations, dyn.MustPathFromString("bar")) + + assert.False(t, exists) +} + +func TestLoadOutput(t *testing.T) { + location := dyn.Location{File: "my_job.py", Line: 1, Column: 1} + bundleRoot := t.TempDir() + output := `{ + "resources": { + "jobs": { + "my_job": { + "name": "my_job", + "tasks": [ + { + "task_key": "my_task", + "notebook_task": { + "notebook_path": "my_notebook" + } + } + ] + } + } + } + }` + + locations := newPythonLocations() + putPythonLocation( + locations, + dyn.MustPathFromString("resources.jobs.my_job"), + location, + ) + + value, diags := loadOutput( + bundleRoot, + bytes.NewReader([]byte(output)), + locations, + ) + + assert.Equal(t, diag.Diagnostics{}, diags) + + name, err := dyn.Get(value, "resources.jobs.my_job.name") + require.NoError(t, err) + require.Equal(t, []dyn.Location{location}, name.Locations()) + + // until we implement path normalization, we have to keep locations of values + // that change semantic depending on their location + // + // note: it's important to have absolute path including 'bundleRoot' + // because mutator pipeline already has expanded locations into absolute path + notebookPath, err := dyn.Get(value, "resources.jobs.my_job.tasks[0].notebook_task.notebook_path") + require.NoError(t, err) + require.Len(t, notebookPath.Locations(), 1) + require.Equal(t, filepath.Join(bundleRoot, generatedFileName), notebookPath.Locations()[0].File) +} + +func TestParsePythonLocations(t *testing.T) { + expected := dyn.Location{File: "foo.py", Line: 1, Column: 2} + + input := `{"path": "foo", "file": "foo.py", "line": 1, "column": 2}` + reader := bytes.NewReader([]byte(input)) + locations, err := parsePythonLocations(reader) + + assert.NoError(t, err) + + assert.True(t, locations.keys["foo"].exists) + assert.Equal(t, expected, locations.keys["foo"].location) +} diff --git a/bundle/config/mutator/python/python_mutator.go b/bundle/config/mutator/python/python_mutator.go index 8009ab243..cd2e286e5 100644 --- a/bundle/config/mutator/python/python_mutator.go +++ b/bundle/config/mutator/python/python_mutator.go @@ -7,11 +7,14 @@ import ( "errors" "fmt" "io" + "io/fs" "os" "path/filepath" "reflect" "strings" + "github.com/databricks/cli/bundle/config/mutator/paths" + "github.com/databricks/databricks-sdk-go/logger" "github.com/fatih/color" @@ -124,6 +127,15 @@ type opts struct { enabled bool venvPath string + + loadLocations bool +} + +type runPythonMutatorOpts struct { + cacheDir string + bundleRootPath string + pythonPath string + loadLocations bool } // getOpts adapts deprecated PyDABs and upcoming Python configuration @@ -148,8 +160,9 @@ func getOpts(b *bundle.Bundle, phase phase) (opts, error) { // don't execute for phases for 'python' section if phase == PythonMutatorPhaseInit || phase == PythonMutatorPhaseLoad { return opts{ - enabled: true, - venvPath: experimental.PyDABs.VEnvPath, + enabled: true, + venvPath: experimental.PyDABs.VEnvPath, + loadLocations: false, // not supported in PyDABs }, nil } else { return opts{}, nil @@ -158,8 +171,9 @@ func getOpts(b *bundle.Bundle, phase phase) (opts, error) { // don't execute for phases for 'pydabs' section if phase == PythonMutatorPhaseLoadResources || phase == PythonMutatorPhaseApplyMutators { return opts{ - enabled: true, - venvPath: experimental.Python.VEnvPath, + enabled: true, + venvPath: experimental.Python.VEnvPath, + loadLocations: true, }, nil } else { return opts{}, nil @@ -194,7 +208,12 @@ func (m *pythonMutator) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagno return dyn.InvalidValue, fmt.Errorf("failed to create cache dir: %w", err) } - rightRoot, diags := m.runPythonMutator(ctx, cacheDir, b.BundleRootPath, pythonPath, leftRoot) + rightRoot, diags := m.runPythonMutator(ctx, leftRoot, runPythonMutatorOpts{ + cacheDir: cacheDir, + bundleRootPath: b.BundleRootPath, + pythonPath: pythonPath, + loadLocations: opts.loadLocations, + }) mutateDiags = diags if diags.HasError() { return dyn.InvalidValue, mutateDiagsHasError @@ -238,13 +257,14 @@ func createCacheDir(ctx context.Context) (string, error) { return os.MkdirTemp("", "-python") } -func (m *pythonMutator) runPythonMutator(ctx context.Context, cacheDir, rootPath, pythonPath string, root dyn.Value) (dyn.Value, diag.Diagnostics) { - inputPath := filepath.Join(cacheDir, "input.json") - outputPath := filepath.Join(cacheDir, "output.json") - diagnosticsPath := filepath.Join(cacheDir, "diagnostics.json") +func (m *pythonMutator) runPythonMutator(ctx context.Context, root dyn.Value, opts runPythonMutatorOpts) (dyn.Value, diag.Diagnostics) { + inputPath := filepath.Join(opts.cacheDir, "input.json") + outputPath := filepath.Join(opts.cacheDir, "output.json") + diagnosticsPath := filepath.Join(opts.cacheDir, "diagnostics.json") + locationsPath := filepath.Join(opts.cacheDir, "locations.json") args := []string{ - pythonPath, + opts.pythonPath, "-m", "databricks.bundles.build", "--phase", @@ -257,6 +277,10 @@ func (m *pythonMutator) runPythonMutator(ctx context.Context, cacheDir, rootPath diagnosticsPath, } + if opts.loadLocations { + args = append(args, "--locations", locationsPath) + } + if err := writeInputFile(inputPath, root); err != nil { return dyn.InvalidValue, diag.Errorf("failed to write input file: %s", err) } @@ -271,7 +295,7 @@ func (m *pythonMutator) runPythonMutator(ctx context.Context, cacheDir, rootPath _, processErr := process.Background( ctx, args, - process.WithDir(rootPath), + process.WithDir(opts.bundleRootPath), process.WithStderrWriter(stderrWriter), process.WithStdoutWriter(stdoutWriter), ) @@ -307,7 +331,12 @@ func (m *pythonMutator) runPythonMutator(ctx context.Context, cacheDir, rootPath return dyn.InvalidValue, diag.Errorf("failed to load diagnostics: %s", pythonDiagnosticsErr) } - output, outputDiags := loadOutputFile(rootPath, outputPath) + locations, err := loadLocationsFile(locationsPath) + if err != nil { + return dyn.InvalidValue, diag.Errorf("failed to load locations: %s", err) + } + + output, outputDiags := loadOutputFile(opts.bundleRootPath, outputPath, locations) pythonDiagnostics = pythonDiagnostics.Extend(outputDiags) // we pass through pythonDiagnostic because it contains warnings @@ -351,7 +380,21 @@ func writeInputFile(inputPath string, input dyn.Value) error { return os.WriteFile(inputPath, rootConfigJson, 0o600) } -func loadOutputFile(rootPath, outputPath string) (dyn.Value, diag.Diagnostics) { +// loadLocationsFile loads locations.json containing source locations for generated YAML. +func loadLocationsFile(locationsPath string) (*pythonLocations, error) { + locationsFile, err := os.Open(locationsPath) + if errors.Is(err, fs.ErrNotExist) { + return newPythonLocations(), nil + } else if err != nil { + return nil, fmt.Errorf("failed to open locations file: %w", err) + } + + defer locationsFile.Close() + + return parsePythonLocations(locationsFile) +} + +func loadOutputFile(rootPath, outputPath string, locations *pythonLocations) (dyn.Value, diag.Diagnostics) { outputFile, err := os.Open(outputPath) if err != nil { return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to open output file: %w", err)) @@ -359,15 +402,19 @@ func loadOutputFile(rootPath, outputPath string) (dyn.Value, diag.Diagnostics) { defer outputFile.Close() + return loadOutput(rootPath, outputFile, locations) +} + +func loadOutput(rootPath string, outputFile io.Reader, locations *pythonLocations) (dyn.Value, diag.Diagnostics) { // we need absolute path because later parts of pipeline assume all paths are absolute // and this file will be used as location to resolve relative paths. // - // virtualPath has to stay in rootPath, because locations outside root path are not allowed: + // virtualPath has to stay in bundleRootPath, because locations outside root path are not allowed: // // Error: path /var/folders/.../python/dist/*.whl is not contained in bundle root path // // for that, we pass virtualPath instead of outputPath as file location - virtualPath, err := filepath.Abs(filepath.Join(rootPath, "__generated_by_python__.yml")) + virtualPath, err := filepath.Abs(filepath.Join(rootPath, generatedFileName)) if err != nil { return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to get absolute path: %w", err)) } @@ -377,7 +424,29 @@ func loadOutputFile(rootPath, outputPath string) (dyn.Value, diag.Diagnostics) { return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to parse output file: %w", err)) } - return strictNormalize(config.Root{}, generated) + // paths are resolved relative to locations of their values, if we change location + // we have to update each path, until we simplify that, we don't update locations + // for such values, so we don't change how paths are resolved + // + // we can remove this once we: + // - add variable interpolation before and after PythonMutator + // - implement path normalization (aka path normal form) + _, err = paths.VisitJobPaths(generated, func(p dyn.Path, kind paths.PathKind, v dyn.Value) (dyn.Value, error) { + putPythonLocation(locations, p, v.Location()) + return v, nil + }) + if err != nil { + return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to update locations: %w", err)) + } + + // generated has dyn.Location as if it comes from generated YAML file + // earlier we loaded locations.json with source locations in Python code + generatedWithLocations, err := mergePythonLocations(generated, locations) + if err != nil { + return dyn.InvalidValue, diag.FromErr(fmt.Errorf("failed to update locations: %w", err)) + } + + return strictNormalize(config.Root{}, generatedWithLocations) } func strictNormalize(dst any, generated dyn.Value) (dyn.Value, diag.Diagnostics) { diff --git a/bundle/config/mutator/python/python_mutator_test.go b/bundle/config/mutator/python/python_mutator_test.go index d51572c8a..322fb79e8 100644 --- a/bundle/config/mutator/python/python_mutator_test.go +++ b/bundle/config/mutator/python/python_mutator_test.go @@ -7,7 +7,6 @@ import ( "os" "os/exec" "path/filepath" - "reflect" "runtime" "testing" @@ -93,6 +92,8 @@ func TestPythonMutator_loadResources(t *testing.T) { } }`, `{"severity": "warning", "summary": "job doesn't have any tasks", "location": {"file": "src/examples/file.py", "line": 10, "column": 5}}`, + `{"path": "resources.jobs.job0", "file": "src/examples/job0.py", "line": 3, "column": 5} + {"path": "resources.jobs.job1", "file": "src/examples/job1.py", "line": 5, "column": 7}`, ) mutator := PythonMutator(PythonMutatorPhaseLoadResources) @@ -110,6 +111,25 @@ func TestPythonMutator_loadResources(t *testing.T) { assert.Equal(t, "job_1", job1.Name) } + // output of locations.json should be applied to underlying dyn.Value + err := b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) { + name1, err := dyn.GetByPath(v, dyn.MustPathFromString("resources.jobs.job1.name")) + if err != nil { + return dyn.InvalidValue, err + } + + assert.Equal(t, []dyn.Location{ + { + File: "src/examples/job1.py", + Line: 5, + Column: 7, + }, + }, name1.Locations()) + + return v, nil + }) + assert.NoError(t, err) + assert.Equal(t, 1, len(diags)) assert.Equal(t, "job doesn't have any tasks", diags[0].Summary) assert.Equal(t, []dyn.Location{ @@ -157,7 +177,7 @@ func TestPythonMutator_loadResources_disallowed(t *testing.T) { } } } - }`, "") + }`, "", "") mutator := PythonMutator(PythonMutatorPhaseLoadResources) diag := bundle.Apply(ctx, b, mutator) @@ -202,7 +222,7 @@ func TestPythonMutator_applyMutators(t *testing.T) { } } } - }`, "") + }`, "", "") mutator := PythonMutator(PythonMutatorPhaseApplyMutators) diag := bundle.Apply(ctx, b, mutator) @@ -224,7 +244,7 @@ func TestPythonMutator_applyMutators(t *testing.T) { description, err := dyn.GetByPath(v, dyn.MustPathFromString("resources.jobs.job0.description")) require.NoError(t, err) - expectedVirtualPath, err := filepath.Abs("__generated_by_python__.yml") + expectedVirtualPath, err := filepath.Abs(generatedFileName) require.NoError(t, err) assert.Equal(t, expectedVirtualPath, description.Location().File) @@ -263,7 +283,7 @@ func TestPythonMutator_badOutput(t *testing.T) { } } } - }`, "") + }`, "", "") mutator := PythonMutator(PythonMutatorPhaseLoadResources) diag := bundle.Apply(ctx, b, mutator) @@ -312,7 +332,7 @@ func TestGetOps_Python(t *testing.T) { }, PythonMutatorPhaseLoadResources) assert.NoError(t, err) - assert.Equal(t, opts{venvPath: ".venv", enabled: true}, actual) + assert.Equal(t, opts{venvPath: ".venv", enabled: true, loadLocations: true}, actual) } func TestGetOps_PyDABs(t *testing.T) { @@ -328,7 +348,7 @@ func TestGetOps_PyDABs(t *testing.T) { }, PythonMutatorPhaseInit) assert.NoError(t, err) - assert.Equal(t, opts{venvPath: ".venv", enabled: true}, actual) + assert.Equal(t, opts{venvPath: ".venv", enabled: true, loadLocations: false}, actual) } func TestGetOps_empty(t *testing.T) { @@ -661,7 +681,7 @@ or activate the environment before running CLI commands: assert.Equal(t, expected, out) } -func withProcessStub(t *testing.T, args []string, output, diagnostics string) context.Context { +func withProcessStub(t *testing.T, args []string, output, diagnostics, locations string) context.Context { ctx := context.Background() ctx, stub := process.WithStub(ctx) @@ -673,32 +693,51 @@ func withProcessStub(t *testing.T, args []string, output, diagnostics string) co inputPath := filepath.Join(cacheDir, "input.json") outputPath := filepath.Join(cacheDir, "output.json") + locationsPath := filepath.Join(cacheDir, "locations.json") diagnosticsPath := filepath.Join(cacheDir, "diagnostics.json") - args = append(args, "--input", inputPath) - args = append(args, "--output", outputPath) - args = append(args, "--diagnostics", diagnosticsPath) - stub.WithCallback(func(actual *exec.Cmd) error { _, err := os.Stat(inputPath) assert.NoError(t, err) - if reflect.DeepEqual(actual.Args, args) { - err := os.WriteFile(outputPath, []byte(output), 0o600) - require.NoError(t, err) + actualInputPath := getArg(actual.Args, "--input") + actualOutputPath := getArg(actual.Args, "--output") + actualDiagnosticsPath := getArg(actual.Args, "--diagnostics") + actualLocationsPath := getArg(actual.Args, "--locations") - err = os.WriteFile(diagnosticsPath, []byte(diagnostics), 0o600) - require.NoError(t, err) + require.Equal(t, inputPath, actualInputPath) + require.Equal(t, outputPath, actualOutputPath) + require.Equal(t, diagnosticsPath, actualDiagnosticsPath) - return nil - } else { - return fmt.Errorf("unexpected command: %v", actual.Args) + // locations is an optional argument + if locations != "" { + require.Equal(t, locationsPath, actualLocationsPath) + + err = os.WriteFile(locationsPath, []byte(locations), 0o600) + require.NoError(t, err) } + + err = os.WriteFile(outputPath, []byte(output), 0o600) + require.NoError(t, err) + + err = os.WriteFile(diagnosticsPath, []byte(diagnostics), 0o600) + require.NoError(t, err) + + return nil }) return ctx } +func getArg(args []string, name string) string { + for i := range args { + if args[i] == name { + return args[i+1] + } + } + return "" +} + func loadYaml(name, content string) *bundle.Bundle { v, diag := config.LoadFromBytes(name, []byte(content))