diff --git a/browser/mapping.go b/browser/mapping.go index c39e38a74..4acbd702e 100644 --- a/browser/mapping.go +++ b/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, diff --git a/browser/module.go b/browser/module.go index ca20cf521..f5afc27c3 100644 --- a/browser/module.go +++ b/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(), diff --git a/browser/modulevu.go b/browser/modulevu.go index 1edabf2be..443e2aefc 100644 --- a/browser/modulevu.go +++ b/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. diff --git a/common/element_handle.go b/common/element_handle.go index 8e51a492f..c769f4627 100644 --- a/common/element_handle.go +++ b/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) diff --git a/common/page.go b/common/page.go index d4fe32bb0..2306a0dbd 100644 --- a/common/page.go +++ b/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) diff --git a/common/screenshotter.go b/common/screenshotter.go index 0afd0ee44..4bb9dd8d1 100644 --- a/common/screenshotter.go +++ b/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) } } diff --git a/storage/file_persister.go b/storage/file_persister.go new file mode 100644 index 000000000..2a1e2921d --- /dev/null +++ b/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 +} diff --git a/storage/file_persister_test.go b/storage/file_persister_test.go new file mode 100644 index 000000000..31f340091 --- /dev/null +++ b/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)) + }) + } +} diff --git a/tests/element_handle_test.go b/tests/element_handle_test.go index d48548615..182f72226 100644 --- a/tests/element_handle_test.go +++ b/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) diff --git a/tests/page_test.go b/tests/page_test.go index b98614f98..dbd304818 100644 --- a/tests/page_test.go +++ b/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)