databricks-cli/bundle/config/mutator/trampoline_test.go

package mutator

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"testing"

	"github.com/databricks/cli/bundle"
	"github.com/databricks/cli/bundle/config"
	"github.com/databricks/cli/bundle/config/paths"
	"github.com/databricks/cli/bundle/config/resources"
	"github.com/databricks/databricks-sdk-go/service/jobs"
	"github.com/stretchr/testify/require"
)

type functions struct{}

func (f *functions) GetTasks(b *bundle.Bundle) []TaskWithJobKey {
	tasks := make([]TaskWithJobKey, 0)
	for k := range b.Config.Resources.Jobs["test"].Tasks {
		tasks = append(tasks, TaskWithJobKey{
			JobKey: "test",
			Task:   &b.Config.Resources.Jobs["test"].Tasks[k],
		})
	}

	return tasks
}

func (f *functions) GetTemplateData(task *jobs.Task) (map[string]any, error) {
	if task.PythonWheelTask == nil {
		return nil, fmt.Errorf("PythonWheelTask cannot be nil")
	}

	data := make(map[string]any)
	data["MyName"] = "Trampoline"
	return data, nil
}

func (f *functions) CleanUp(task *jobs.Task) error {
	task.PythonWheelTask = nil
	return nil
}

func TestGenerateTrampoline(t *testing.T) {
	tmpDir := t.TempDir()

	tasks := []jobs.Task{
		{
			TaskKey: "to_trampoline",
			PythonWheelTask: &jobs.PythonWheelTask{
				PackageName: "test",
				EntryPoint:  "run",
			}},
	}

	b := &bundle.Bundle{
		Config: config.Root{
			Path: tmpDir,
			Bundle: config.Bundle{
				Target: "development",
			},
			Resources: config.Resources{
				Jobs: map[string]*resources.Job{
					"test": {
						Paths: paths.Paths{
							ConfigFilePath: tmpDir,
						},
						JobSettings: &jobs.JobSettings{
							Tasks: tasks,
						},
					},
				},
			},
		},
	}
	ctx := context.Background()

	funcs := functions{}
	trampoline := NewTrampoline("test_trampoline", &funcs, "Hello from {{.MyName}}")
	err := bundle.Apply(ctx, b, trampoline)
	require.NoError(t, err)

	dir, err := b.InternalDir(ctx)
	require.NoError(t, err)
	filename := filepath.Join(dir, "notebook_test_to_trampoline.py")

	bytes, err := os.ReadFile(filename)
	require.NoError(t, err)

	require.Equal(t, "Hello from Trampoline", string(bytes))

	task := b.Config.Resources.Jobs["test"].Tasks[0]
	require.Equal(t, task.NotebookTask.NotebookPath, ".databricks/bundle/development/.internal/notebook_test_to_trampoline")
	require.Nil(t, task.PythonWheelTask)
}