Add Screenshot function to Selection.

Fixes sclevine#104. Adds Screenshot function to Selection and GetScreenshot to the element API.

api/element.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"encoding/base64"
 	"errors"
 	"path"
 	"strings"
@@ -145,3 +146,24 @@ func (e *Element) GetLocation() (x, y int, err error) {
 func round(number float64) int {
 	return int(number + 0.5)
 }
+
+func (e *Element) GetRect() (x, y, width, height int, err error) {
+	var rect struct {
+		X      float64 `json:"x"`
+		Y      float64 `json:"y"`
+		Height float64 `json:"height"`
+		Width  float64 `json:"width"`
+	}
+	if err := e.Send("GET", "rect", nil, &rect); err != nil {
+		return 0, 0, 0, 0, err
+	}
+	return round(rect.X), round(rect.Y), round(rect.Width), round(rect.Height), nil
+}
+
+func (e *Element) GetScreenshot() ([]byte, error) {
+	var base64Image string
+	if err := e.Send("GET", "screenshot", nil, &base64Image); err != nil {
+		return nil, err
+	}
+	return base64.StdEncoding.DecodeString(base64Image)
+}

api/element_test.go

@@ -386,4 +386,65 @@ var _ = Describe("Element", func() {
 			})
 		})
 	})
+
+	Describe("#GetRect", func() {
+		It("should successfully send a GET request to the rect endpoint", func() {
+			_, _, _, _, err := element.GetRect()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(bus.SendCall.Method).To(Equal("GET"))
+			Expect(bus.SendCall.Endpoint).To(Equal("element/some-id/rect"))
+		})
+
+		It("should return the rounded rect of the element", func() {
+			bus.SendCall.Result = `{"x": 100.7, "y": 200, "height": 55.05, "width": 33}`
+			x, y, width, height, err := element.GetRect()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(x).To(Equal(101))
+			Expect(y).To(Equal(200))
+			Expect(width).To(Equal(33))
+			Expect(height).To(Equal(55))
+		})
+
+		Context("when the bus indicates a failure", func() {
+			It("should return an error indicating the bus failed to retrieve the rect", func() {
+				bus.SendCall.Err = errors.New("some error")
+				_, _, _, _, err := element.GetRect()
+				Expect(err).To(MatchError("some error"))
+			})
+		})
+	})
+
+	Describe("#GetScreenshot", func() {
+		It("should successfully send a GET request to the screenshot endpoint", func() {
+			_, err := element.GetScreenshot()
+			Expect(err).NotTo(HaveOccurred())
+			Expect(bus.SendCall.Method).To(Equal("GET"))
+			Expect(bus.SendCall.Endpoint).To(Equal("element/some-id/screenshot"))
+		})
+
+		Context("when the image is valid base64", func() {
+			It("should return the decoded image", func() {
+				bus.SendCall.Result = `"c29tZS1wbmc="`
+				image, err := element.GetScreenshot()
+				Expect(err).NotTo(HaveOccurred())
+				Expect(string(image)).To(Equal("some-png"))
+			})
+		})
+
+		Context("when the image is not valid base64", func() {
+			It("should return an error", func() {
+				bus.SendCall.Result = `"..."`
+				_, err := element.GetScreenshot()
+				Expect(err).To(MatchError("illegal base64 data at input byte 0"))
+			})
+		})
+
+		Context("when the bus indicates a failure", func() {
+			It("should return an error", func() {
+				bus.SendCall.Err = errors.New("some error")
+				_, err := element.GetScreenshot()
+				Expect(err).To(MatchError("some error"))
+			})
+		})
+	})
 })

injector_test.go

@@ -2,14 +2,14 @@ package agouti
 
 import "github.com/sclevine/agouti/internal/target"
 
