Skip to content

Commit

Permalink
Added test to submit and run various Python tasks on multiple DBR ver…
Browse files Browse the repository at this point in the history 
…sions (#806)

## Changes
These tests allow us to get information for execution context
(PYTHONPATH, CWD) for various Python tasks and different cluster setups.

Note: this test won't be executed automatically as part of nightly
builds since it requires RUN_PYTHON_TASKS_TEST env to be executed.

## Tests
Integration test run successfully.

---------

Co-authored-by: Pieter Noordhuis <pieter.noordhuis@databricks.com>
  • Loading branch information
@andrewnester @pietern
andrewnester and pietern authored Oct 3, 2023
1 parent 452565c commit 79e271f
Show file tree
Hide file tree
Showing 14 changed files with 465 additions and 79 deletions.
61 changes: 2 additions & 59 deletions internal/filer_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import (
"bytes"
"context"
"errors"
"fmt"
"io"
"io/fs"
"net/http"
Expand All @@ -15,8 +14,6 @@ import (
"github.com/databricks/cli/libs/filer"
"github.com/databricks/databricks-sdk-go"
"github.com/databricks/databricks-sdk-go/apierr"
"github.com/databricks/databricks-sdk-go/service/files"
"github.com/databricks/databricks-sdk-go/service/workspace"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
Expand Down Expand Up @@ -209,41 +206,12 @@ func runFilerReadDirTest(t *testing.T, ctx context.Context, f filer.Filer) {
assert.False(t, entries[0].IsDir())
}

func temporaryWorkspaceDir(t *testing.T, w *databricks.WorkspaceClient) string {
ctx := context.Background()
me, err := w.CurrentUser.Me(ctx)
require.NoError(t, err)

path := fmt.Sprintf("/Users/%s/%s", me.UserName, RandomName("integration-test-wsfs-"))

// Ensure directory exists, but doesn't exist YET!
// Otherwise we could inadvertently remove a directory that already exists on cleanup.
t.Logf("mkdir %s", path)
err = w.Workspace.MkdirsByPath(ctx, path)
require.NoError(t, err)

// Remove test directory on test completion.
t.Cleanup(func() {
t.Logf("rm -rf %s", path)
err := w.Workspace.Delete(ctx, workspace.Delete{
Path: path,
Recursive: true,
})
if err == nil || apierr.IsMissing(err) {
return
}
t.Logf("unable to remove temporary workspace directory %s: %#v", path, err)
})

return path
}

func setupWorkspaceFilesTest(t *testing.T) (context.Context, filer.Filer) {
t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))

ctx := context.Background()
w := databricks.Must(databricks.NewWorkspaceClient())
tmpdir := temporaryWorkspaceDir(t, w)
tmpdir := TemporaryWorkspaceDir(t, w)
f, err := filer.NewWorkspaceFilesClient(w, tmpdir)
require.NoError(t, err)

Expand All @@ -267,37 +235,12 @@ func TestAccFilerWorkspaceFilesReadDir(t *testing.T) {
runFilerReadDirTest(t, ctx, f)
}

func temporaryDbfsDir(t *testing.T, w *databricks.WorkspaceClient) string {
ctx := context.Background()
path := fmt.Sprintf("/tmp/%s", RandomName("integration-test-dbfs-"))

// This call fails if the path already exists.
t.Logf("mkdir dbfs:%s", path)
err := w.Dbfs.MkdirsByPath(ctx, path)
require.NoError(t, err)

// Remove test directory on test completion.
t.Cleanup(func() {
t.Logf("rm -rf dbfs:%s", path)
err := w.Dbfs.Delete(ctx, files.Delete{
Path: path,
Recursive: true,
})
if err == nil || apierr.IsMissing(err) {
return
}
t.Logf("unable to remove temporary dbfs directory %s: %#v", path, err)
})

return path
}

func setupFilerDbfsTest(t *testing.T) (context.Context, filer.Filer) {
t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV"))

ctx := context.Background()
w := databricks.Must(databricks.NewWorkspaceClient())
tmpdir := temporaryDbfsDir(t, w)
tmpdir := TemporaryDbfsDir(t, w)
f, err := filer.NewDbfsClient(w, tmpdir)
require.NoError(t, err)
return ctx, f
Expand Down
4 changes: 2 additions & 2 deletions internal/fs_cat_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ func TestAccFsCatForDbfs(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down Expand Up @@ -54,7 +54,7 @@ func TestAccFsCatDoesNotSupportOutputModeJson(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down
4 changes: 2 additions & 2 deletions internal/fs_cp_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ func setupDbfsFiler(t *testing.T) (filer.Filer, string) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)
f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)

Expand Down Expand Up @@ -256,7 +256,7 @@ func TestAccFsCpErrorsWhenSourceIsDirWithoutRecursiveFlag(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

_, _, err = RequireErrorRun(t, "fs", "cp", "dbfs:"+tmpDir, "dbfs:/tmp")
assert.Equal(t, fmt.Sprintf("source path %s is a directory. Please specify the --recursive flag", tmpDir), err.Error())
Expand Down
8 changes: 4 additions & 4 deletions internal/fs_ls_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ func TestAccFsLsForDbfs(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down Expand Up @@ -58,7 +58,7 @@ func TestAccFsLsForDbfsWithAbsolutePaths(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down Expand Up @@ -94,7 +94,7 @@ func TestAccFsLsForDbfsOnFile(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand All @@ -114,7 +114,7 @@ func TestAccFsLsForDbfsOnEmptyDir(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

stdout, stderr := RequireSuccessfulRun(t, "fs", "ls", "dbfs:"+tmpDir, "--output=json")
assert.Equal(t, "", stderr.String())
Expand Down
8 changes: 4 additions & 4 deletions internal/fs_mkdir_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ func TestAccFsMkdirCreatesDirectory(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand All @@ -44,7 +44,7 @@ func TestAccFsMkdirCreatesMultipleDirectories(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down Expand Up @@ -80,7 +80,7 @@ func TestAccFsMkdirWhenDirectoryAlreadyExists(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

// create directory "a"
f, err := filer.NewDbfsClient(w, tmpDir)
Expand All @@ -101,7 +101,7 @@ func TestAccFsMkdirWhenFileExistsAtPath(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

// create file hello
f, err := filer.NewDbfsClient(w, tmpDir)
Expand Down
8 changes: 4 additions & 4 deletions internal/fs_rm_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ func TestAccFsRmForFile(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down Expand Up @@ -52,7 +52,7 @@ func TestAccFsRmForEmptyDirectory(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down Expand Up @@ -84,7 +84,7 @@ func TestAccFsRmForNonEmptyDirectory(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down Expand Up @@ -120,7 +120,7 @@ func TestAccFsRmForNonEmptyDirectoryWithRecursiveFlag(t *testing.T) {
w, err := databricks.NewWorkspaceClient()
require.NoError(t, err)

tmpDir := temporaryDbfsDir(t, w)
tmpDir := TemporaryDbfsDir(t, w)

f, err := filer.NewDbfsClient(w, tmpDir)
require.NoError(t, err)
Expand Down
Loading

0 comments on commit 79e271f

Please sign in to comment.