Add support for image format during screenshot

go.mod

@@ -8,4 +8,5 @@ require (
  github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80
  github.com/mailru/easyjson v0.7.7
  github.com/orisano/pixelmatch v0.0.0-20220722002657-fb0b55479cde
+ golang.org/x/image v0.14.0
 )

screenshot.go

@@ -33,9 +33,21 @@ func Screenshot(sel interface{}, picbuf *[]byte, opts ...QueryOption) QueryActio
  return ScreenshotScale(sel, 1, picbuf, opts...)
 }
 
+// ScreenshotWithFormat is a chromedp element query action that captures a screenshot
+// in the specified image format of the first element node matching the selector.
+func ScreenshotWithFormat(sel interface{}, picbuf *[]byte, imgFormat page.CaptureScreenshotFormat, opts ...QueryOption) QueryAction {
+ return ScreenshotScaleWithFormat(sel, imgFormat, 1, picbuf, opts...)
+}
+
 // ScreenshotScale is like [Screenshot] but accepts a scale parameter that
 // specifies the page scale factor.
 func ScreenshotScale(sel interface{}, scale float64, picbuf *[]byte, opts ...QueryOption) QueryAction {
+ return ScreenshotScaleWithFormat(sel, page.CaptureScreenshotFormatPng, scale, picbuf, opts...)
+}
+
+// ScreenshotScaleWithFormat is like [Screenshot] but accepts the image format
+// and scale parameter that specifies the page scale factor.
+func ScreenshotScaleWithFormat(sel interface{}, imageFormat page.CaptureScreenshotFormat, scale float64, picbuf *[]byte, opts ...QueryOption) QueryAction {
  if picbuf == nil {
  panic("picbuf cannot be nil")
  }
@@ -62,7 +74,7 @@ func ScreenshotScale(sel interface{}, scale float64, picbuf *[]byte, opts ...Que
 
  // take screenshot of the box
  buf, err := page.CaptureScreenshot().
- WithFormat(page.CaptureScreenshotFormatPng).
+ WithFormat(imageFormat).
  WithCaptureBeyondViewport(true).
  WithFromSurface(true).
  WithClip(&clip).
@@ -101,6 +113,23 @@ func CaptureScreenshot(res *[]byte) Action {
  })
 }
 
+// CaptureScreenshotWithFormat is an action that captures/takes a screenshot of the
+// current browser viewport for the provided image format.
+func CaptureScreenshotWithFormat(res *[]byte, format page.CaptureScreenshotFormat) Action {
+ if res == nil {
+ panic("res cannot be nil")
+ }
+
+ return ActionFunc(func(ctx context.Context) error {
+ var err error
+ *res, err = page.CaptureScreenshot().
+ WithFormat(format).
+ WithFromSurface(true).
+ Do(ctx)
+ return err
+ })
+}
+
 // FullScreenshot takes a full screenshot with the specified image quality of
 // the entire browser viewport.
 //

screenshot_test.go

@@ -10,6 +10,9 @@ import (
  "path"
  "testing"
 
+ _ "golang.org/x/image/webp"
+
+ "github.com/chromedp/cdproto/page"
  "github.com/orisano/pixelmatch"
 )
 
@@ -79,6 +82,140 @@ func TestScreenshot(t *testing.T) {
  }
 }
 