-func NewTestSelection(session apiSession, elements elementRepository, firstSelector string) *Selection {
+func NewTestSelection(session apiSession, elements elementRepository, firstSelector string, cropper Cropper) *Selection {
 	selector := target.Selector{Type: target.CSS, Value: firstSelector, Single: true}
-	return &Selection{selectable{session, target.Selectors{selector}}, elements}
+	return &Selection{selectable{session, target.Selectors{selector}}, elements, cropper}
 }
 
-func NewTestMultiSelection(session apiSession, elements elementRepository, firstSelector string) *MultiSelection {
+func NewTestMultiSelection(session apiSession, elements elementRepository, firstSelector string, cropper Cropper) *MultiSelection {
 	selector := target.Selector{Type: target.CSS, Value: firstSelector}
-	selection := Selection{selectable{session, target.Selectors{selector}}, elements}
+	selection := Selection{selectable{session, target.Selectors{selector}}, elements, cropper}
 	return &MultiSelection{selection}
 }
 

internal/element/repository.go

@@ -34,6 +34,8 @@ type Element interface {
 	Value(text string) error
 	Submit() error
 	GetLocation() (x, y int, err error)
+	GetRect() (x, y, width, height int, err error)
+	GetScreenshot() ([]byte, error)
 }
 
 func (e *Repository) GetAtLeastOne() ([]Element, error) {

internal/integration/selection_test.go

@@ -1,13 +1,17 @@
 package integration_test
 
 import (
+	"image"
+	"image/png"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"os"
+
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 	"github.com/sclevine/agouti"
 	. "github.com/sclevine/agouti/matchers"
-	"io/ioutil"
-	"net/http"
-	"net/http/httptest"
 )
 
 func testSelection(browserName string, newPage pageFunc) {
@@ -187,5 +191,16 @@ func testSelection(browserName string, newPage pageFunc) {
 				Eventually(func() bool { return submitted }).Should(BeTrue())
 			})
 		})
+
+		It("should support taking screenshots", func() {
+			selection := page.Find("a")
+			Expect(selection.Screenshot(".test.screenshot.png")).To(Succeed())
+			defer os.Remove(".test.screenshot.png")
+			file, _ := os.Open(".test.screenshot.png")
+			img, err := png.Decode(file)
+			Expect(err).NotTo(HaveOccurred())
+			// Check screenshot is of element only
+			Expect(img.Bounds().Size()).To(Equal(image.Point{X: 71, Y: 19}))
+		})
 	})
 }

internal/mocks/crop.go

@@ -0,0 +1,22 @@
+package mocks
+
+import (
+	"image"
+)
+
+type Cropper struct {
+	Image       image.Image
+	Width       int
+	Height      int
+	Anchor      image.Point
+	ReturnImage image.Image
+	Err         error
+}
+
+func (c *Cropper) Crop(img image.Image, width, height int, anchor image.Point) (image.Image, error) {
+	c.Image = img
+	c.Width = width
+	c.Height = height
+	c.Anchor = anchor
+	return c.ReturnImage, c.Err
+}

internal/mocks/element.go

@@ -87,6 +87,19 @@ type Element struct {
 		ReturnY int
 		Err     error
 	}
