Allow Firefox full page screenshot functionality to be augmented by R…

…emoteWebDriver

java/src/org/openqa/selenium/firefox/AddHasFullPageScreenshot.java

@@ -0,0 +1,78 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.openqa.selenium.firefox;
+
+import com.google.auto.service.AutoService;
+import com.google.common.collect.ImmutableMap;
+import org.openqa.selenium.Capabilities;
+import org.openqa.selenium.OutputType;
+import org.openqa.selenium.remote.AdditionalHttpCommands;
+import org.openqa.selenium.remote.AugmenterProvider;
+import org.openqa.selenium.remote.CommandInfo;
+import org.openqa.selenium.remote.ExecuteMethod;
+import org.openqa.selenium.remote.http.HttpMethod;
+
+import java.util.Map;
+import java.util.function.Predicate;
+
+import static org.openqa.selenium.remote.BrowserType.FIREFOX;
+
+@AutoService({AdditionalHttpCommands.class, AugmenterProvider.class})
+public class AddHasFullPageScreenshot<X> implements AugmenterProvider<HasFullPageScreenshot>, AdditionalHttpCommands {
+
+  public static final String FULL_PAGE_SCREENSHOT = "fullPageScreenshot";
+
+  private static final Map<String, CommandInfo> COMMANDS = ImmutableMap.of(
+    FULL_PAGE_SCREENSHOT, new CommandInfo("/session/:sessionId/moz/screenshot/full", HttpMethod.GET));
+
+  @Override
+  public Map<String, CommandInfo> getAdditionalCommands() {
+    return COMMANDS;
+  }
+
+  @Override
+  public Predicate<Capabilities> isApplicable() {
+    return caps -> FIREFOX.equals(caps.getBrowserName());
+  }
+
+  @Override
+  public Class<HasFullPageScreenshot> getDescribedInterface() {
+    return HasFullPageScreenshot.class;
+  }
+
+  @Override
+  public HasFullPageScreenshot getImplementation(Capabilities capabilities, ExecuteMethod executeMethod) {
+    return new HasFullPageScreenshot() {
+      @Override public <X> X getFullPageScreenshotAs(OutputType<X> outputType) {
+        Object result = executeMethod.execute(FULL_PAGE_SCREENSHOT, null);
+
+        if (result instanceof String) {
+          String base64EncodedPng = (String) result;
+          return outputType.convertFromBase64Png(base64EncodedPng);
+        } else if (result instanceof byte[]) {
+          return outputType.convertFromPngBytes((byte[]) result);
+        } else {
+          throw new RuntimeException(String.format(
+            "Unexpected result for %s command: %s",
+            FULL_PAGE_SCREENSHOT,
+            result == null ? "null" : result.getClass().getName() + " instance"));
+        }
+      }
+    };
+  }
+}

java/src/org/openqa/selenium/firefox/FirefoxDriver.java

@@ -42,11 +42,9 @@
 import org.openqa.selenium.remote.CommandInfo;
 import org.openqa.selenium.remote.FileDetector;
 import org.openqa.selenium.remote.RemoteWebDriver;
-import org.openqa.selenium.remote.Response;
 import org.openqa.selenium.remote.html5.RemoteWebStorage;
 import org.openqa.selenium.remote.http.ClientConfig;
 import org.openqa.selenium.remote.http.HttpClient;
-import org.openqa.selenium.remote.http.HttpMethod;
 import org.openqa.selenium.remote.service.DriverCommandExecutor;
 import org.openqa.selenium.remote.service.DriverService;
 
@@ -58,7 +56,6 @@
 import java.util.Set;
 import java.util.stream.StreamSupport;
 
