Skip to content

Commit 20ee23e

Browse files
fjakobsfjakobs
committed
add unit tests
1 parent 13996bd commit 20ee23e

File tree

4 files changed

+61
-13
lines changed

4 files changed

+61
-13
lines changed

‎libs/template/file.go‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ type file interface {
2323
DstPath() *destinationPath
2424

2525
// Write file to disk at the destination path.
26-
PersistToDisk(ctx context.Context) error
26+
PersistToDisk() error
2727
}
2828

2929
type destinationPath struct {
@@ -62,7 +62,7 @@ func (f *copyFile) DstPath() *destinationPath {
6262
return f.dstPath
6363
}
6464

65-
func (f *copyFile) PersistToDisk(ctx context.Context) error {
65+
func (f *copyFile) PersistToDisk() error {
6666
path := f.DstPath().absPath()
6767
err := os.MkdirAll(filepath.Dir(path), 0755)
6868
if err != nil {
@@ -78,10 +78,12 @@ func (f *copyFile) PersistToDisk(ctx context.Context) error {
7878
if err != nil {
7979
return err
8080
}
81-
return writeFile(ctx, path, content)
81+
return writeFile(f.ctx, path, content, f.perm)
8282
}
8383

8484
type inMemoryFile struct {
85+
ctx context.Context
86+
8587
dstPath *destinationPath
8688

8789
content []byte
@@ -94,27 +96,31 @@ func (f *inMemoryFile) DstPath() *destinationPath {
9496
return f.dstPath
9597
}
9698

97-
func (f *inMemoryFile) PersistToDisk(ctx context.Context) error {
99+
func (f *inMemoryFile) PersistToDisk() error {
98100
path := f.DstPath().absPath()
99101

100102
err := os.MkdirAll(filepath.Dir(path), 0755)
101103
if err != nil {
102104
return err
103105
}
104106

105-
return writeFile(ctx, path, f.content)
107+
return writeFile(f.ctx, path, f.content, f.perm)
106108
}
107109

108110
func runsOnDatabricks(ctx context.Context) bool {
109111
_, ok := env.Lookup(ctx, "DATABRICKS_RUNTIME_VERSION")
110112
return ok
111113
}
112114

113-
func writeFile(ctx context.Context, path string, content []byte) error {
114-
if strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb") {
115+
func shouldUseImportNotebook(ctx context.Context, path string) bool {
116+
return strings.HasPrefix(path, "/Workspace/") && runsOnDatabricks(ctx) && strings.HasSuffix(path, ".ipynb")
117+
}
118+
119+
func writeFile(ctx context.Context, path string, content []byte, perm fs.FileMode) error {
120+
if shouldUseImportNotebook(ctx, path) {
115121
return importNotebook(ctx, path, content)
116122
} else {
117-
return os.WriteFile(path, content, 0644)
123+
return os.WriteFile(path, content, perm)
118124
}
119125
}
120126

‎libs/template/file_test.go‎

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,11 @@ import (
88
"runtime"
99
"testing"
1010

11+
"github.com/databricks/databricks-sdk-go/experimental/mocks"
12+
"github.com/databricks/databricks-sdk-go/service/workspace"
13+
"github.com/stretchr/testify/mock"
14+
15+
"github.com/databricks/cli/cmd/root"
1116
"github.com/databricks/cli/libs/filer"
1217
"github.com/stretchr/testify/assert"
1318
"github.com/stretchr/testify/require"
@@ -17,14 +22,15 @@ func testInMemoryFile(t *testing.T, perm fs.FileMode) {
1722
tmpDir := t.TempDir()
1823

1924
f := &inMemoryFile{
25+
ctx: context.Background(),
2026
dstPath: &destinationPath{
2127
root: tmpDir,
2228
relPath: "a/b/c",
2329
},
2430
perm: perm,
2531
content: []byte("123"),
2632
}
27-
err := f.PersistToDisk(context.Background())
33+
err := f.PersistToDisk()
2834
assert.NoError(t, err)
2935

3036
assertFileContent(t, filepath.Join(tmpDir, "a/b/c"), "123")
@@ -38,10 +44,9 @@ func testCopyFile(t *testing.T, perm fs.FileMode) {
3844
require.NoError(t, err)
3945
err = os.WriteFile(filepath.Join(tmpDir, "source"), []byte("qwerty"), perm)
4046
require.NoError(t, err)
41-
ctx := context.Background()
4247

4348
f := &copyFile{
44-
ctx: ctx,
49+
ctx: context.Background(),
4550
dstPath: &destinationPath{
4651
root: tmpDir,
4752
relPath: "a/b/c",
@@ -50,7 +55,7 @@ func testCopyFile(t *testing.T, perm fs.FileMode) {
5055
srcPath: "source",
5156
srcFiler: templateFiler,
5257
}
53-
err = f.PersistToDisk(ctx)
58+
err = f.PersistToDisk()
5459
assert.NoError(t, err)
5560

5661
assertFileContent(t, filepath.Join(tmpDir, "a/b/c"), "qwerty")
@@ -110,3 +115,35 @@ func TestTemplateCopyFilePersistToDiskForWindows(t *testing.T) {
110115
// fs.FileMode values we can use for different operating systems.
111116
testCopyFile(t, 0666)
112117
}
118+
119+
func TestShouldUseImportNotebook(t *testing.T) {
120+
ctx := context.Background()
121+
assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar"))
122+
assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar.ipynb"))
123+
assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar"))
124+
assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar.ipynb"))
125+
126+
t.Setenv("DATABRICKS_RUNTIME_VERSION", "14.3")
127+
assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar"))
128+
assert.False(t, shouldUseImportNotebook(ctx, "./foo/bar.ipynb"))
129+
assert.False(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar"))
130+
assert.True(t, shouldUseImportNotebook(ctx, "/Workspace/foo/bar.ipynb"))
131+
}
132+
133+
func TestImportNotebook(t *testing.T) {
134+
ctx := context.Background()
135+
136+
m := mocks.NewMockWorkspaceClient(t)
137+
ctx = root.SetWorkspaceClient(ctx, m.WorkspaceClient)
138+
139+
workspaceApi := m.GetMockWorkspaceAPI()
140+
workspaceApi.EXPECT().Import(mock.Anything, workspace.Import{
141+
Content: "cXdlcnR5", // base64 of "qwerty"
142+
Format: "AUTO",
143+
Overwrite: false,
144+
Path: "/Workspace/foo/bar.ipynb",
145+
}).Return(nil)
146+
147+
err := importNotebook(ctx, "/Workspace/foo/bar.ipynb", []byte("qwerty"))
148+
assert.NoError(t, err)
149+
}

‎libs/template/renderer.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,7 @@ func (r *renderer) computeFile(relPathTemplate string) (file, error) {
194194
}
195195

196196
return &inMemoryFile{
197+
ctx: r.ctx,
197198
dstPath: &destinationPath{
198199
root: r.instanceRoot,
199200
relPath: relPath,
@@ -320,7 +321,7 @@ func (r *renderer) persistToDisk() error {
320321

321322
// Persist files to disk
322323
for _, file := range filesToPersist {
323-
err := file.PersistToDisk(r.ctx)
324+
err := file.PersistToDisk()
324325
if err != nil {
325326
return err
326327
}

‎libs/template/renderer_test.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,7 @@ func TestRendererPersistToDisk(t *testing.T) {
329329
skipPatterns: []string{"a/b/c", "mn*"},
330330
files: []file{
331331
&inMemoryFile{
332+
ctx: ctx,
332333
dstPath: &destinationPath{
333334
root: tmpDir,
334335
relPath: "a/b/c",
@@ -337,6 +338,7 @@ func TestRendererPersistToDisk(t *testing.T) {
337338
content: nil,
338339
},
339340
&inMemoryFile{
341+
ctx: ctx,
340342
dstPath: &destinationPath{
341343
root: tmpDir,
342344
relPath: "mno",
@@ -345,6 +347,7 @@ func TestRendererPersistToDisk(t *testing.T) {
345347
content: nil,
346348
},
347349
&inMemoryFile{
350+
ctx: ctx,
348351
dstPath: &destinationPath{
349352
root: tmpDir,
350353
relPath: "a/b/d",
@@ -353,6 +356,7 @@ func TestRendererPersistToDisk(t *testing.T) {
353356
content: []byte("123"),
354357
},
355358
&inMemoryFile{
359+
ctx: ctx,
356360
dstPath: &destinationPath{
357361
root: tmpDir,
358362
relPath: "mmnn",

0 commit comments

Comments
 (0)