+
+	GetRectCall struct {
+		ReturnX      int
+		ReturnY      int
+		ReturnHeight int
+		ReturnWidth  int
+		Err          error
+	}
+
+	GetScreenshotCall struct {
+		ReturnImage []byte
+		Err         error
+	}
 }
 
 func (e *Element) GetElement(selector api.Selector) (*api.Element, error) {
@@ -161,3 +174,11 @@ func (e *Element) IsEqualTo(other *api.Element) (bool, error) {
 func (e *Element) GetLocation() (x, y int, err error) {
 	return e.GetLocationCall.ReturnX, e.GetLocationCall.ReturnY, e.GetLocationCall.Err
 }
+
+func (e *Element) GetRect() (x, y, width, height int, err error) {
+	return e.GetRectCall.ReturnX, e.GetRectCall.ReturnY, e.GetRectCall.ReturnWidth, e.GetRectCall.ReturnHeight, e.GetRectCall.Err
+}
+
+func (e *Element) GetScreenshot() (s []byte, err error) {
+	return e.GetScreenshotCall.ReturnImage, e.GetScreenshotCall.Err
+}

multiselection_test.go

@@ -18,7 +18,7 @@ var _ = Describe("MultiSelection", func() {
 	BeforeEach(func() {
 		bus = &mocks.Bus{}
 		session = &api.Session{Bus: bus}
-		selection = NewTestMultiSelection(session, nil, "#selector")
+		selection = NewTestMultiSelection(session, nil, "#selector", nil)
 	})
 
 	Describe("#At", func() {

selection.go

@@ -1,9 +1,15 @@
 package agouti
 
 import (
+	"bytes"
 	"fmt"
+	"image"
+	"image/png"
+	"io/ioutil"
+	"path/filepath"
 
 	"github.com/sclevine/agouti/api"
+	"github.com/sclevine/agouti/crop"
 	"github.com/sclevine/agouti/internal/element"
 	"github.com/sclevine/agouti/internal/target"
 )
@@ -24,6 +30,7 @@ import (
 type Selection struct {
 	selectable
 	elements elementRepository
+	cropper  Cropper
 }
 
 type elementRepository interface {
@@ -39,6 +46,7 @@ func newSelection(session apiSession, selectors target.Selectors) *Selection {
 			Client:    session,
 			Selectors: selectors,
 		},
+		crop.CropperFunc(crop.Crop),
 	}
 }
 
@@ -115,3 +123,60 @@ func (s *Selection) MouseToElement() error {
 
 	return nil
 }
+
+// Screenshot takes a screenshot of exactly one element
+// and saves it to the provided filename.
+// The provided filename may be an absolute or relative path.
+func (s *Selection) Screenshot(filename string) error {
+	selectedElement, err := s.elements.GetExactlyOne()
+	if err != nil {
+		return fmt.Errorf("failed to select element from %s: %s", s, err)
+	}
+
+	absFilePath, err := filepath.Abs(filename)
+	if err != nil {
+		return fmt.Errorf("failed to find absolute path for filename: %s", err)
+	}
+
+	screenshot, err := selectedElement.GetScreenshot()
+	if err != nil {
+		// Fallback to getting full size screenshot and cropping
+		data, err := s.session.GetScreenshot()
+		if err != nil {
+			return fmt.Errorf("failed to retrieve screenshot: %s", err)
+		}
+
+		img, err := png.Decode(bytes.NewBuffer(data))
+		if err != nil {
+			return fmt.Errorf("failed to decode screenshot: %s", err)
+		}
+
+		x, y, width, height, err := selectedElement.GetRect()
+		if err != nil {
+			return fmt.Errorf("failed to retrieve bounds for selection: %s", err)
+		}
+
+		fmt.Println(x, y, width, height)
+
+		croppedImg, err := s.cropper.Crop(
+			img, width, height,
+			image.Point{X: x, Y: y})
+		if err != nil {
+			return fmt.Errorf("failed to crop screenshot: %s", err)
+		}
+
+		b := new(bytes.Buffer)
+		err = png.Encode(b, croppedImg)
+		if err != nil {
+			return fmt.Errorf("failed to encode screenshot: %s", err)
+		}
+
+		screenshot = b.Bytes()
+	}
+
+	if err := ioutil.WriteFile(absFilePath, screenshot, 0666); err != nil {
+		return fmt.Errorf("failed to save screenshot: %s", err)
+	}
+
+	return nil
+}

selection_actions_test.go

@@ -27,7 +27,7 @@ var _ = Describe("Selection Actions", func() {
 		firstElement = &mocks.Element{}
 		secondElement = &mocks.Element{}
 		elementRepository = &mocks.ElementRepository{}
-		selection = NewTestMultiSelection(session, elementRepository, "#selector")
+		selection = NewTestMultiSelection(session, elementRepository, "#selector", nil)
 		elementRepository.GetAtLeastOneCall.ReturnElements = []element.Element{firstElement, secondElement}
 	})
 

selection_frames_test.go

@@ -21,7 +21,7 @@ var _ = Describe("Selection Frames", func() {
 	BeforeEach(func() {
 		session = &mocks.Session{}
 		elementRepository = &mocks.ElementRepository{}
-		selection = NewTestSelection(session, elementRepository, "#selector")
+		selection = NewTestSelection(session, elementRepository, "#selector", nil)
 	})
 
 	Describe("#SwitchToFrame", func() {

selection_properties_test.go

@@ -26,7 +26,7 @@ var _ = Describe("Selection Properties", func() {
 		firstElement = &mocks.Element{}
 		secondElement = &mocks.Element{}
 		elementRepository = &mocks.ElementRepository{}
-		selection = NewTestMultiSelection(session, elementRepository, "#selector")
+		selection = NewTestMultiSelection(session, elementRepository, "#selector", nil)
 	})
 
 	Describe("#Text", func() {