Fixes #68: Adding after scroll timeout, before taking screenshot using Chrome commands

Author: glibas

src/main/java/com/assertthat/selenium_shutterbug/core/Shutterbug.java

@@ -65,11 +65,26 @@ public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll) {
      *
      * @param driver WebDriver instance
      * @param scroll ScrollStrategy How you need to scroll
-     * @param  scrollTimeout Timeout to wait after scrolling and before taking screen shot
+     * @param  betweenScrollTimeout Timeout to wait after scrolling and before taking screenshot
      * @return PageSnapshot instance
      */
-    public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int scrollTimeout) {
-        return shootPage(driver,scroll,scrollTimeout,false);
+    public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int betweenScrollTimeout) {
+        return shootPage(driver,scroll,betweenScrollTimeout,false);
+    }
+
+    /**
+     * To be used when screen shooting the page
+     * and need to scroll while making screen shots, either vertically or
+     * horizontally or both directions (Chrome).
+     *
+     * @param driver WebDriver instance
+     * @param scroll ScrollStrategy How you need to scroll
+     * @param  betweenScrollTimeout Timeout to wait after scrolling and before taking screenshot
+     * @param  afterScrollTimeout Timeout to wait after scrolling and before taking screenshot
+     * @return PageSnapshot instance
+     */
+    public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int betweenScrollTimeout, int afterScrollTimeout) {
+        return shootPage(driver,scroll,betweenScrollTimeout,false, afterScrollTimeout);
     }
 
     /**
@@ -93,13 +108,45 @@ public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, bo
      *
      * @param driver WebDriver instance
      * @param scroll ScrollStrategy How you need to scroll
-     * @param  scrollTimeout Timeout to wait after scrolling and before taking screen shot
+     * @param  useDevicePixelRatio whether or not take into account device pixel ratio
+     * @param  afterScrollTimeout Timeout to wait after scrolling and before taking screenshot
+     * @return PageSnapshot instance
+     */
+    public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, boolean useDevicePixelRatio, int afterScrollTimeout) {
+        return shootPage(driver,scroll,0,useDevicePixelRatio, afterScrollTimeout);
+    }
+
+    /**
+     * To be used when screen shooting the page
+     * and need to scroll while making screen shots, either vertically or
+     * horizontally or both directions (Chrome).
+     *
+     * @param driver WebDriver instance
+     * @param scroll ScrollStrategy How you need to scroll
+     * @param  betweenScrollTimeout Timeout to wait between each scrolling operation
+     * @param useDevicePixelRatio whether or not take into account device pixel ratio
+     * @return PageSnapshot instance
+     */
+    public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int betweenScrollTimeout, boolean useDevicePixelRatio) {
+        return shootPage(driver,scroll,betweenScrollTimeout,useDevicePixelRatio,0);
+    }
+
+    /**
+     * To be used when screen shooting the page
+     * and need to scroll while making screen shots, either vertically or
+     * horizontally or both directions (Chrome).
+     *
+     * @param driver WebDriver instance
+     * @param scroll ScrollStrategy How you need to scroll
+     * @param  betweenScrollTimeout Timeout to wait between each scrolling operation
+     * @param  afterScrollTimeout Timeout to wait after scrolling and before taking screenshot
      * @param useDevicePixelRatio whether or not take into account device pixel ratio
      * @return PageSnapshot instance
      */
-    public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int scrollTimeout, boolean useDevicePixelRatio) {
+    public static PageSnapshot shootPage(WebDriver driver, ScrollStrategy scroll, int betweenScrollTimeout, boolean useDevicePixelRatio, int afterScrollTimeout) {
         Browser browser = new Browser(driver, useDevicePixelRatio);
-        browser.setScrollTimeout(scrollTimeout);
+        browser.setBetweenScrollTimeout(betweenScrollTimeout);
+        browser.setAfterScrollTimeout(afterScrollTimeout);
 
         PageSnapshot pageScreenshot = new PageSnapshot(driver, browser.getDevicePixelRatio());
         switch (scroll) {

src/main/java/com/assertthat/selenium_shutterbug/utils/web/Browser.java

@@ -58,7 +58,9 @@ public class Browser {
     private int docWidth = -1;
     private int viewportWidth = -1;
     private int viewportHeight = -1;
-    private int scrollTimeout;
+    private int betweenScrollTimeout;
+
+    private int afterScrollTimeout;
     private Double devicePixelRatio = 1.0;
 
     public Browser(WebDriver driver, boolean useDevicePixelRatio) {
@@ -81,10 +83,15 @@ public static void wait(int ms) {
         }
     }
 
-    public void setScrollTimeout(int scrollTimeout) {
-        this.scrollTimeout = scrollTimeout;
+    public void setBetweenScrollTimeout(int betweenScrollTimeout) {
+        this.betweenScrollTimeout = betweenScrollTimeout;
+    }
+
+    public void setAfterScrollTimeout(int afterScrollTimeout) {
+        this.afterScrollTimeout = afterScrollTimeout;
     }
 
+
     public BufferedImage takeScreenshot() {
         File srcFile = ((TakesScreenshot) this.getUnderlyingDriver()).getScreenshotAs(OutputType.FILE);
         try {
@@ -161,7 +168,7 @@ public BufferedImage takeScreenshotEntirePageDefault() {
             this.scrollTo(0, j * _viewportHeight);
             for (int i = 0; i < horizontalIterations; i++) {
                 this.scrollTo(i * _viewportWidth, _viewportHeight * j);
-                wait(scrollTimeout);
+                wait(betweenScrollTimeout);
                 Image image = takeScreenshot();
                 g.drawImage(image, this.getCurrentScrollX(), this.getCurrentScrollY(), null);
                 if (_docWidth == image.getWidth(null) && _docHeight == image.getHeight(null)) {
@@ -183,10 +190,11 @@ public BufferedImage takeScreenshotEntirePageUsingChromeCommand() {
         int verticalIterations = (int) Math.ceil(((double) this.getDocHeight()) / this.getViewportHeight());
         for (int j = 0; j < verticalIterations; j++) {
             this.scrollTo(0, j * this.getViewportHeight());
-            wait(scrollTimeout);
+            wait(betweenScrollTimeout);
         }
         Object metrics = this.evaluate(FileUtil.getJsScript(ALL_METRICS));
         this.sendCommand("Emulation.setDeviceMetricsOverride", metrics);
+        wait(afterScrollTimeout);
         Object result = this.sendCommand("Page.captureScreenshot", ImmutableMap.of("format", "png", "fromSurface", true));
         this.sendCommand("Emulation.clearDeviceMetricsOverride", ImmutableMap.of());
         return decodeBase64EncodedPng((String) ((Map<String, ?>) result).get("data"));