Skip to content

Commit

Permalink
Add integration test for mlops-stacks initialization (#1155)
Browse files Browse the repository at this point in the history
## Changes
This PR:
1. Adds an integration test for mlops-stacks that checks the
initialization and deployment of the project was successful.
2. Fixes a bug in the initialization of templates from non-tty. We need
to process the input parameters in order since their descriptions can
refer to input parameters that came before in the interactive UX.

## Tests
The integration test passes in CI.
  • Loading branch information
shreyas-goenka authored Mar 12, 2024
1 parent c781856 commit d4329f4
Show file tree
Hide file tree
Showing 5 changed files with 115 additions and 30 deletions.
5 changes: 3 additions & 2 deletions bundle/bundle_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"testing"

"github.com/databricks/cli/bundle/env"
"github.com/databricks/cli/internal/testutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
Expand Down Expand Up @@ -86,7 +87,7 @@ func TestBundleMustLoadFailureWithEnv(t *testing.T) {
}

func TestBundleMustLoadFailureIfNotFound(t *testing.T) {
chdir(t, t.TempDir())
testutil.Chdir(t, t.TempDir())
_, err := MustLoad(context.Background())
require.Error(t, err, "unable to find bundle root")
}
Expand All @@ -105,7 +106,7 @@ func TestBundleTryLoadFailureWithEnv(t *testing.T) {
}

func TestBundleTryLoadOkIfNotFound(t *testing.T) {
chdir(t, t.TempDir())
testutil.Chdir(t, t.TempDir())
b, err := TryLoad(context.Background())
assert.NoError(t, err)
assert.Nil(t, b)
Expand Down
35 changes: 8 additions & 27 deletions bundle/root_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,30 +8,11 @@ import (

"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/bundle/env"
"github.com/databricks/cli/internal/testutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

// Changes into specified directory for the duration of the test.
// Returns the current working directory.
func chdir(t *testing.T, dir string) string {
wd, err := os.Getwd()
require.NoError(t, err)

abs, err := filepath.Abs(dir)
require.NoError(t, err)

err = os.Chdir(abs)
require.NoError(t, err)

t.Cleanup(func() {
err := os.Chdir(wd)
require.NoError(t, err)
})

return wd
}

func TestRootFromEnv(t *testing.T) {
ctx := context.Background()
dir := t.TempDir()
Expand Down Expand Up @@ -83,7 +64,7 @@ func TestRootLookup(t *testing.T) {
t.Setenv(env.RootVariable, "")
os.Unsetenv(env.RootVariable)

chdir(t, t.TempDir())
testutil.Chdir(t, t.TempDir())

// Create databricks.yml file.
f, err := os.Create(config.FileNames[0])
Expand All @@ -95,7 +76,7 @@ func TestRootLookup(t *testing.T) {
require.NoError(t, err)

// It should find the project root from $PWD.
wd := chdir(t, "./a/b/c")
wd := testutil.Chdir(t, "./a/b/c")
root, err := mustGetRoot(ctx)
require.NoError(t, err)
require.Equal(t, wd, root)
Expand All @@ -109,14 +90,14 @@ func TestRootLookupError(t *testing.T) {
os.Unsetenv(env.RootVariable)

// It can't find a project root from a temporary directory.
_ = chdir(t, t.TempDir())
_ = testutil.Chdir(t, t.TempDir())
_, err := mustGetRoot(ctx)
require.ErrorContains(t, err, "unable to locate bundle root")
}

func TestLoadYamlWhenIncludesEnvPresent(t *testing.T) {
ctx := context.Background()
chdir(t, filepath.Join(".", "tests", "basic"))
testutil.Chdir(t, filepath.Join(".", "tests", "basic"))
t.Setenv(env.IncludesVariable, "test")

bundle, err := MustLoad(ctx)
Expand All @@ -131,7 +112,7 @@ func TestLoadYamlWhenIncludesEnvPresent(t *testing.T) {
func TestLoadDefautlBundleWhenNoYamlAndRootAndIncludesEnvPresent(t *testing.T) {
ctx := context.Background()
dir := t.TempDir()
chdir(t, dir)
testutil.Chdir(t, dir)
t.Setenv(env.RootVariable, dir)
t.Setenv(env.IncludesVariable, "test")

Expand All @@ -143,7 +124,7 @@ func TestLoadDefautlBundleWhenNoYamlAndRootAndIncludesEnvPresent(t *testing.T) {
func TestErrorIfNoYamlNoRootEnvAndIncludesEnvPresent(t *testing.T) {
ctx := context.Background()
dir := t.TempDir()
chdir(t, dir)
testutil.Chdir(t, dir)
t.Setenv(env.IncludesVariable, "test")

_, err := MustLoad(ctx)
Expand All @@ -153,7 +134,7 @@ func TestErrorIfNoYamlNoRootEnvAndIncludesEnvPresent(t *testing.T) {
func TestErrorIfNoYamlNoIncludesEnvAndRootEnvPresent(t *testing.T) {
ctx := context.Background()
dir := t.TempDir()
chdir(t, dir)
testutil.Chdir(t, dir)
t.Setenv(env.RootVariable, dir)

_, err := MustLoad(ctx)
Expand Down
77 changes: 77 additions & 0 deletions internal/init_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,15 @@ package internal

import (
"context"
"encoding/json"
"fmt"
"os"
"path/filepath"
"strconv"
"testing"

"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/internal/testutil"
"github.com/databricks/cli/libs/auth"
"github.com/databricks/databricks-sdk-go"
"github.com/stretchr/testify/assert"
Expand All @@ -21,6 +25,79 @@ func TestAccBundleInitErrorOnUnknownFields(t *testing.T) {
assert.EqualError(t, err, "failed to compute file content for bar.tmpl. variable \"does_not_exist\" not defined")
}

// This test tests the MLOps Stacks DAB e2e and thus there's a couple of special
// considerations to take note of:
//
// 1. Upstream changes to the MLOps Stacks DAB can cause this test to fail.
// In which case we should do one of:
// (a) Update this test to reflect the changes
// (b) Update the MLOps Stacks DAB to not break this test. Skip this test
// temporarily until the MLOps Stacks DAB is updated
//
// 2. While rare and to be avoided if possible, the CLI reserves the right to
// make changes that can break the MLOps Stacks DAB. In which case we should
// skip this test until the MLOps Stacks DAB is updated to work again.
func TestAccBundleInitOnMlopsStacks(t *testing.T) {
t.Parallel()
env := GetEnvOrSkipTest(t, "CLOUD_ENV")
tmpDir1 := t.TempDir()
tmpDir2 := t.TempDir()

w, err := databricks.NewWorkspaceClient(&databricks.Config{})
require.NoError(t, err)

projectName := RandomName("project_name_")

// Create a config file with the project name and root dir
initConfig := map[string]string{
"input_project_name": projectName,
"input_root_dir": "repo_name",
"input_include_models_in_unity_catalog": "no",
"input_cloud": env,
}
b, err := json.Marshal(initConfig)
require.NoError(t, err)
os.WriteFile(filepath.Join(tmpDir1, "config.json"), b, 0644)

// Run bundle init
assert.NoFileExists(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md"))
RequireSuccessfulRun(t, "bundle", "init", "mlops-stacks", "--output-dir", tmpDir2, "--config-file", filepath.Join(tmpDir1, "config.json"))

// Assert that the README.md file was created
assert.FileExists(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md"))
assertLocalFileContents(t, filepath.Join(tmpDir2, "repo_name", projectName, "README.md"), fmt.Sprintf("# %s", projectName))

// Validate the stack
testutil.Chdir(t, filepath.Join(tmpDir2, "repo_name", projectName))
RequireSuccessfulRun(t, "bundle", "validate")

// Deploy the stack
RequireSuccessfulRun(t, "bundle", "deploy")
t.Cleanup(func() {
// Delete the stack
RequireSuccessfulRun(t, "bundle", "destroy", "--auto-approve")
})

// Get summary of the bundle deployment
stdout, _ := RequireSuccessfulRun(t, "bundle", "summary", "--output", "json")
summary := &config.Root{}
err = json.Unmarshal(stdout.Bytes(), summary)
require.NoError(t, err)

// Assert resource Ids are not empty
assert.NotEmpty(t, summary.Resources.Experiments["experiment"].ID)
assert.NotEmpty(t, summary.Resources.Models["model"].ID)
assert.NotEmpty(t, summary.Resources.Jobs["batch_inference_job"].ID)
assert.NotEmpty(t, summary.Resources.Jobs["model_training_job"].ID)

// Assert the batch inference job actually exists
batchJobId, err := strconv.ParseInt(summary.Resources.Jobs["batch_inference_job"].ID, 10, 64)
require.NoError(t, err)
job, err := w.Jobs.GetByJobId(context.Background(), batchJobId)
assert.NoError(t, err)
assert.Equal(t, fmt.Sprintf("dev-%s-batch-inference-job", projectName), job.Settings.Name)
}

func TestAccBundleInitHelpers(t *testing.T) {
env := GetEnvOrSkipTest(t, "CLOUD_ENV")
t.Log(env)
Expand Down
23 changes: 23 additions & 0 deletions internal/testutil/env.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,12 @@ package testutil

import (
"os"
"path/filepath"
"runtime"
"strings"
"testing"

"github.com/stretchr/testify/require"
)

// CleanupEnvironment sets up a pristine environment containing only $PATH and $HOME.
Expand Down Expand Up @@ -44,3 +47,23 @@ func GetEnvOrSkipTest(t *testing.T, name string) string {
}
return value
}

// Changes into specified directory for the duration of the test.
// Returns the current working directory.
func Chdir(t *testing.T, dir string) string {
wd, err := os.Getwd()
require.NoError(t, err)

abs, err := filepath.Abs(dir)
require.NoError(t, err)

err = os.Chdir(abs)
require.NoError(t, err)

t.Cleanup(func() {
err := os.Chdir(wd)
require.NoError(t, err)
})

return wd
}
5 changes: 4 additions & 1 deletion libs/template/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,10 @@ func (c *config) assignValuesFromFile(path string) error {

// Assigns default values from schema to input config map
func (c *config) assignDefaultValues(r *renderer) error {
for name, property := range c.schema.Properties {
for _, p := range c.schema.OrderedProperties() {
name := p.Name
property := p.Schema

// Config already has a value assigned
if _, ok := c.values[name]; ok {
continue
Expand Down

0 comments on commit d4329f4

Please sign in to comment.