databricks-cli/bundle/config/mutator/translate_notebook_paths_te...

135 lines
3.4 KiB
Go

package mutator_test
import (
"context"
"os"
"path/filepath"
"testing"
"github.com/databricks/bricks/bundle"
"github.com/databricks/bricks/bundle/config"
"github.com/databricks/bricks/bundle/config/mutator"
"github.com/databricks/bricks/bundle/config/resources"
"github.com/databricks/databricks-sdk-go/service/jobs"
"github.com/databricks/databricks-sdk-go/service/pipelines"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func touchFile(t *testing.T, path string) {
f, err := os.Create(path)
require.NoError(t, err)
f.Close()
}
func TestNotebookPaths(t *testing.T) {
dir := t.TempDir()
touchFile(t, filepath.Join(dir, "my_job_notebook.py"))
touchFile(t, filepath.Join(dir, "my_pipeline_notebook.py"))
bundle := &bundle.Bundle{
Config: config.Root{
Path: dir,
Resources: config.Resources{
Jobs: map[string]*resources.Job{
"job": {
JobSettings: &jobs.JobSettings{
Tasks: []jobs.JobTaskSettings{
{
NotebookTask: &jobs.NotebookTask{
NotebookPath: "./my_job_notebook.py",
},
},
{
NotebookTask: &jobs.NotebookTask{
NotebookPath: "./doesnt_exist.py",
},
},
{
NotebookTask: &jobs.NotebookTask{
NotebookPath: "./my_job_notebook.py",
},
},
{
PythonWheelTask: &jobs.PythonWheelTask{
PackageName: "foo",
},
},
},
},
},
},
Pipelines: map[string]*resources.Pipeline{
"pipeline": {
PipelineSpec: &pipelines.PipelineSpec{
Libraries: []pipelines.PipelineLibrary{
{
Notebook: &pipelines.NotebookLibrary{
Path: "./my_pipeline_notebook.py",
},
},
{
Notebook: &pipelines.NotebookLibrary{
Path: "./doesnt_exist.py",
},
},
{
Notebook: &pipelines.NotebookLibrary{
Path: "./my_pipeline_notebook.py",
},
},
{
Jar: "foo",
},
},
},
},
},
},
},
}
_, err := mutator.TranslateNotebookPaths().Apply(context.Background(), bundle)
require.NoError(t, err)
// Assert that the notebook artifact was defined.
assert.Len(t, bundle.Config.Artifacts, 2)
for _, artifact := range bundle.Config.Artifacts {
assert.Contains(t, artifact.Notebook.Path, "notebook.py")
}
// Assert that the path in the tasks now refer to the artifact.
assert.Equal(
t,
"${artifacts.my_job_notebook_py.notebook.remote_path}",
bundle.Config.Resources.Jobs["job"].Tasks[0].NotebookTask.NotebookPath,
)
assert.Equal(
t,
"./doesnt_exist.py",
bundle.Config.Resources.Jobs["job"].Tasks[1].NotebookTask.NotebookPath,
)
assert.Equal(
t,
"${artifacts.my_job_notebook_py.notebook.remote_path}",
bundle.Config.Resources.Jobs["job"].Tasks[2].NotebookTask.NotebookPath,
)
// Assert that the path in the libraries now refer to the artifact.
assert.Equal(
t,
"${artifacts.my_pipeline_notebook_py.notebook.remote_path}",
bundle.Config.Resources.Pipelines["pipeline"].Libraries[0].Notebook.Path,
)
assert.Equal(
t,
"./doesnt_exist.py",
bundle.Config.Resources.Pipelines["pipeline"].Libraries[1].Notebook.Path,
)
assert.Equal(
t,
"${artifacts.my_pipeline_notebook_py.notebook.remote_path}",
bundle.Config.Resources.Pipelines["pipeline"].Libraries[2].Notebook.Path,
)
}