From 63e599ccb26230dce714f39261ddae2ce0b21dfe Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 31 Dec 2024 12:31:37 +0530
Subject: [PATCH] added integration test

---
 integration/libs/filer/filer_test.go | 26 +++++++++++++++++++++++
 libs/filer/dbfs_client_test.go       | 31 ++++++++++++++++++++++++++--
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/integration/libs/filer/filer_test.go b/integration/libs/filer/filer_test.go
index 766f9817b..6fbd6d6ef 100644
--- a/integration/libs/filer/filer_test.go
+++ b/integration/libs/filer/filer_test.go
@@ -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")
+}
diff --git a/libs/filer/dbfs_client_test.go b/libs/filer/dbfs_client_test.go
index 2e8ad86d0..d08c73cf7 100644
--- a/libs/filer/dbfs_client_test.go
+++ b/libs/filer/dbfs_client_test.go
@@ -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())
+}