Add Screenshot function to Selection.

injector_test.go

@@ -1,15 +1,18 @@
 package agouti
 
-import "github.com/sclevine/agouti/internal/target"
+import (
+	"github.com/sclevine/agouti/internal/crop"
+	"github.com/sclevine/agouti/internal/target"
+)
 
-func NewTestSelection(session apiSession, elements elementRepository, firstSelector string) *Selection {
+func NewTestSelection(session apiSession, elements elementRepository, firstSelector string, cropper crop.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 crop.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/mocks/crop.go

@@ -0,0 +1,20 @@
+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/internal/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  crop.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,58 @@ 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)
+		}
+
+		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() {

selection_test.go

@@ -1,7 +1,13 @@
 package agouti_test
 
 import (
+	"bytes"
 	"errors"
+	"image"
+	"image/png"
+	"io/ioutil"
+	"os"
+	"path/filepath"
 
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
@@ -12,6 +18,15 @@ import (
 	"github.com/sclevine/agouti/internal/mocks"
 )
 
+var minimalpng = []byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a,
+	0x0a, 0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52, 0x00,
+	0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00,
+	0x00, 0x00, 0x1f, 0x15, 0xc4, 0x89, 0x00, 0x00, 0x00, 0x11,
+	0x49, 0x44, 0x41, 0x54, 0x78, 0x9c, 0x62, 0x62, 0x60, 0x60,
+	0x60, 0x00, 0x04, 0x00, 0x00, 0xff, 0xff, 0x00, 0x0f, 0x00,
+	0x03, 0xfe, 0x8f, 0xeb, 0xcf, 0x00, 0x00, 0x00, 0x00, 0x49,
+	0x45, 0x4e, 0x44, 0xae, 0x42, 0x60, 0x82}
+
 var _ = Describe("Selection", func() {
 	var (
 		firstElement  *mocks.Element
@@ -25,7 +40,7 @@ var _ = Describe("Selection", func() {
 
 	Describe("#String", func() {
 		It("should return a string representation of the selection", func() {
-			selection := NewTestMultiSelection(nil, nil, "#selector")
+			selection := NewTestMultiSelection(nil, nil, "#selector", nil)
 			Expect(selection.AllByXPath("#subselector").String()).To(Equal("selection 'CSS: #selector | XPath: #subselector'"))
 		})
 	})
@@ -38,7 +53,7 @@ var _ = Describe("Selection", func() {
 
 		BeforeEach(func() {
 			elementRepository = &mocks.ElementRepository{}
-			selection = NewTestSelection(nil, elementRepository, "#selector")
+			selection = NewTestSelection(nil, elementRepository, "#selector", nil)
 		})
 
 		It("should return a []*api.Elements retrieved from the element repository", func() {
@@ -64,7 +79,7 @@ var _ = Describe("Selection", func() {
 
 		BeforeEach(func() {
 			elementRepository = &mocks.ElementRepository{}
-			selection = NewTestMultiSelection(nil, elementRepository, "#selector")
+			selection = NewTestMultiSelection(nil, elementRepository, "#selector", nil)
 			elementRepository.GetCall.ReturnElements = []element.Element{firstElement, secondElement}
 		})
 
@@ -92,11 +107,11 @@ var _ = Describe("Selection", func() {
 		BeforeEach(func() {
 			firstElementRepository = &mocks.ElementRepository{}
 			firstElementRepository.GetExactlyOneCall.ReturnElement = firstElement
-			firstSelection = NewTestSelection(nil, firstElementRepository, "#first_selector")
+			firstSelection = NewTestSelection(nil, firstElementRepository, "#first_selector", nil)
 
 			secondElementRepository = &mocks.ElementRepository{}
 			secondElementRepository.GetExactlyOneCall.ReturnElement = secondElement
-			secondSelection = NewTestSelection(nil, secondElementRepository, "#second_selector")
+			secondSelection = NewTestSelection(nil, secondElementRepository, "#second_selector", nil)
 		})
 
 		It("should compare the selection elements for equality", func() {
@@ -116,7 +131,7 @@ var _ = Describe("Selection", func() {
 
 		Context("when the provided object is a *MultiSelection", func() {
 			It("should not fail", func() {
-				multiSelection := NewTestMultiSelection(nil, secondElementRepository, "#multi_selector")
+				multiSelection := NewTestMultiSelection(nil, secondElementRepository, "#multi_selector", nil)
 				Expect(firstSelection.EqualsElement(multiSelection)).To(BeFalse())
 				Expect(firstElement.IsEqualToCall.Element).To(ExactlyEqual(secondElement))
 			})
@@ -165,7 +180,7 @@ var _ = Describe("Selection", func() {
 			elementRepository = &mocks.ElementRepository{}
 			elementRepository.GetExactlyOneCall.ReturnElement = secondElement
 			session = &mocks.Session{}
-			selection = NewTestSelection(session, elementRepository, "#selector")
+			selection = NewTestSelection(session, elementRepository, "#selector", nil)
 		})
 
 		It("should successfully instruct the session to move the mouse over the selection", func() {
@@ -190,4 +205,106 @@ var _ = Describe("Selection", func() {
 			})
 		})
 	})
+
+	Describe("#Screenshot", func() {
+		var (
+			selection         *Selection
+			session           *mocks.Session
+			cropper           *mocks.Cropper
+			firstElement      *mocks.Element
+			elementRepository *mocks.ElementRepository
+		)
+
+		BeforeEach(func() {
+			firstElement = &mocks.Element{}
+			elementRepository = &mocks.ElementRepository{}
+			elementRepository.GetExactlyOneCall.ReturnElement = firstElement
+			session = &mocks.Session{}
+			cropper = &mocks.Cropper{}
+			selection = NewTestSelection(session, elementRepository, "#selector", cropper)
+		})
+
+		It("should successfully return the screenshot", 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 selection fails to retrieve a screenshot", func() {
+			BeforeEach(func() {
+				firstElement.GetScreenshotCall.Err = errors.New("some error")
+			})
+
+			It("should fall back to using the session screenshot and cropping", func() {
+				session.GetScreenshotCall.ReturnImage = minimalpng
+				cropper.ReturnImage, _ = png.Decode(bytes.NewBuffer(minimalpng))
+				filename, _ := filepath.Abs(".test.screenshot.png")
+				Expect(selection.Screenshot(".test.screenshot.png")).To(Succeed())
+				defer os.Remove(filename)
+				result, _ := ioutil.ReadFile(filename)
+				Expect(result).To(Equal(minimalpng))
+			})
+
+			Context("and the session fails to retrieve a screenshot", func() {
+				It("should return an error", func() {
+					session.GetScreenshotCall.Err = errors.New("some error")
+					err := selection.Screenshot(".test.screenshot.png")
+					Expect(err).To(MatchError("failed to retrieve screenshot: some error"))
+				})
+			})
+
+			Context("and the session screenshot cannot be decoded", func() {
+				It("should return an error", func() {
+					session.GetScreenshotCall.ReturnImage = []byte("some-image")
+					err := selection.Screenshot(".test.screenshot.png")
+					Expect(err).To(MatchError("failed to decode screenshot: png: invalid format: not a PNG file"))
+				})
+			})
+
+			Context("and the selections bounding rectangle cannot be retrieved", func() {
+				It("should return an error", func() {
+					session.GetScreenshotCall.ReturnImage = minimalpng
+					firstElement.GetRectCall.Err = errors.New("some error")
+					err := selection.Screenshot(".test.screenshot.png")
+					Expect(err).To(MatchError("failed to retrieve bounds for selection: some error"))
+				})
+			})
+
+			Context("and the image cannot be cropped", func() {
+				It("should return an error", func() {
+					session.GetScreenshotCall.ReturnImage = minimalpng
+					cropper.Err = errors.New("some error")
+					err := selection.Screenshot(".test.screenshot.png")
+					Expect(err).To(MatchError("failed to crop screenshot: some error"))
+				})
+			})
+
+			Context("and the image cannot be encoded", func() {
+				It("should return an error", func() {
+					session.GetScreenshotCall.ReturnImage = minimalpng
+					cropper.ReturnImage = image.Rectangle{}
+					err := selection.Screenshot(".test.screenshot.png")
+					Expect(err).To(MatchError("failed to encode screenshot: png: invalid format: invalid image size: 0x0"))
+				})
+			})
+		})
+	})
 })