Make the RemoteWebDriver implement TakesScreenshot.

There's only one driver that doesn't implement the interface (the HtmUnitDriver)
and this is the most common cause for using the Augmenter, which isn't very
discoverable.

Author: shs96c

java/client/src/org/openqa/selenium/remote/RemoteWebDriver.java

@@ -31,9 +31,11 @@
 import org.openqa.selenium.Dimension;
 import org.openqa.selenium.HasCapabilities;
 import org.openqa.selenium.JavascriptExecutor;
+import org.openqa.selenium.OutputType;
 import org.openqa.selenium.Platform;
 import org.openqa.selenium.Point;
 import org.openqa.selenium.SearchContext;
+import org.openqa.selenium.TakesScreenshot;
 import org.openqa.selenium.UnsupportedCommandException;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebDriverException;
@@ -72,7 +74,7 @@
 public class RemoteWebDriver implements WebDriver, JavascriptExecutor,
     FindsById, FindsByClassName, FindsByLinkText, FindsByName,
     FindsByCssSelector, FindsByTagName, FindsByXPath,
-    HasInputDevices, HasCapabilities {
+    HasInputDevices, HasCapabilities, TakesScreenshot {
 
   // TODO(dawagner): This static logger should be unified with the per-instance localLogs
   private static final Logger logger = Logger.getLogger(RemoteWebDriver.class.getName());
@@ -291,6 +293,23 @@ public String getCurrentUrl() {
     return execute(DriverCommand.GET_CURRENT_URL).getValue().toString();
   }
 
+  @Override
+  public <X> X getScreenshotAs(OutputType<X> outputType) throws WebDriverException {
+    Response response = execute(DriverCommand.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);
+      return outputType.convertFromBase64Png(base64EncodedPng);
+    } else {
+      throw new RuntimeException(String.format("Unexpected result for %s command: %s",
+          DriverCommand.SCREENSHOT,
+          result == null ? "null" : result.getClass().getName() + " instance"));
+    }
+  }
+
   public List<WebElement> findElements(By by) {
     return by.findElements(this);
   }

java/client/test/org/openqa/selenium/remote/BaseAugmenterTest.java

@@ -56,38 +56,31 @@ public void shouldReturnANormalWebDriverUntouched() {
   }
 
   @Test
-  public void shouldLeaveARemoteWebDriverWhichCannotTakeSnapshotsAlone() throws Exception {
+  public void shouldAddInterfaceFromCapabilityIfNecessary() {
     DesiredCapabilities caps = new DesiredCapabilities();
+    caps.setCapability("magic.numbers", true);
     WebDriver driver = new RemoteWebDriver(new StubExecutor(caps), caps);
 
-    WebDriver returned = getAugmenter().augment(driver);
-    System.out.println("This is the returned webdriver " + returned);
-    assertSame(driver, returned);
-    assertFalse(returned instanceof TakesScreenshot);
-  }
-
-  @Test
-  public void shouldAddTheTakesSnapshotInterfaceIfNecessary() {
-    DesiredCapabilities caps = new DesiredCapabilities();
-    caps.setCapability(CapabilityType.TAKES_SCREENSHOT, true);
-    WebDriver driver = new RemoteWebDriver(new StubExecutor(caps), caps);
-
-    WebDriver returned = getAugmenter().augment(driver);
+    BaseAugmenter augmenter = getAugmenter();
+    augmenter.addDriverAugmentation("magic.numbers", new AddsMagicNumberHolder());
+    WebDriver returned = augmenter.augment(driver);
 
     assertNotSame(driver, returned);
     assertTrue(returned instanceof TakesScreenshot);
   }
 
   @Test
-  public void shouldNotAddTheTakesSnapshotInterfaceWhenBooleanValueIsFalse() {
+  public void shouldNotAddInterfaceWhenBooleanValueForItIsFalse() {
     DesiredCapabilities caps = new DesiredCapabilities();
-    caps.setCapability(CapabilityType.TAKES_SCREENSHOT, false);
+    caps.setCapability("magic.numbers", false);
     WebDriver driver = new RemoteWebDriver(new StubExecutor(caps), caps);
 
-    WebDriver returned = getAugmenter().augment(driver);
+    BaseAugmenter augmenter = getAugmenter();
+    augmenter.addDriverAugmentation("magic.numbers", new AddsMagicNumberHolder());
+    WebDriver returned = augmenter.augment(driver);
 
     assertSame(driver, returned);
-    assertFalse(returned instanceof TakesScreenshot);
+    assertFalse(returned instanceof MagicNumberHolder);
   }
 
   @Test
@@ -121,12 +114,14 @@ public Object invoke(ExecuteMethod executeMethod, Object self, Method method,
   @Test
   public void shouldDelegateUnmatchedMethodCallsToDriverImplementation() {
     DesiredCapabilities caps = new DesiredCapabilities();
-    caps.setCapability(CapabilityType.TAKES_SCREENSHOT, true);
+    caps.setCapability("magic.numbers", true);
     StubExecutor stubExecutor = new StubExecutor(caps);
     stubExecutor.expect(DriverCommand.GET_TITLE, new HashMap<String, Object>(), "Title");
     WebDriver driver = new RemoteWebDriver(stubExecutor, caps);
 
-    WebDriver returned = getAugmenter().augment(driver);
+    BaseAugmenter augmenter = getAugmenter();
+    augmenter.addDriverAugmentation("magic.numbers", new AddsMagicNumberHolder());
+    WebDriver returned = augmenter.augment(driver);
 
     assertEquals("Title", returned.getTitle());
   }
@@ -135,12 +130,14 @@ public void shouldDelegateUnmatchedMethodCallsToDriverImplementation() {
   public void proxyShouldNotAppearInStackTraces() {
     final DesiredCapabilities caps = new DesiredCapabilities();
     // This will force the class to be enhanced
-    caps.setCapability(CapabilityType.TAKES_SCREENSHOT, true);
+    caps.setCapability("magic.numbers", true);
 
     DetonatingDriver driver = new DetonatingDriver();
     driver.setCapabilities(caps);
 
-    WebDriver returned = getAugmenter().augment(driver);
+    BaseAugmenter augmenter = getAugmenter();
+    augmenter.addDriverAugmentation("magic.numbers", new AddsMagicNumberHolder());
+    WebDriver returned = augmenter.augment(driver);
 
     returned.findElement(By.id("ignored"));
   }
@@ -222,7 +219,7 @@ public void shouldNotChokeOnFinalFields() {
   public void shouldBeAbleToAugmentMultipleTimes() {
     DesiredCapabilities caps = new DesiredCapabilities();
     caps.setCapability("canRotate", true);
-    caps.setCapability("reallyTakesScreenshot", true);
+    caps.setCapability("magic.numbers", true);
 
     StubExecutor stubExecutor = new StubExecutor(caps);
     stubExecutor.expect(DriverCommand.GET_SCREEN_ORIENTATION,
@@ -236,15 +233,15 @@ public void shouldBeAbleToAugmentMultipleTimes() {
     WebDriver augmented = augmenter.augment(driver);
     assertNotSame(augmented, driver);
     assertTrue(augmented instanceof Rotatable);
-    assertFalse(augmented instanceof TakesScreenshot);
+    assertFalse(augmented instanceof MagicNumberHolder);
 
     augmenter = getAugmenter();
-    augmenter.addDriverAugmentation("reallyTakesScreenshot", new AddTakesScreenshot());
+    augmenter.addDriverAugmentation("magic.numbers", new AddsMagicNumberHolder());
 
     WebDriver augmentedAgain = augmenter.augment(augmented);
     assertNotSame(augmentedAgain, augmented);
     assertTrue(augmentedAgain instanceof Rotatable);
-    assertTrue(augmentedAgain instanceof TakesScreenshot);
+    assertTrue(augmentedAgain instanceof MagicNumberHolder);
 
     ((Rotatable) augmentedAgain).getOrientation();  // Should not throw.
 
@@ -355,4 +352,22 @@ public Capabilities getCapabilities() {
   }
 
   public abstract BaseAugmenter getAugmenter();
+
+  private static class AddsMagicNumberHolder implements AugmenterProvider {
+    @Override
+    public Class<?> getDescribedInterface() {
+      return MagicNumberHolder.class;
+    }
+
+    @Override
+    public InterfaceImplementation getImplementation(Object value) {
+      return new InterfaceImplementation() {
+        @Override
+        public Object invoke(ExecuteMethod executeMethod, Object self, Method method,
+                             Object... args) {
+          return null;
+        }
+      };
+    }
+  }
 }