mirror of https://github.com/databricks/cli.git
added integration test
This commit is contained in:
parent
0ce50fadf0
commit
63e599ccb2
|
@ -7,7 +7,9 @@ import (
|
|||
"errors"
|
||||
"io"
|
||||
"io/fs"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
@ -895,3 +897,27 @@ func TestWorkspaceFilesExtensions_ExportFormatIsPreserved(t *testing.T) {
|
|||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDbfsFilerForStreamingUploads(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
f, _ := setupDbfsFiler(t)
|
||||
|
||||
// Set MaxDbfsPutFileSize to 1 to force streaming uploads
|
||||
prevV := filer.MaxDbfsPutFileSize
|
||||
filer.MaxDbfsPutFileSize = 1
|
||||
t.Cleanup(func() {
|
||||
filer.MaxDbfsPutFileSize = prevV
|
||||
})
|
||||
|
||||
// Write a file to local disk.
|
||||
tmpDir := t.TempDir()
|
||||
err := os.WriteFile(filepath.Join(tmpDir, "foo.txt"), []byte("foobar"), 0o644)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Write a file with streaming upload
|
||||
err = f.Write(ctx, "foo.txt", strings.NewReader("foo"))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Assert contents
|
||||
filerTest{t, f}.assertContents(ctx, "foo.txt", "foo")
|
||||
}
|
||||
|
|
|
@ -67,7 +67,6 @@ type mockDbfsHandle struct {
|
|||
builder strings.Builder
|
||||
}
|
||||
|
||||
func (h *mockDbfsHandle) Write(data []byte) (n int, err error) { return 0, nil }
|
||||
func (h *mockDbfsHandle) Read(data []byte) (n int, err error) { return 0, nil }
|
||||
func (h *mockDbfsHandle) Close() error { return nil }
|
||||
func (h *mockDbfsHandle) WriteTo(w io.Writer) (n int64, err error) { return 0, nil }
|
||||
|
@ -81,6 +80,10 @@ func (h *mockDbfsHandle) ReadFrom(r io.Reader) (n int64, err error) {
|
|||
return int64(num), err
|
||||
}
|
||||
|
||||
func (h *mockDbfsHandle) Write(data []byte) (n int, err error) {
|
||||
return h.builder.Write(data)
|
||||
}
|
||||
|
||||
func TestDbfsClientForLargerFiles(t *testing.T) {
|
||||
// write file to local disk
|
||||
tmp := t.TempDir()
|
||||
|
@ -106,7 +109,6 @@ func TestDbfsClientForLargerFiles(t *testing.T) {
|
|||
}
|
||||
|
||||
h := &mockDbfsHandle{}
|
||||
|
||||
m.GetMockDbfsAPI().EXPECT().GetStatusByPath(mock.Anything, "dbfs:/a/b/c").Return(nil, nil)
|
||||
m.GetMockDbfsAPI().EXPECT().Open(mock.Anything, "dbfs:/a/b/c/hello.txt", files.FileModeWrite).Return(h, nil)
|
||||
|
||||
|
@ -122,3 +124,28 @@ func TestDbfsClientForLargerFiles(t *testing.T) {
|
|||
// verify the file content was written to the mock handle
|
||||
assert.Equal(t, "hello world", h.builder.String())
|
||||
}
|
||||
|
||||
func TestDbfsClientForNonLocalFiles(t *testing.T) {
|
||||
// setup DBFS client with mocks
|
||||
m := mocks.NewMockWorkspaceClient(t)
|
||||
mockApiClient := &mockDbfsApiClient{t: t}
|
||||
dbfsClient := DbfsClient{
|
||||
apiClient: mockApiClient,
|
||||
workspaceClient: m.WorkspaceClient,
|
||||
root: NewWorkspaceRootPath("dbfs:/a/b/c"),
|
||||
}
|
||||
|
||||
h := &mockDbfsHandle{}
|
||||
m.GetMockDbfsAPI().EXPECT().GetStatusByPath(mock.Anything, "dbfs:/a/b/c").Return(nil, nil)
|
||||
m.GetMockDbfsAPI().EXPECT().Open(mock.Anything, "dbfs:/a/b/c/hello.txt", files.FileModeWrite).Return(h, nil)
|
||||
|
||||
// write file to DBFS
|
||||
err := dbfsClient.Write(context.Background(), "hello.txt", strings.NewReader("hello world"))
|
||||
require.NoError(t, err)
|
||||
|
||||
// verify mock API client is NOT called
|
||||
require.False(t, mockApiClient.isCalled)
|
||||
|
||||
// verify the file content was written to the mock handle
|
||||
assert.Equal(t, "hello world", h.builder.String())
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue