Add a localPersister

browser/mapping.go

@@ -279,8 +279,10 @@ func mapElementHandle(vu moduleVU, eh *common.ElementHandle) mapping {
 			}
 			return mapFrame(vu, f), nil
 		},
-		"press":                  eh.Press,
-		"screenshot":             eh.Screenshot,
+		"press": eh.Press,
+		"screenshot": func(opts goja.Value) goja.ArrayBuffer {
+			return eh.Screenshot(opts, vu.LocalFilePersister)
+		},
 		"scrollIntoViewIfNeeded": eh.ScrollIntoViewIfNeeded,
 		"selectOption":           eh.SelectOption,
 		"selectText":             eh.SelectText,
@@ -686,8 +688,10 @@ func mapPage(vu moduleVU, p *common.Page) mapping {
 			r := mapResponse(vu, p.Reload(opts))
 			return rt.ToValue(r).ToObject(rt)
 		},
-		"route":                       p.Route,
-		"screenshot":                  p.Screenshot,
+		"route": p.Route,
+		"screenshot": func(opts goja.Value) goja.ArrayBuffer {
+			return p.Screenshot(opts, vu.LocalFilePersister)
+		},
 		"selectOption":                p.SelectOption,
 		"setContent":                  p.SetContent,
 		"setDefaultNavigationTimeout": p.SetDefaultNavigationTimeout,

browser/module.go

@@ -13,6 +13,7 @@ import (
 	"github.com/grafana/xk6-browser/common"
 	"github.com/grafana/xk6-browser/env"
 	"github.com/grafana/xk6-browser/k6ext"
+	"github.com/grafana/xk6-browser/storage"
 
 	k6modules "go.k6.io/k6/js/modules"
 )
@@ -76,7 +77,8 @@ func (m *RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
 					m.PidRegistry,
 					m.tracesMetadata,
 				),
-				taskQueueRegistry: newTaskQueueRegistry(vu),
+				taskQueueRegistry:  newTaskQueueRegistry(vu),
+				LocalFilePersister: &storage.LocalFilePersister{},
 			}),
 			Devices:         common.GetDevices(),
 			NetworkProfiles: common.GetNetworkProfiles(),

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
+
+	*storage.LocalFilePersister
 }
 
 // browser returns the VU browser instance for the current iteration.

common/element_handle.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/grafana/xk6-browser/common/js"
 	"github.com/grafana/xk6-browser/k6ext"
+	"github.com/grafana/xk6-browser/storage"
 )
 
 const resultDone = "done"
@@ -1176,7 +1177,8 @@ func (h *ElementHandle) setChecked(apiCtx context.Context, checked bool, p *Posi
 	return nil
 }
 
-func (h *ElementHandle) Screenshot(opts goja.Value) goja.ArrayBuffer {
+// Screenshot will instruct Chrome to save a screenshot of the current element and save it to specified file.
+func (h *ElementHandle) Screenshot(opts goja.Value, fp *storage.LocalFilePersister) goja.ArrayBuffer {
 	spanCtx, span := TraceAPICall(
 		h.ctx,
 		h.frame.page.targetID.String(),
@@ -1191,7 +1193,7 @@ func (h *ElementHandle) Screenshot(opts goja.Value) goja.ArrayBuffer {
 	}
 	span.SetAttributes(attribute.String("screenshot.path", parsedOpts.Path))
 
-	s := newScreenshotter(spanCtx)
+	s := newScreenshotter(spanCtx, fp)
 	buf, err := s.screenshotElement(h, parsedOpts)
 	if err != nil {
 		k6ext.Panic(h.ctx, "taking screenshot: %w", err)

common/page.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/grafana/xk6-browser/k6ext"
 	"github.com/grafana/xk6-browser/log"
+	"github.com/grafana/xk6-browser/storage"
 
 	k6modules "go.k6.io/k6/js/modules"
 )
@@ -1120,7 +1121,7 @@ func (p *Page) Route(url goja.Value, handler goja.Callable) {
 }
 
 // Screenshot will instruct Chrome to save a screenshot of the current page and save it to specified file.
-func (p *Page) Screenshot(opts goja.Value) goja.ArrayBuffer {
+func (p *Page) Screenshot(opts goja.Value, fp *storage.LocalFilePersister) goja.ArrayBuffer {
 	spanCtx, span := TraceAPICall(p.ctx, p.targetID.String(), "page.screenshot")
 	defer span.End()
 
@@ -1130,7 +1131,7 @@ func (p *Page) Screenshot(opts goja.Value) goja.ArrayBuffer {
 	}
 	span.SetAttributes(attribute.String("screenshot.path", parsedOpts.Path))
 
-	s := newScreenshotter(spanCtx)
+	s := newScreenshotter(spanCtx, fp)
 	buf, err := s.screenshotPage(p, parsedOpts)
 	if err != nil {
 		k6ext.Panic(p.ctx, "capturing screenshot: %w", err)

common/screenshotter.go

@@ -7,14 +7,14 @@ import (
 	"errors"
 	"fmt"
 	"math"
-	"os"
-	"path/filepath"
 	"strings"
 
 	"github.com/chromedp/cdproto/cdp"
 	"github.com/chromedp/cdproto/emulation"
 	cdppage "github.com/chromedp/cdproto/page"
 	"github.com/dop251/goja"
+
+	"github.com/grafana/xk6-browser/storage"
 )
 
 // ImageFormat represents an image file format.
@@ -61,11 +61,12 @@ func (f *ImageFormat) UnmarshalJSON(b []byte) error {
 }
 
 type screenshotter struct {
-	ctx context.Context
+	ctx       context.Context
+	persister *storage.LocalFilePersister
 }
 
-func newScreenshotter(ctx context.Context) *screenshotter {
-	return &screenshotter{ctx}
+func newScreenshotter(ctx context.Context, fp *storage.LocalFilePersister) *screenshotter {
+	return &screenshotter{ctx, fp}
 }
 
 func (s *screenshotter) fullPageSize(p *Page) (*Size, error) {
@@ -215,14 +216,9 @@ func (s *screenshotter) screenshot(
 	}
 
 	// Save screenshot capture to file
-	// TODO: we should not write to disk here but put it on some queue for async disk writes
 	if path != "" {
-		dir := filepath.Dir(path)
-		if err := os.MkdirAll(dir, 0o755); err != nil {
-			return nil, fmt.Errorf("creating screenshot directory %q: %w", dir, err)
-		}
-		if err := os.WriteFile(path, buf, 0o644); err != nil {
-			return nil, fmt.Errorf("saving screenshot to %q: %w", path, err)
+		if err := s.persister.Persist(path, bytes.NewBuffer(buf)); err != nil {
+			return nil, fmt.Errorf("persisting screenshot: %w", err)
 		}
 	}
 

storage/file_persister.go

@@ -0,0 +1,45 @@
+package storage
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+)
+
+// 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(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() {
+		if cerr := f.Close(); cerr != nil && err == nil {
+			err = fmt.Errorf("closing the local file %q: %w", cp, cerr)
+		}
+	}()
+
+	bf := bufio.NewWriter(f)
+
+	if _, err := io.Copy(bf, data); err != nil {
+		return fmt.Errorf("copying data to file: %w", err)
+	}
+
+	if err := bf.Flush(); err != nil {
+		return fmt.Errorf("flushing data to disk: %w", err)
+	}
+
+	return nil
+}

storage/file_persister_test.go

@@ -0,0 +1,76 @@
+package storage
+
+import (
+	"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 := t.TempDir()
+			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)
+			}
+
+			var l LocalFilePersister
+			err := l.Persist(p, strings.NewReader(tt.data))
+			assert.NoError(t, err)
+
+			i, err := os.Stat(p)
+			require.NoError(t, err)
+			assert.False(t, i.IsDir())
+
+			bb, err := os.ReadFile(filepath.Clean(p))
+			require.NoError(t, err)
+
+			if tt.truncates {
+				assert.NotEqual(t, tt.existingData, string(bb))
+			}
+
+			assert.Equal(t, tt.data, string(bb))
+		})
+	}
+}

tests/element_handle_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"github.com/grafana/xk6-browser/common"
+	"github.com/grafana/xk6-browser/storage"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -374,7 +375,7 @@ func TestElementHandleScreenshot(t *testing.T) {
 	elem, err := p.Query("div")
 	require.NoError(t, err)
 
-	buf := elem.Screenshot(nil)
+	buf := elem.Screenshot(nil, &storage.LocalFilePersister{})
 
 	reader := bytes.NewReader(buf.Bytes())
 	img, err := png.Decode(reader)

tests/page_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/grafana/xk6-browser/browser"
 	"github.com/grafana/xk6-browser/common"
 	"github.com/grafana/xk6-browser/k6ext/k6test"
+	"github.com/grafana/xk6-browser/storage"
 )
 
 type emulateMediaOpts struct {
@@ -477,7 +478,7 @@ func TestPageScreenshotFullpage(t *testing.T) {
 		FullPage bool `js:"fullPage"`
 	}{
 		FullPage: true,
-	}))
+	}), &storage.LocalFilePersister{})
 
 	reader := bytes.NewReader(buf.Bytes())
 	img, err := png.Decode(reader)