From 0d789120b7a8f529b56271612d043a46135712cc Mon Sep 17 00:00:00 2001 From: ankur22 Date: Mon, 29 Jan 2024 10:28:39 +0000 Subject: [PATCH] Refactor to work with the filePersister We're refactoring out the existing use of working directly with the os package to persist screenshots to the local disk from screenshotter. This change enables to work with the localPersister whereby the underlying details of the where and how are hidden allowing us to easily use either the localPersister (current) or the remotePersister (future). --- browser/mapping.go | 12 ++++++++---- common/element_handle.go | 6 ++++-- common/page.go | 5 +++-- common/screenshotter.go | 19 ++++++++----------- tests/element_handle_test.go | 3 ++- tests/page_test.go | 3 ++- 6 files changed, 27 insertions(+), 21 deletions(-) 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/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..d63d71eeb 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) { @@ -217,12 +218,8 @@ 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/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)