steve/databricks-cli
steve
/
databricks-cli
mirror of https://github.com/databricks/cli.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use `acc.WorkspaceTest` helper from bundle integration tests (#1181) 

## Changes

This helper:
* Constructs a context
* Constructs a `*databricks.WorkspaceClient`
* Ensures required environment variables are present to run an
integration test
* Enables debugging integration tests from VS Code

Debugging integration tests (from VS Code) is made possible by a prelude
in the helper that checks if the calling process is a debug binary, and
if so, sources environment variables from
`~/..databricks/debug-env.json` (if present).

## Tests

Integration tests still pass.

---------

Co-authored-by: Andrew Nester <andrew.nester@databricks.com>
This commit is contained in:
Pieter Noordhuis 2024-02-07 12:18:56 +01:00 committed by GitHub
parent 6edab93233
commit f8b0f783ea
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 91 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
internal/bundle/deploy_then_remove_resources_test.go
View File

@ -1,38 +1,33 @@
package bundle
import (
"context"
"os"
"path/filepath"
"testing"
"github.com/databricks/cli/internal"
"github.com/databricks/databricks-sdk-go"
"github.com/databricks/cli/internal/acc"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAccBundleDeployThenRemoveResources(t *testing.T) {
env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV")
t.Log(env)
ctx, wt := acc.WorkspaceTest(t)
w := wt.W
uniqueId := uuid.New().String()
bundleRoot, err := initTestTemplate(t, "deploy_then_remove_resources", map[string]any{
bundleRoot, err := initTestTemplate(t, ctx, "deploy_then_remove_resources", map[string]any{
"unique_id": uniqueId,
})
require.NoError(t, err)
// deploy pipeline
err = deployBundle(t, bundleRoot)
require.NoError(t, err)
w, err := databricks.NewWorkspaceClient()
err = deployBundle(t, ctx, bundleRoot)
require.NoError(t, err)
// assert pipeline is created
pipelineName := "test-bundle-pipeline-" + uniqueId
pipeline, err := w.Pipelines.GetByName(context.Background(), pipelineName)
pipeline, err := w.Pipelines.GetByName(ctx, pipelineName)
require.NoError(t, err)
assert.Equal(t, pipeline.Name, pipelineName)
@ -41,15 +36,15 @@ func TestAccBundleDeployThenRemoveResources(t *testing.T) {
require.NoError(t, err)
// deploy again
err = deployBundle(t, bundleRoot)
err = deployBundle(t, ctx, bundleRoot)
require.NoError(t, err)
// assert pipeline is deleted
_, err = w.Pipelines.GetByName(context.Background(), pipelineName)
_, err = w.Pipelines.GetByName(ctx, pipelineName)
assert.ErrorContains(t, err, "does not exist")
t.Cleanup(func() {
err = destroyBundle(t, bundleRoot)
err = destroyBundle(t, ctx, bundleRoot)
require.NoError(t, err)
})
}

9
internal/bundle/empty_bundle_test.go
View File

@ -6,14 +6,13 @@ import (
"path/filepath"
"testing"
"github.com/databricks/cli/internal"
"github.com/databricks/cli/internal/acc"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
)
func TestAccEmptyBundleDeploy(t *testing.T) {
env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV")
t.Log(env)
ctx, _ := acc.WorkspaceTest(t)
// create empty bundle
tmpDir := t.TempDir()
@ -27,11 +26,11 @@ func TestAccEmptyBundleDeploy(t *testing.T) {
f.Close()
// deploy empty bundle
err = deployBundle(t, tmpDir)
err = deployBundle(t, ctx, tmpDir)
require.NoError(t, err)
t.Cleanup(func() {
err = destroyBundle(t, tmpDir)
err = destroyBundle(t, ctx, tmpDir)
require.NoError(t, err)
})
}

43
internal/bundle/generate_job_test.go
View File

@ -10,6 +10,7 @@ import (
"testing"
"github.com/databricks/cli/internal"
"github.com/databricks/cli/internal/acc"
"github.com/databricks/cli/internal/testutil"
"github.com/databricks/cli/libs/filer"
"github.com/databricks/databricks-sdk-go"
@ -20,23 +21,22 @@ import (
)
func TestAccGenerateFromExistingJobAndDeploy(t *testing.T) {
env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV")
t.Log(env)
ctx, wt := acc.WorkspaceTest(t)
gt := &generateJobTest{T: t, w: wt.W}
uniqueId := uuid.New().String()
bundleRoot, err := initTestTemplate(t, "with_includes", map[string]any{
bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{
"unique_id": uniqueId,
})
require.NoError(t, err)
jobId := createTestJob(t)
jobId := gt.createTestJob(ctx)
t.Cleanup(func() {
destroyJob(t, jobId)
require.NoError(t, err)
gt.destroyJob(ctx, jobId)
})
t.Setenv("BUNDLE_ROOT", bundleRoot)
c := internal.NewCobraTestRunner(t, "bundle", "generate", "job",
c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "generate", "job",
"--existing-job-id", fmt.Sprint(jobId),
"--config-dir", filepath.Join(bundleRoot, "resources"),
"--source-dir", filepath.Join(bundleRoot, "src"))
@ -61,15 +61,22 @@ func TestAccGenerateFromExistingJobAndDeploy(t *testing.T) {
require.Contains(t, generatedYaml, "spark_version: 13.3.x-scala2.12")
require.Contains(t, generatedYaml, "num_workers: 1")
err = deployBundle(t, bundleRoot)
err = deployBundle(t, ctx, bundleRoot)
require.NoError(t, err)
err = destroyBundle(t, bundleRoot)
err = destroyBundle(t, ctx, bundleRoot)
require.NoError(t, err)
}
func createTestJob(t *testing.T) int64 {
type generateJobTest struct {
T *testing.T
w *databricks.WorkspaceClient
}
func (gt *generateJobTest) createTestJob(ctx context.Context) int64 {
t := gt.T
w := gt.w
var nodeTypeId string
switch testutil.GetCloud(t) {
case testutil.AWS:
@ -80,10 +87,6 @@ func createTestJob(t *testing.T) int64 {
nodeTypeId = "n1-standard-4"
}
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)
ctx := context.Background()
tmpdir := internal.TemporaryWorkspaceDir(t, w)
f, err := filer.NewWorkspaceFilesClient(w, tmpdir)
require.NoError(t, err)
@ -112,13 +115,9 @@ func createTestJob(t *testing.T) int64 {
return resp.JobId
}
func destroyJob(t *testing.T, jobId int64) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)
ctx := context.Background()
err = w.Jobs.Delete(ctx, jobs.DeleteJob{
func (gt *generateJobTest) destroyJob(ctx context.Context, jobId int64) {
err := gt.w.Jobs.Delete(ctx, jobs.DeleteJob{
JobId: jobId,
})
require.NoError(t, err)
require.NoError(gt.T, err)
}

40
internal/bundle/generate_pipeline_test.go
View File

@ -10,6 +10,7 @@ import (
"testing"
"github.com/databricks/cli/internal"
"github.com/databricks/cli/internal/acc"
"github.com/databricks/cli/libs/filer"
"github.com/databricks/databricks-sdk-go"
"github.com/databricks/databricks-sdk-go/service/pipelines"
@ -18,23 +19,22 @@ import (
)
func TestAccGenerateFromExistingPipelineAndDeploy(t *testing.T) {
env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV")
t.Log(env)
ctx, wt := acc.WorkspaceTest(t)
gt := &generatePipelineTest{T: t, w: wt.W}
uniqueId := uuid.New().String()
bundleRoot, err := initTestTemplate(t, "with_includes", map[string]any{
bundleRoot, err := initTestTemplate(t, ctx, "with_includes", map[string]any{
"unique_id": uniqueId,
})
require.NoError(t, err)
pipelineId := createTestPipeline(t)
pipelineId := gt.createTestPipeline(ctx)
t.Cleanup(func() {
destroyPipeline(t, pipelineId)
require.NoError(t, err)
gt.destroyPipeline(ctx, pipelineId)
})
t.Setenv("BUNDLE_ROOT", bundleRoot)
c := internal.NewCobraTestRunner(t, "bundle", "generate", "pipeline",
c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "generate", "pipeline",
"--existing-pipeline-id", fmt.Sprint(pipelineId),
"--config-dir", filepath.Join(bundleRoot, "resources"),
"--source-dir", filepath.Join(bundleRoot, "src"))
@ -61,18 +61,22 @@ func TestAccGenerateFromExistingPipelineAndDeploy(t *testing.T) {
require.Contains(t, generatedYaml, "- file:")
require.Contains(t, generatedYaml, fmt.Sprintf("path: %s", filepath.Join("..", "src", "test.py")))
err = deployBundle(t, bundleRoot)
err = deployBundle(t, ctx, bundleRoot)
require.NoError(t, err)
err = destroyBundle(t, bundleRoot)
err = destroyBundle(t, ctx, bundleRoot)
require.NoError(t, err)
}
func createTestPipeline(t *testing.T) string {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)
type generatePipelineTest struct {
T *testing.T
w *databricks.WorkspaceClient
}
func (gt *generatePipelineTest) createTestPipeline(ctx context.Context) string {
t := gt.T
w := gt.w
ctx := context.Background()
tmpdir := internal.TemporaryWorkspaceDir(t, w)
f, err := filer.NewWorkspaceFilesClient(w, tmpdir)
require.NoError(t, err)
@ -103,13 +107,9 @@ func createTestPipeline(t *testing.T) string {
return resp.PipelineId
}
func destroyPipeline(t *testing.T, pipelineId string) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)
ctx := context.Background()
err = w.Pipelines.Delete(ctx, pipelines.DeletePipelineRequest{
func (gt *generatePipelineTest) destroyPipeline(ctx context.Context, pipelineId string) {
err := gt.w.Pipelines.Delete(ctx, pipelines.DeletePipelineRequest{
PipelineId: pipelineId,
})
require.NoError(t, err)
require.NoError(gt.T, err)
}

18
internal/bundle/helpers.go
View File

@ -15,7 +15,7 @@ import (
"github.com/databricks/cli/libs/template"
)
func initTestTemplate(t *testing.T, templateName string, config map[string]any) (string, error) {
func initTestTemplate(t *testing.T, ctx context.Context, templateName string, config map[string]any) (string, error) {
templateRoot := filepath.Join("bundles", templateName)
bundleRoot := t.TempDir()
@ -24,7 +24,7 @@ func initTestTemplate(t *testing.T, templateName string, config map[string]any)
return "", err
}
ctx := root.SetWorkspaceClient(context.Background(), nil)
ctx = root.SetWorkspaceClient(ctx, nil)
cmd := cmdio.NewIO(flags.OutputJSON, strings.NewReader(""), os.Stdout, os.Stderr, "bundles")
ctx = cmdio.InContext(ctx, cmd)
@ -46,15 +46,14 @@ func writeConfigFile(t *testing.T, config map[string]any) (string, error) {
return filepath, err
}
func deployBundle(t *testing.T, path string) error {
func deployBundle(t *testing.T, ctx context.Context, path string) error {
t.Setenv("BUNDLE_ROOT", path)
c := internal.NewCobraTestRunner(t, "bundle", "deploy", "--force-lock")
c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "deploy", "--force-lock")
_, _, err := c.Run()
return err
}
func runResource(t *testing.T, path string, key string) (string, error) {
ctx := context.Background()
func runResource(t *testing.T, ctx context.Context, path string, key string) (string, error) {
ctx = cmdio.NewContext(ctx, cmdio.Default())
c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "run", key)
@ -62,8 +61,7 @@ func runResource(t *testing.T, path string, key string) (string, error) {
return stdout.String(), err
}
func runResourceWithParams(t *testing.T, path string, key string, params ...string) (string, error) {
ctx := context.Background()
func runResourceWithParams(t *testing.T, ctx context.Context, path string, key string, params ...string) (string, error) {
ctx = cmdio.NewContext(ctx, cmdio.Default())
args := make([]string, 0)
@ -74,9 +72,9 @@ func runResourceWithParams(t *testing.T, path string, key string, params ...stri
return stdout.String(), err
}
func destroyBundle(t *testing.T, path string) error {
func destroyBundle(t *testing.T, ctx context.Context, path string) error {
t.Setenv("BUNDLE_ROOT", path)
c := internal.NewCobraTestRunner(t, "bundle", "destroy", "--auto-approve")
c := internal.NewCobraTestRunnerWithContext(t, ctx, "bundle", "destroy", "--auto-approve")
_, _, err := c.Run()
return err
}

18
internal/bundle/job_metadata_test.go
View File

@ -12,23 +12,21 @@ import (
"github.com/databricks/cli/bundle/config"
"github.com/databricks/cli/bundle/metadata"
"github.com/databricks/cli/internal"
"github.com/databricks/cli/internal/acc"
"github.com/databricks/cli/libs/env"
"github.com/databricks/cli/libs/filer"
"github.com/databricks/databricks-sdk-go"
"github.com/google/uuid"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAccJobsMetadataFile(t *testing.T) {
env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV")
t.Log(env)
ctx, wt := acc.WorkspaceTest(t)
w := wt.W
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)
nodeTypeId := internal.GetNodeTypeId(env)
nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
uniqueId := uuid.New().String()
bundleRoot, err := initTestTemplate(t, "job_metadata", map[string]any{
bundleRoot, err := initTestTemplate(t, ctx, "job_metadata", map[string]any{
"unique_id": uniqueId,
"node_type_id": nodeTypeId,
"spark_version": "13.2.x-snapshot-scala2.12",
@ -36,12 +34,12 @@ func TestAccJobsMetadataFile(t *testing.T) {
require.NoError(t, err)
// deploy bundle
err = deployBundle(t, bundleRoot)
err = deployBundle(t, ctx, bundleRoot)
require.NoError(t, err)
// Cleanup the deployed bundle
t.Cleanup(func() {
err = destroyBundle(t, bundleRoot)
err = destroyBundle(t, ctx, bundleRoot)
require.NoError(t, err)
})

24
internal/bundle/local_state_staleness_test.go
View File

@ -5,7 +5,8 @@ import (
"testing"
"github.com/databricks/cli/internal"
"github.com/databricks/databricks-sdk-go"
"github.com/databricks/cli/internal/acc"
"github.com/databricks/cli/libs/env"
"github.com/databricks/databricks-sdk-go/listing"
"github.com/databricks/databricks-sdk-go/service/jobs"
"github.com/google/uuid"
@ -14,11 +15,8 @@ import (
)
func TestAccLocalStateStaleness(t *testing.T) {
env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV")
t.Log(env)
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)
ctx, wt := acc.WorkspaceTest(t)
w := wt.W
// The approach for this test is as follows:
// 1) First deploy of bundle instance A
@ -27,10 +25,10 @@ func TestAccLocalStateStaleness(t *testing.T) {
// Because of deploy (2), the locally cached state of bundle instance A should be stale.
// Then for deploy (3), it must use the remote state over the stale local state.
nodeTypeId := internal.GetNodeTypeId(env)
nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
uniqueId := uuid.New().String()
initialize := func() string {
root, err := initTestTemplate(t, "basic", map[string]any{
root, err := initTestTemplate(t, ctx, "basic", map[string]any{
"unique_id": uniqueId,
"node_type_id": nodeTypeId,
"spark_version": "13.2.x-snapshot-scala2.12",
@ -38,26 +36,28 @@ func TestAccLocalStateStaleness(t *testing.T) {
require.NoError(t, err)
t.Cleanup(func() {
err = destroyBundle(t, root)
err = destroyBundle(t, ctx, root)
require.NoError(t, err)
})
return root
}
var err error
bundleA := initialize()
bundleB := initialize()
// 1) Deploy bundle A
err = deployBundle(t, bundleA)
err = deployBundle(t, ctx, bundleA)
require.NoError(t, err)
// 2) Deploy bundle B
err = deployBundle(t, bundleB)
err = deployBundle(t, ctx, bundleB)
require.NoError(t, err)
// 3) Deploy bundle A again
err = deployBundle(t, bundleA)
err = deployBundle(t, ctx, bundleA)
require.NoError(t, err)
// Assert that there is only a single job in the workspace corresponding to this bundle.

25
internal/bundle/python_wheel_test.go
View File

@ -4,24 +4,17 @@ import (
"testing"
"github.com/databricks/cli/internal"
"github.com/databricks/cli/internal/acc"
"github.com/databricks/cli/libs/env"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
)
func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bool) {
env := internal.GetEnvOrSkipTest(t, "CLOUD_ENV")
t.Log(env)
ctx, _ := acc.WorkspaceTest(t)
var nodeTypeId string
if env == "gcp" {
nodeTypeId = "n1-standard-4"
} else if env == "aws" {
nodeTypeId = "i3.xlarge"
} else {
nodeTypeId = "Standard_DS4_v2"
}
bundleRoot, err := initTestTemplate(t, "python_wheel_task", map[string]any{
nodeTypeId := internal.GetNodeTypeId(env.Get(ctx, "CLOUD_ENV"))
bundleRoot, err := initTestTemplate(t, ctx, "python_wheel_task", map[string]any{
"node_type_id": nodeTypeId,
"unique_id": uuid.New().String(),
"spark_version": sparkVersion,
@ -29,20 +22,20 @@ func runPythonWheelTest(t *testing.T, sparkVersion string, pythonWheelWrapper bo
})
require.NoError(t, err)
err = deployBundle(t, bundleRoot)
err = deployBundle(t, ctx, bundleRoot)
require.NoError(t, err)
t.Cleanup(func() {
destroyBundle(t, bundleRoot)
destroyBundle(t, ctx, bundleRoot)
})
out, err := runResource(t, bundleRoot, "some_other_job")
out, err := runResource(t, ctx, bundleRoot, "some_other_job")
require.NoError(t, err)
require.Contains(t, out, "Hello from my func")
require.Contains(t, out, "Got arguments:")
require.Contains(t, out, "['my_test_code', 'one', 'two']")
out, err = runResourceWithParams(t, bundleRoot, "some_other_job", "--python-params=param1,param2")
out, err = runResourceWithParams(t, ctx, bundleRoot, "some_other_job", "--python-params=param1,param2")
require.NoError(t, err)
require.Contains(t, out, "Hello from my func")
require.Contains(t, out, "Got arguments:")