Bugfix 39 (#40)

* Unwrapping EventFiringWebDriver

* Fix not making screenshot on RemoteWebDriver

* Fixes #39. Review fixes

Author: anna-shchurok

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

@@ -16,6 +16,7 @@
 import org.openqa.selenium.remote.RemoteWebDriver;
 import org.openqa.selenium.remote.Response;
 import org.openqa.selenium.remote.http.HttpMethod;
+import org.openqa.selenium.support.events.EventFiringWebDriver;
 
 import javax.imageio.ImageIO;
 import java.awt.*;
@@ -61,14 +62,7 @@ public Browser(WebDriver driver, boolean useDevicePixelRatio) {
             this.devicePixelRatio = devicePixelRatio instanceof Double? (Double)devicePixelRatio: (Long)devicePixelRatio*1.0;
         }
     }
-    public Browser(ChromeDriver driver) throws InvocationTargetException, IllegalAccessException, NoSuchMethodException {
-        this.driver = driver;
-        CommandInfo cmd = new CommandInfo("/session/:sessionId/chromium/send_command_and_get_result", HttpMethod.POST);
-        Method defineCommand = HttpCommandExecutor.class.getDeclaredMethod("defineCommand", String.class, CommandInfo.class);
-        defineCommand.setAccessible(true);
-        defineCommand.invoke(driver.getCommandExecutor(), "sendCommand", cmd);
-    }
-  
+
     public Double getDevicePixelRatio() {
         return devicePixelRatio;
     }
@@ -134,12 +128,13 @@ public BufferedImage takeScreenshotEntirePage() {
     }
 
     public BufferedImage takeScreenshotEntirePageUsingChromeCommand() {
-        if(!(this.driver instanceof ChromeDriver))
-            throw new UnsupportedOperationException("You must use Chrome driver for this operation");
+        this.driver = getDriverAfterValidation(this.driver);
 
-        Browser driver;
         try {
-            driver = new Browser((ChromeDriver) this.driver);
+            CommandInfo cmd = new CommandInfo("/session/:sessionId/chromium/send_command_and_get_result", HttpMethod.POST);
+            Method defineCommand = HttpCommandExecutor.class.getDeclaredMethod("defineCommand", String.class, CommandInfo.class);
+            defineCommand.setAccessible(true);
+            defineCommand.invoke(((RemoteWebDriver) this.driver).getCommandExecutor(), "sendCommand", cmd);
         } catch (InvocationTargetException | IllegalAccessException | NoSuchMethodException e) {
             throw new RuntimeException(e);
         }
@@ -149,9 +144,9 @@ public BufferedImage takeScreenshotEntirePageUsingChromeCommand() {
             this.scrollTo(0, j * this.getViewportHeight());
             wait(scrollTimeout);
         }
-        Object metrics = driver.evaluate(FileUtil.getJsScript(ALL_METRICS));
-        driver.sendCommand("Emulation.setDeviceMetricsOverride", metrics);
-        Object result = driver.sendCommand("Page.captureScreenshot", ImmutableMap.of("format", "png", "fromSurface", true));
+        Object metrics = this.evaluate(FileUtil.getJsScript(ALL_METRICS));
+        this.sendCommand("Emulation.setDeviceMetricsOverride", metrics);
+        Object result = this.sendCommand("Page.captureScreenshot", ImmutableMap.of("format", "png", "fromSurface", true));
         String base64EncodedPng = (String) ((Map<String, ?>) result).get("data");
         InputStream in = new ByteArrayInputStream(OutputType.BYTES.convertFromBase64Png(base64EncodedPng));
         BufferedImage bImageFromConvert;
@@ -163,6 +158,21 @@ public BufferedImage takeScreenshotEntirePageUsingChromeCommand() {
         return bImageFromConvert;
     }
 
+    private RemoteWebDriver getDriverAfterValidation(WebDriver driver) {
+        if (driver instanceof EventFiringWebDriver) {
+            driver = ((EventFiringWebDriver) this.driver).getWrappedDriver();
+        }
+
+        if (driver instanceof ChromeDriver) {
+            return (ChromeDriver) driver;
+        } else if (driver instanceof RemoteWebDriver) {
+            if (((RemoteWebDriver) driver).getCapabilities().getBrowserName().equals("chrome")) {
+                return (RemoteWebDriver) driver;
+            }
+        }
+        throw new UnsupportedOperationException("ScrollStrategy.WHOLE_PAGE_CHROME is only applicable for Chrome driver");
+    }
+
     public BufferedImage takeScreenshotScrollHorizontally() {
         BufferedImage combinedImage = new BufferedImage(this.getDocWidth(), this.getViewportHeight(), BufferedImage.TYPE_INT_ARGB);
         Graphics2D g = combinedImage.createGraphics();