-import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.openqa.selenium.remote.CapabilityType.PROXY;
 
 /**
@@ -140,23 +137,14 @@ public static final class Capability {
     public static final String MARIONETTE = "marionette";
   }
 
-  static class ExtraCommands {
-    static String FULL_PAGE_SCREENSHOT = "fullPageScreenshot";
-  }
-
-  private static final ImmutableMap<String, CommandInfo> EXTRA_COMMANDS = ImmutableMap.of(
-      ExtraCommands.FULL_PAGE_SCREENSHOT,
-      new CommandInfo("/session/:sessionId/moz/screenshot/full", HttpMethod.GET)
-  );
-
   private static class FirefoxDriverCommandExecutor extends DriverCommandExecutor {
     public FirefoxDriverCommandExecutor(DriverService service) {
       super(service, getExtraCommands());
     }
 
     private static Map<String, CommandInfo> getExtraCommands() {
       return ImmutableMap.<String, CommandInfo>builder()
-        .putAll(EXTRA_COMMANDS)
+        .putAll(new AddHasFullPageScreenshot().getAdditionalCommands())
         .putAll(new AddHasExtensions().getAdditionalCommands())
         .build();
     }
@@ -166,6 +154,7 @@ private static Map<String, CommandInfo> getExtraCommands() {
   protected FirefoxBinary binary;
   private final RemoteWebStorage webStorage;
   private final HasExtensions extensions;
+  private final HasFullPageScreenshot fullPageScreenshot;
   private final Optional<URI> cdpUri;
   private DevTools devTools;
 
@@ -207,6 +196,7 @@ private FirefoxDriver(FirefoxDriverCommandExecutor executor, FirefoxOptions opti
     super(executor, dropCapabilities(options));
     webStorage = new RemoteWebStorage(getExecuteMethod());
     extensions = new AddHasExtensions().getImplementation(getCapabilities(), getExecuteMethod());
+    fullPageScreenshot = new AddHasFullPageScreenshot().getImplementation(getCapabilities(), getExecuteMethod());
 
     Capabilities capabilities = super.getCapabilities();
     HttpClient.Factory clientFactory = HttpClient.Factory.createDefault();
@@ -291,19 +281,7 @@ public void uninstallExtension(String extensionId) {
    * @throws WebDriverException on failure.
    */
   public <X> X getFullPageScreenshotAs(OutputType<X> outputType) throws WebDriverException {
-    Response response = execute(ExtraCommands.FULL_PAGE_SCREENSHOT);
-    Object result = response.getValue();
-    if (result instanceof String) {
-      String base64EncodedPng = (String) result;
-      return outputType.convertFromBase64Png(base64EncodedPng);
-    } else if (result instanceof byte[]) {
-      String base64EncodedPng = new String((byte[]) result, UTF_8);
-      return outputType.convertFromBase64Png(base64EncodedPng);
-    } else {
-      throw new RuntimeException(String.format("Unexpected result for %s command: %s",
-                                               ExtraCommands.FULL_PAGE_SCREENSHOT,
-                                               result == null ? "null" : result.getClass().getName() + " instance"));
-    }
+    return fullPageScreenshot.getFullPageScreenshotAs(outputType);
   }
 
   private static Boolean forceMarionetteFromSystemProperty() {

java/src/org/openqa/selenium/firefox/HasFullPageScreenshot.java

@@ -0,0 +1,36 @@
+// Licensed to the Software Freedom Conservancy (SFC) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The SFC licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.openqa.selenium.firefox;
+
+import org.openqa.selenium.Beta;
+import org.openqa.selenium.OutputType;
+
+/**
+ * Used by classes to indicate that they can take a full page screenshot.
+ */
+@Beta
+public interface HasFullPageScreenshot {
+
+  /**
+   * Capture the full page screenshot and store it in the specified location.
+   *
+   * @param outputType target type, @see OutputType
+   * @return Object in which is stored information about the screenshot.
+   */
+    public <X> X getFullPageScreenshotAs(OutputType<X> outputType);
+}

java/test/org/openqa/selenium/firefox/FirefoxDriverTest.java

@@ -26,6 +26,7 @@
 import org.openqa.selenium.Dimension;
 import org.openqa.selenium.ImmutableCapabilities;
 import org.openqa.selenium.JavascriptExecutor;
+import org.openqa.selenium.OutputType;
 import org.openqa.selenium.ParallelTestRunner;
 import org.openqa.selenium.ParallelTestRunner.Worker;
 import org.openqa.selenium.WebDriver;
@@ -562,6 +563,43 @@ public void shouldAllowDriverToBeAugmentedWithHasExtensions() throws MalformedUR
     }
   }
 
+  @Test
+  public void shouldTakeFullPageScreenshot() {
+    File tempFile = ((FirefoxDriver) driver).getFullPageScreenshotAs(OutputType.FILE);
+    assertThat(tempFile.exists()).isTrue();
+    assertThat(tempFile.length()).isGreaterThan(0);
+  }
+
+  @Test
+  public void shouldAllowRemoteWebDriverBuilderToUseHasFullPageScreenshot() {
+    WebDriver driver = RemoteWebDriver.builder()
+      .oneOf(new FirefoxOptions())
+      .address("http://localhost:4444/")
+      .build();
+
+    try {
+      File tempFile = ((HasFullPageScreenshot) driver).getFullPageScreenshotAs(OutputType.FILE);
+      assertThat(tempFile.exists()).isTrue();
+      assertThat(tempFile.length()).isGreaterThan(0);
+    } finally {
+      driver.quit();
+    }
+  }
+
+  @Test
+  public void shouldAllowDriverToBeAugmentedWithFullHasPageScreenshot() throws MalformedURLException {
+    WebDriver driver = new RemoteWebDriver(new URL("http://localhost:4444/"), new FirefoxOptions());
+    WebDriver augmentedDriver = new Augmenter().augment(driver);
+
+    try {
+      File tempFile = ((HasFullPageScreenshot) augmentedDriver).getFullPageScreenshotAs(OutputType.FILE);
+      assertThat(tempFile.exists()).isTrue();
+      assertThat(tempFile.length()).isGreaterThan(0);
+    } finally {
+      driver.quit();
+    }
+  }
+
   private static class CustomFirefoxProfile extends FirefoxProfile {}
 
 }