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,13 @@ func (e *Element) GetLocation() (x, y int, err error) {
 func round(number float64) int {
 	return int(number + 0.5)
 }
+
+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,38 @@ var _ = Describe("Element", func() {
 			})
 		})
 	})
+
+	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"))
+			})
+		})
+	})
 })

internal/element/repository.go

@@ -34,6 +34,7 @@ type Element interface {
 	Value(text string) error
 	Submit() error
 	GetLocation() (x, y 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,17 @@ 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
+			// TODO: Correct this number estimate
+			Expect(img.Bounds().Size()).To(Equal(image.Point{X: 60, Y: 18}))
+		})
 	})
 }

internal/mocks/element.go

@@ -87,6 +87,11 @@ type Element struct {
 		ReturnY int
 		Err     error
 	}
+
+	GetScreenshotCall struct {
+		ReturnImage []byte
+		Err         error
+	}
 }
 
 func (e *Element) GetElement(selector api.Selector) (*api.Element, error) {
@@ -161,3 +166,7 @@ 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) GetScreenshot() (s []byte, err error) {
+	return e.GetScreenshotCall.ReturnImage, e.GetScreenshotCall.Err
+}

selection.go

@@ -2,6 +2,8 @@ package agouti
 
 import (
 	"fmt"
+	"io/ioutil"
+	"path/filepath"
 
 	"github.com/sclevine/agouti/api"
 	"github.com/sclevine/agouti/internal/element"
@@ -115,3 +117,29 @@ 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 {
+		return fmt.Errorf("failed to retrieve screenshot: %s", err)
+	}
+
+	if err := ioutil.WriteFile(absFilePath, screenshot, 0666); err != nil {
+		return fmt.Errorf("failed to save screenshot: %s", err)
+	}
+
+	return nil
+}

selection_test.go

@@ -2,6 +2,9 @@ package agouti_test
 
 import (
 	"errors"
+	"io/ioutil"
+	"os"
+	"path/filepath"
 
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
@@ -190,4 +193,53 @@ var _ = Describe("Selection", func() {
 			})
 		})
 	})
+
+	Describe("#Screenshot", func() {
+		var (
+			selection         *Selection
+			session           *mocks.Session
+			firstElement      *mocks.Element
+			elementRepository *mocks.ElementRepository
+		)
+
+		BeforeEach(func() {
+			firstElement = &mocks.Element{}
+			elementRepository = &mocks.ElementRepository{}
+			elementRepository.GetExactlyOneCall.ReturnElement = firstElement
+			session = &mocks.Session{}
+			selection = NewTestSelection(session, elementRepository, "#selector")
+		})
+
+		It("should successfully return the text", func() {
+			firstElement.GetScreenshotCall.ReturnImage = []byte("some-image")
+			filename, _ := filepath.Abs(".test.screenshot.png")
+			Expect(selection.Screenshot(".test.screenshot.png")).To(Succeed())
+			defer os.Remove(filename)
+			result, _ := ioutil.ReadFile(filename)
+			Expect(string(result)).To(Equal("some-image"))
+		})
+
+		Context("when a new screenshot file cannot be saved", func() {
+			It("should return an error", func() {
+				err := selection.Screenshot("")
+				Expect(err.Error()).To(ContainSubstring("failed to save screenshot: open"))
+			})
+		})
+
+		Context("when the element repository fails to return exactly one element", func() {
+			It("should return an error", func() {
+				elementRepository.GetExactlyOneCall.Err = errors.New("some error")
+				err := selection.Screenshot(".test.screenshot.png")
+				Expect(err).To(MatchError("failed to select element from selection 'CSS: #selector [single]': some error"))
+			})
+		})
+
+		Context("when the session fails to retrieve a screenshot", func() {
+			It("should return an error", func() {
+				firstElement.GetScreenshotCall.Err = errors.New("some error")
+				err := selection.Screenshot(".test.screenshot.png")
+				Expect(err).To(MatchError("failed to retrieve screenshot: some error"))
+			})
+		})
+	})
 })