added integration test

This commit is contained in:
Shreyas Goenka 2024-12-31 12:31:37 +05:30
parent 0ce50fadf0
commit 63e599ccb2
No known key found for this signature in database
GPG Key ID: 92A07DF49CCB0622
2 changed files with 55 additions and 2 deletions

View File

@ -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")
}

View File

@ -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())
}