Add local file persister

browser/modulevu.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/grafana/xk6-browser/common"
 	"github.com/grafana/xk6-browser/k6ext"
+	"github.com/grafana/xk6-browser/storage"
 
 	k6modules "go.k6.io/k6/js/modules"
 )
@@ -20,6 +21,8 @@ type moduleVU struct {
 	*browserRegistry
 
 	*taskQueueRegistry
+
+	FilePersister storage.FilePersister
 }
 
 // browser returns the VU browser instance for the current iteration.

storage/file_persister.go

@@ -0,0 +1,45 @@
+package storage
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+)
+
+// FilePersister will persist files. It abstracts away the where and how of
+// writing files to the source destination.
+type FilePersister interface {
+	Persist(ctx context.Context, path string, data io.Reader) error
+}
+
+// LocalFilePersister will persist files to the local disk.
+type LocalFilePersister struct{}
+
+// Persist will write the contents of data to the local disk on the specified path.
+// TODO: we should not write to disk here but put it on some queue for async disk writes.
+func (l *LocalFilePersister) Persist(_ context.Context, path string, data io.Reader) (err error) {
+	cp := filepath.Clean(path)
+
+	dir := filepath.Dir(cp)
+	if err = os.MkdirAll(dir, 0o755); err != nil {
+		return fmt.Errorf("creating a local directory %q: %w", dir, err)
+	}
+
+	f, err := os.OpenFile(cp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600)
+	if err != nil {
+		return fmt.Errorf("creating a local file %q: %w", cp, err)
+	}
+	defer func() {
+		tempErr := f.Close()
+		// Only return the close error if there isn't already an existing error.
+		if tempErr != nil && err == nil {
+			err = fmt.Errorf("closing the local file %q: %w", cp, tempErr)
+		}
+	}()
+
+	_, err = io.Copy(f, data)
+
+	return
+}

storage/file_persister_test.go

@@ -0,0 +1,87 @@
+package storage
+
+import (
+	"context"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestLocalFilePersister(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name         string
+		path         string
+		existingData string
+		data         string
+		truncates    bool
+	}{
+		{
+			name: "just_file",
+			path: "test.txt",
+			data: "some data",
+		},
+		{
+			name: "with_dir",
+			path: "path/test.txt",
+			data: "some data",
+		},
+		{
+			name:         "truncates",
+			path:         "test.txt",
+			data:         "some data",
+			truncates:    true,
+			existingData: "existing data",
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			dir, err := os.MkdirTemp("", "*")
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = os.RemoveAll(dir) })
+			p := filepath.Join(dir, tt.path)
+
+			// We want to make sure that the persister truncates the existing
+			// data and therefore overwrites existing data. This sets up a file
+			// with some existing data that should be overwritten.
+			if tt.truncates {
+				err = os.WriteFile(p, []byte(tt.existingData), 0o600)
+				require.NoError(t, err)
+			}
+
+			l := &LocalFilePersister{}
+			err = l.Persist(context.Background(), p, strings.NewReader(tt.data))
+			assert.NoError(t, err)
+
+			i, err := os.Stat(p)
+			require.NoError(t, err)
+			assert.False(t, i.IsDir())
+
+			f, err := os.Open(filepath.Clean(p))
+			require.NoError(t, err)
+			defer func() {
+				err = f.Close()
+				require.NoError(t, err)
+			}()
+
+			bb, err := io.ReadAll(f)
+			require.NoError(t, err)
+
+			if tt.truncates {
+				assert.NotEqual(t, tt.existingData, string(bb))
+			}
+
+			assert.Equal(t, tt.data, string(bb))
+		})
+	}
+}