+func TestScreenshotWithFormatWebp(t *testing.T) {
+ t.Parallel()
+
+ tests := []struct {
+ name string
+ sel string
+ want string
+ }{
+ {
+ name: "padding border",
+ sel: "#padding-border",
+ want: "element-padding-border.webp",
+ },
+ {
+ name: "larger than viewport",
+ sel: "#larger-than-viewport",
+ want: "element-larger-than-viewport.webp",
+ },
+ {
+ name: "outside viewport",
+ sel: "#outside-viewport",
+ want: "element-scrolled-into-view.webp",
+ },
+ {
+ name: "rotate element",
+ sel: "#rotated",
+ want: "element-rotate.webp",
+ },
+ {
+ name: "fractional dimensions",
+ sel: "#fractional-dimensions",
+ want: "element-fractional.webp",
+ },
+ {
+ name: "fractional offset",
+ sel: "#fractional-offset",
+ want: "element-fractional-offset.webp",
+ },
+ }
+
+ for _, test := range tests {
+ test := test
+ t.Run(test.name, func(t *testing.T) {
+ t.Parallel()
+ ctx, cancel := testAllocate(t, "screenshot.html")
+ defer cancel()
+
+ var buf []byte
+ if err := Run(ctx,
+ EmulateViewport(500, 500),
+ EvaluateAsDevTools("document.documentElement.scrollTo(20, 30)", nil),
+ ScreenshotWithFormat(test.sel, &buf, page.CaptureScreenshotFormatWebp, ByQuery),
+ ); err != nil {
+ t.Fatal(err)
+ }
+
+ diff, err := matchPixel(buf, test.want)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if diff != 0 {
+ t.Fatalf("screenshot does not match. diff: %v", diff)
+ }
+ })
+ }
+}
+
+func TestScreenshotWithFormatJpeg(t *testing.T) {
+ t.Parallel()
+
+ tests := []struct {
+ name string
+ sel string
+ want string
+ }{
+ {
+ name: "padding border",
+ sel: "#padding-border",
+ want: "element-padding-border.jpg",
+ },
+ {
+ name: "larger than viewport",
+ sel: "#larger-than-viewport",
+ want: "element-larger-than-viewport.jpg",
+ },
+ {
+ name: "outside viewport",
+ sel: "#outside-viewport",
+ want: "element-scrolled-into-view.jpg",
+ },
+ {
+ name: "rotate element",
+ sel: "#rotated",
+ want: "element-rotate.jpg",
+ },
+ {
+ name: "fractional dimensions",
+ sel: "#fractional-dimensions",
+ want: "element-fractional.jpg",
+ },
+ {
+ name: "fractional offset",
+ sel: "#fractional-offset",
+ want: "element-fractional-offset.jpg",
+ },
+ }
+
+ for _, test := range tests {
+ test := test
+ t.Run(test.name, func(t *testing.T) {
+ t.Parallel()
+ ctx, cancel := testAllocate(t, "screenshot.html")
+ defer cancel()
+
+ var buf []byte
+ if err := Run(ctx,
+ EmulateViewport(500, 500),
+ EvaluateAsDevTools("document.documentElement.scrollTo(20, 30)", nil),
+ ScreenshotWithFormat(test.sel, &buf, page.CaptureScreenshotFormatJpeg, ByQuery),
+ ); err != nil {
+ t.Fatal(err)
+ }
+
+ diff, err := matchPixel(buf, test.want)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if diff != 0 {
+ t.Fatalf("screenshot does not match. diff: %v", diff)
+ }
+ })
+ }
+}
+
 func TestScreenshotScale(t *testing.T) {
  t.Parallel()
 
@@ -178,6 +315,52 @@ func TestCaptureScreenshot(t *testing.T) {
  }
 }
 
+func TestCaptureScreenshotWithFormatJpeg(t *testing.T) {
+ t.Parallel()
+
+ ctx, cancel := testAllocate(t, "grid.html")
+ defer cancel()
+
+ var buf []byte
+ if err := Run(ctx,
+ EmulateViewport(500, 500),
+ CaptureScreenshotWithFormat(&buf, page.CaptureScreenshotFormatJpeg),
+ ); err != nil {
+ t.Fatal(err)
+ }
+
+ diff, err := matchPixel(buf, "sanity.jpeg")
+ if err != nil {
+ t.Fatal(err)
+ }
+ if diff != 0 {
+ t.Fatalf("screenshot does not match. diff: %v", diff)
+ }
+}
+
+func TestCaptureScreenshotWithFormatWebp(t *testing.T) {
+ t.Parallel()
+
+ ctx, cancel := testAllocate(t, "grid.html")
+ defer cancel()
+
+ var buf []byte
+ if err := Run(ctx,
+ EmulateViewport(500, 500),
+ CaptureScreenshotWithFormat(&buf, page.CaptureScreenshotFormatWebp),
+ ); err != nil {
+ t.Fatal(err)
+ }
+
+ diff, err := matchPixel(buf, "sanity.webp")
+ if err != nil {
+ t.Fatal(err)
+ }
+ if diff != 0 {
+ t.Fatalf("screenshot does not match. diff: %v", diff)
+ }
+}
+
 func TestFullScreenshot(t *testing.T) {
  t.Parallel()