add unit tests

databricks · Sep 4, 2024 · 20ee23e · 20ee23e
1 parent 13996bd
commit 20ee23e
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 13 deletions.
diff --git a/libs/template/file.go b/libs/template/file.go
@@ -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,27 +96,31 @@ 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)
 	if err != nil {
 		return err
 	}
 
-	return writeFile(ctx, path, f.content)
+	return writeFile(f.ctx, path, f.content, f.perm)
 }
 
 func runsOnDatabricks(ctx context.Context) bool {
 	_, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION")
 	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)
 	}
 }
 

diff --git a/libs/template/file_test.go b/libs/template/file_test.go
@@ -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,14 +22,15 @@ func testInMemoryFile(t *testing.T, perm fs.FileMode) {
 	tmpDir := t.TempDir()
 
 	f := &inMemoryFile{
+		ctx: context.Background(),
 		dstPath: &destinationPath{
 			root:    tmpDir,
 			relPath: "a/b/c",
 		},
 		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)
+}
diff --git a/libs/template/renderer.go b/libs/template/renderer.go
@@ -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
 		}

diff --git a/libs/template/renderer_test.go b/libs/template/renderer_test.go
@@ -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",