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

add unit tests

This commit is contained in:
Fabian Jakobs 2024-09-04 11:16:50 +02:00
parent 62b2451c10
commit 08b6b10ff0
No known key found for this signature in database
4 changed files with 61 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
libs/template/file.go
View File

@ -23,7 +23,7 @@ type file interface {
DstPath() *destinationPath
// Write file to disk at the destination path.
PersistToDisk(ctx context.Context) error
PersistToDisk() error
}
type destinationPath struct {
@ -62,7 +62,7 @@ func (f *copyFile) DstPath() *destinationPath {
return f.dstPath
}
func (f *copyFile) PersistToDisk(ctx context.Context) error {
func (f *copyFile) PersistToDisk() error {
path := f.DstPath().absPath()
err := os.MkdirAll(filepath.Dir(path), 0755)
if err != nil {
@ -78,10 +78,12 @@ func (f *copyFile) PersistToDisk(ctx context.Context) error {
if err != nil {
return err
}
return writeFile(ctx, path, content)
return writeFile(f.ctx, path, content, f.perm)
}
type inMemoryFile struct {
ctx context.Context
dstPath *destinationPath
content []byte
@ -94,7 +96,7 @@ func (f *inMemoryFile) DstPath() *destinationPath {
return f.dstPath
}
func (f *inMemoryFile) PersistToDisk(ctx context.Context) error {
func (f *inMemoryFile) PersistToDisk() error {
path := f.DstPath().absPath()
err := os.MkdirAll(filepath.Dir(path), 0755)
@ -102,7 +104,7 @@ func (f *inMemoryFile) PersistToDisk(ctx context.Context) error {
return err
}
return writeFile(ctx, path, f.content)
return writeFile(f.ctx, path, f.content, f.perm)
}
func runsOnDatabricks(ctx context.Context) bool {
@ -110,11 +112,15 @@ func runsOnDatabricks(ctx context.Context) bool {
return ok
}
func writeFile(ctx context.Context, path string, content []byte) error {
if strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") {
func shouldUseImportNotebook(ctx context.Context, path string) bool {
return strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb")
}
func writeFile(ctx context.Context, path string, content []byte, perm fs.FileMode) error {
if shouldUseImportNotebook(ctx, path) {
return importNotebook(ctx, path, content)
} else {
return os.WriteFile(path, content, 0644)
return os.WriteFile(path, content, perm)
}
}

45
libs/template/file_test.go
View File

@ -8,6 +8,11 @@ import (
"runtime"
"testing"
"github.com/databricks/databricks-sdk-go/experimental/mocks"
"github.com/databricks/databricks-sdk-go/service/workspace"
"github.com/stretchr/testify/mock"
"github.com/databricks/cli/cmd/root"
"github.com/databricks/cli/libs/filer"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@ -17,6 +22,7 @@ func testInMemoryFile(t *testing.T, perm fs.FileMode) {
tmpDir := t.TempDir()
f := &inMemoryFile{
ctx: context.Background(),
dstPath: &destinationPath{
root: tmpDir,
relPath: "a/b/c",
@ -24,7 +30,7 @@ func testInMemoryFile(t *testing.T, perm fs.FileMode) {
perm: perm,
content: []byte("123"),
}
err := f.PersistToDisk(context.Background())
err := f.PersistToDisk()
assert.NoError(t, err)
assertFileContent(t, filepath.Join(tmpDir, "a/b/c"), "123")
@ -38,10 +44,9 @@ func testCopyFile(t *testing.T, perm fs.FileMode) {
require.NoError(t, err)
err = os.WriteFile(filepath.Join(tmpDir, "source"), []byte("qwerty"), perm)
require.NoError(t, err)
ctx := context.Background()
f := &copyFile{
ctx: ctx,
ctx: context.Background(),
dstPath: &destinationPath{
root: tmpDir,
relPath: "a/b/c",
@ -50,7 +55,7 @@ func testCopyFile(t *testing.T, perm fs.FileMode) {
srcPath: "source",
srcFiler: templateFiler,
}
err = f.PersistToDisk(ctx)
err = f.PersistToDisk()
assert.NoError(t, err)
assertFileContent(t, filepath.Join(tmpDir, "a/b/c"), "qwerty")
@ -110,3 +115,35 @@ func TestTemplateCopyFilePersistToDiskForWindows(t *testing.T) {
// fs.FileMode values we can use for different operating systems.
testCopyFile(t, 0666)
}
func TestShouldUseImportNotebook(t *testing.T) {
ctx := context.Background()
assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar"))
assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar.ipynb"))
assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar"))
assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar.ipynb"))
t.Setenv("DATABRICKS_RUNTIME_VERSION", "14.3")
assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar"))
assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar.ipynb"))
assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar"))
assert.True(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar.ipynb"))
}
func TestImportNotebook(t *testing.T) {
ctx := context.Background()
m := mocks.NewMockWorkspaceClient(t)
ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
workspaceApi := m.GetMockWorkspaceAPI()
workspaceApi.EXPECT().Import(mock.Anything, workspace.Import{
Content: "cXdlcnR5", // base64 of "qwerty"
Format: "AUTO",
Overwrite: false,
Path: "/Workspace/foo/bar.ipynb",
}).Return(nil)
err := importNotebook(ctx, "/Workspace/foo/bar.ipynb", []byte("qwerty"))
assert.NoError(t, err)
}

3
libs/template/renderer.go
View File

@ -194,6 +194,7 @@ func (r *renderer) computeFile(relPathTemplate string) (file, error) {
}
return &inMemoryFile{
ctx: r.ctx,
dstPath: &destinationPath{
root: r.instanceRoot,
relPath: relPath,
@ -320,7 +321,7 @@ func (r *renderer) persistToDisk() error {
// Persist files to disk
for _, file := range filesToPersist {
err := file.PersistToDisk(r.ctx)
err := file.PersistToDisk()
if err != nil {
return err
}

4
libs/template/renderer_test.go
View File

@ -329,6 +329,7 @@ func TestRendererPersistToDisk(t *testing.T) {
skipPatterns: []string{"a/b/c", "mn*"},
files: []file{
&inMemoryFile{
ctx: ctx,
dstPath: &destinationPath{
root: tmpDir,
relPath: "a/b/c",
@ -337,6 +338,7 @@ func TestRendererPersistToDisk(t *testing.T) {
content: nil,
},
&inMemoryFile{
ctx: ctx,
dstPath: &destinationPath{
root: tmpDir,
relPath: "mno",
@ -345,6 +347,7 @@ func TestRendererPersistToDisk(t *testing.T) {
content: nil,
},
&inMemoryFile{
ctx: ctx,
dstPath: &destinationPath{
root: tmpDir,
relPath: "a/b/d",
@ -353,6 +356,7 @@ func TestRendererPersistToDisk(t *testing.T) {
content: []byte("123"),
},
&inMemoryFile{
ctx: ctx,
dstPath: &destinationPath{
root: tmpDir,
relPath: "mmnn",