databricks-cli/bundle/deploy/terraform/convert_test.go

package terraform

import (
	"testing"

	"github.com/databricks/cli/bundle/config"
	"github.com/databricks/cli/bundle/config/resources"
	"github.com/databricks/databricks-sdk-go/service/catalog"
	"github.com/databricks/databricks-sdk-go/service/compute"
	"github.com/databricks/databricks-sdk-go/service/jobs"
	"github.com/databricks/databricks-sdk-go/service/ml"
	"github.com/databricks/databricks-sdk-go/service/pipelines"
	"github.com/databricks/databricks-sdk-go/service/serving"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestConvertJob(t *testing.T) {
	var src = resources.Job{
		JobSettings: &jobs.JobSettings{
			Name: "my job",
			JobClusters: []jobs.JobCluster{
				{
					JobClusterKey: "key",
					NewCluster: &compute.ClusterSpec{
						SparkVersion: "10.4.x-scala2.12",
					},
				},
			},
			GitSource: &jobs.GitSource{
				GitProvider: jobs.GitProviderGitHub,
				GitUrl:      "https://github.com/foo/bar",
			},
			Parameters: []jobs.JobParameterDefinition{
				{
					Name:    "param1",
					Default: "default1",
				},
				{
					Name:    "param2",
					Default: "default2",
				},
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Jobs: map[string]*resources.Job{
				"my_job": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.Equal(t, "my job", out.Resource.Job["my_job"].Name)
	assert.Len(t, out.Resource.Job["my_job"].JobCluster, 1)
	assert.Equal(t, "https://github.com/foo/bar", out.Resource.Job["my_job"].GitSource.Url)
	assert.Len(t, out.Resource.Job["my_job"].Parameter, 2)
	assert.Equal(t, "param1", out.Resource.Job["my_job"].Parameter[0].Name)
	assert.Equal(t, "param2", out.Resource.Job["my_job"].Parameter[1].Name)
	assert.Nil(t, out.Data)
}

func TestConvertJobPermissions(t *testing.T) {
	var src = resources.Job{
		Permissions: []resources.Permission{
			{
				Level:    "CAN_VIEW",
				UserName: "jane@doe.com",
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Jobs: map[string]*resources.Job{
				"my_job": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.NotEmpty(t, out.Resource.Permissions["job_my_job"].JobId)
	assert.Len(t, out.Resource.Permissions["job_my_job"].AccessControl, 1)

	p := out.Resource.Permissions["job_my_job"].AccessControl[0]
	assert.Equal(t, "jane@doe.com", p.UserName)
	assert.Equal(t, "CAN_VIEW", p.PermissionLevel)
}

func TestConvertJobTaskLibraries(t *testing.T) {
	var src = resources.Job{
		JobSettings: &jobs.JobSettings{
			Name: "my job",
			Tasks: []jobs.Task{
				{
					TaskKey: "key",
					Libraries: []compute.Library{
						{
							Pypi: &compute.PythonPyPiLibrary{
								Package: "mlflow",
							},
						},
					},
				},
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Jobs: map[string]*resources.Job{
				"my_job": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.Equal(t, "my job", out.Resource.Job["my_job"].Name)
	require.Len(t, out.Resource.Job["my_job"].Task, 1)
	require.Len(t, out.Resource.Job["my_job"].Task[0].Library, 1)
	assert.Equal(t, "mlflow", out.Resource.Job["my_job"].Task[0].Library[0].Pypi.Package)
}

func TestConvertPipeline(t *testing.T) {
	var src = resources.Pipeline{
		PipelineSpec: &pipelines.PipelineSpec{
			Name: "my pipeline",
			Libraries: []pipelines.PipelineLibrary{
				{
					Notebook: &pipelines.NotebookLibrary{
						Path: "notebook path",
					},
				},
				{
					File: &pipelines.FileLibrary{
						Path: "file path",
					},
				},
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Pipelines: map[string]*resources.Pipeline{
				"my_pipeline": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.Equal(t, "my pipeline", out.Resource.Pipeline["my_pipeline"].Name)
	assert.Len(t, out.Resource.Pipeline["my_pipeline"].Library, 2)
	assert.Nil(t, out.Data)
}

func TestConvertPipelinePermissions(t *testing.T) {
	var src = resources.Pipeline{
		Permissions: []resources.Permission{
			{
				Level:    "CAN_VIEW",
				UserName: "jane@doe.com",
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Pipelines: map[string]*resources.Pipeline{
				"my_pipeline": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.NotEmpty(t, out.Resource.Permissions["pipeline_my_pipeline"].PipelineId)
	assert.Len(t, out.Resource.Permissions["pipeline_my_pipeline"].AccessControl, 1)

	p := out.Resource.Permissions["pipeline_my_pipeline"].AccessControl[0]
	assert.Equal(t, "jane@doe.com", p.UserName)
	assert.Equal(t, "CAN_VIEW", p.PermissionLevel)
}

func TestConvertModel(t *testing.T) {
	var src = resources.MlflowModel{
		Model: &ml.Model{
			Name:        "name",
			Description: "description",
			Tags: []ml.ModelTag{
				{
					Key:   "k1",
					Value: "v1",
				},
				{
					Key:   "k2",
					Value: "v2",
				},
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Models: map[string]*resources.MlflowModel{
				"my_model": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.Equal(t, "name", out.Resource.MlflowModel["my_model"].Name)
	assert.Equal(t, "description", out.Resource.MlflowModel["my_model"].Description)
	assert.Len(t, out.Resource.MlflowModel["my_model"].Tags, 2)
	assert.Equal(t, "k1", out.Resource.MlflowModel["my_model"].Tags[0].Key)
	assert.Equal(t, "v1", out.Resource.MlflowModel["my_model"].Tags[0].Value)
	assert.Equal(t, "k2", out.Resource.MlflowModel["my_model"].Tags[1].Key)
	assert.Equal(t, "v2", out.Resource.MlflowModel["my_model"].Tags[1].Value)
	assert.Nil(t, out.Data)
}

func TestConvertModelPermissions(t *testing.T) {
	var src = resources.MlflowModel{
		Permissions: []resources.Permission{
			{
				Level:    "CAN_READ",
				UserName: "jane@doe.com",
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Models: map[string]*resources.MlflowModel{
				"my_model": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.NotEmpty(t, out.Resource.Permissions["mlflow_model_my_model"].RegisteredModelId)
	assert.Len(t, out.Resource.Permissions["mlflow_model_my_model"].AccessControl, 1)

	p := out.Resource.Permissions["mlflow_model_my_model"].AccessControl[0]
	assert.Equal(t, "jane@doe.com", p.UserName)
	assert.Equal(t, "CAN_READ", p.PermissionLevel)
}

func TestConvertExperiment(t *testing.T) {
	var src = resources.MlflowExperiment{
		Experiment: &ml.Experiment{
			Name: "name",
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Experiments: map[string]*resources.MlflowExperiment{
				"my_experiment": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.Equal(t, "name", out.Resource.MlflowExperiment["my_experiment"].Name)
	assert.Nil(t, out.Data)
}

func TestConvertExperimentPermissions(t *testing.T) {
	var src = resources.MlflowExperiment{
		Permissions: []resources.Permission{
			{
				Level:    "CAN_READ",
				UserName: "jane@doe.com",
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			Experiments: map[string]*resources.MlflowExperiment{
				"my_experiment": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.NotEmpty(t, out.Resource.Permissions["mlflow_experiment_my_experiment"].ExperimentId)
	assert.Len(t, out.Resource.Permissions["mlflow_experiment_my_experiment"].AccessControl, 1)

	p := out.Resource.Permissions["mlflow_experiment_my_experiment"].AccessControl[0]
	assert.Equal(t, "jane@doe.com", p.UserName)
	assert.Equal(t, "CAN_READ", p.PermissionLevel)

}

func TestConvertModelServing(t *testing.T) {
	var src = resources.ModelServingEndpoint{
		CreateServingEndpoint: &serving.CreateServingEndpoint{
			Name: "name",
			Config: serving.EndpointCoreConfigInput{
				ServedModels: []serving.ServedModelInput{
					{
						ModelName:          "model_name",
						ModelVersion:       "1",
						ScaleToZeroEnabled: true,
						WorkloadSize:       "Small",
					},
				},
				TrafficConfig: &serving.TrafficConfig{
					Routes: []serving.Route{
						{
							ServedModelName:   "model_name-1",
							TrafficPercentage: 100,
						},
					},
				},
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			ModelServingEndpoints: map[string]*resources.ModelServingEndpoint{
				"my_model_serving_endpoint": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	resource := out.Resource.ModelServing["my_model_serving_endpoint"]
	assert.Equal(t, "name", resource.Name)
	assert.Equal(t, "model_name", resource.Config.ServedModels[0].ModelName)
	assert.Equal(t, "1", resource.Config.ServedModels[0].ModelVersion)
	assert.Equal(t, true, resource.Config.ServedModels[0].ScaleToZeroEnabled)
	assert.Equal(t, "Small", resource.Config.ServedModels[0].WorkloadSize)
	assert.Equal(t, "model_name-1", resource.Config.TrafficConfig.Routes[0].ServedModelName)
	assert.Equal(t, 100, resource.Config.TrafficConfig.Routes[0].TrafficPercentage)
	assert.Nil(t, out.Data)
}

func TestConvertModelServingPermissions(t *testing.T) {
	var src = resources.ModelServingEndpoint{
		Permissions: []resources.Permission{
			{
				Level:    "CAN_VIEW",
				UserName: "jane@doe.com",
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			ModelServingEndpoints: map[string]*resources.ModelServingEndpoint{
				"my_model_serving_endpoint": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.NotEmpty(t, out.Resource.Permissions["model_serving_my_model_serving_endpoint"].ServingEndpointId)
	assert.Len(t, out.Resource.Permissions["model_serving_my_model_serving_endpoint"].AccessControl, 1)

	p := out.Resource.Permissions["model_serving_my_model_serving_endpoint"].AccessControl[0]
	assert.Equal(t, "jane@doe.com", p.UserName)
	assert.Equal(t, "CAN_VIEW", p.PermissionLevel)

}

func TestConvertRegisteredModel(t *testing.T) {
	var src = resources.RegisteredModel{
		CreateRegisteredModelRequest: &catalog.CreateRegisteredModelRequest{
			Name:        "name",
			CatalogName: "catalog",
			SchemaName:  "schema",
			Comment:     "comment",
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			RegisteredModels: map[string]*resources.RegisteredModel{
				"my_registered_model": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	resource := out.Resource.RegisteredModel["my_registered_model"]
	assert.Equal(t, "name", resource.Name)
	assert.Equal(t, "catalog", resource.CatalogName)
	assert.Equal(t, "schema", resource.SchemaName)
	assert.Equal(t, "comment", resource.Comment)
	assert.Nil(t, out.Data)
}

func TestConvertRegisteredModelGrants(t *testing.T) {
	var src = resources.RegisteredModel{
		Grants: []resources.Grant{
			{
				Privileges: []string{"EXECUTE"},
				Principal:  "jane@doe.com",
			},
		},
	}

	var config = config.Root{
		Resources: config.Resources{
			RegisteredModels: map[string]*resources.RegisteredModel{
				"my_registered_model": &src,
			},
		},
	}

	out := BundleToTerraform(&config)
	assert.NotEmpty(t, out.Resource.Grants["registered_model_my_registered_model"].Function)
	assert.Len(t, out.Resource.Grants["registered_model_my_registered_model"].Grant, 1)

	p := out.Resource.Grants["registered_model_my_registered_model"].Grant[0]
	assert.Equal(t, "jane@doe.com", p.Principal)
	assert.Equal(t, "EXECUTE", p.Privileges[0])

}