add before/afterGetScreenshotAs to WebDriverEventListener

Author: morrishoresh

java/client/src/org/openqa/selenium/support/events/AbstractWebDriverEventListener.java

@@ -18,6 +18,7 @@
 package org.openqa.selenium.support.events;
 
 import org.openqa.selenium.By;
+import org.openqa.selenium.OutputType;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebElement;
 
@@ -118,4 +119,13 @@ public void beforeSwitchToWindow(String windowName, WebDriver driver) {
   public void onException(Throwable throwable, WebDriver driver) {
     // Do nothing
   }
+  
+  public <X> void beforeGetScreenshotAs(OutputType<X> target) {
+    // Do nothing
+  }
+
+  public <X> void afterGetScreenshotAs(OutputType<X> target, X screenshot) {
+    // Do nothing
+  }
+
 }

java/client/src/org/openqa/selenium/support/events/EventFiringWebDriver.java

@@ -280,10 +280,13 @@ private Object wrapResult(Object result) {
     return result;
   }
 
-  public <X> X getScreenshotAs(OutputType<X> target) throws WebDriverException {
-    if (driver instanceof TakesScreenshot) {
-      return ((TakesScreenshot) driver).getScreenshotAs(target);
-    }
+   public <X> X getScreenshotAs(OutputType<X> target) throws WebDriverException {
+     if (driver instanceof TakesScreenshot) {
+        dispatcher.beforeGetScreenshotAs(target);
+        X screenshot = ((TakesScreenshot) driver).getScreenshotAs(target);
+        dispatcher.afterGetScreenshotAs(target, screenshot);
+        return screenshot;
+     }
 
     throw new UnsupportedOperationException(
         "Underlying driver instance does not support taking screenshots");

java/client/src/org/openqa/selenium/support/events/WebDriverEventListener.java

@@ -19,6 +19,7 @@
 
 import org.openqa.selenium.Alert;
 import org.openqa.selenium.By;
+import org.openqa.selenium.OutputType;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebElement;
 
@@ -216,4 +217,25 @@ public interface WebDriverEventListener {
    * @param throwable the exception that will be thrown
    */
   void onException(Throwable throwable, WebDriver driver);
+  
+  /**
+   * Called before {@link org.openqa.selenium.TakesScreenshot#getScreenshotAs(OutputType)}
+   * allows the implementation to determine which type of output will be generated
+   *
+   * @param <X> Return type for getScreenshotAs.
+   * @param target target type, @see OutputType
+   */
+  <X> void beforeGetScreenshotAs(OutputType<X> target);
+  
+  /**
+   * Called after {@link org.openqa.selenium.TakesScreenshot#getScreenshotAs(OutputType)}
+   * allows the implementation to determine which type of output was generated
+   * and to access the output itself
+   *
+   * @param <X> Return type for getScreenshotAs.
+   * @param target target type, @see OutputType
+   * @param screenshot screenshot output of the specified type
+   */
+  <X> void afterGetScreenshotAs(OutputType<X> target, X screenshot);
+
 }

java/client/test/org/openqa/selenium/support/events/EventFiringWebDriverTest.java

@@ -43,7 +43,9 @@
 import org.openqa.selenium.By;
 import org.openqa.selenium.JavascriptExecutor;
 import org.openqa.selenium.NoSuchElementException;
+import org.openqa.selenium.OutputType;
 import org.openqa.selenium.StubDriver;
+import org.openqa.selenium.TakesScreenshot;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebDriver.Navigation;
 import org.openqa.selenium.WebDriver.TargetLocator;
@@ -450,4 +452,25 @@ public void shouldReturnLocatorFromToStringMethod() {
   }
 
   private static class ChildDriver extends StubDriver {}
+  
+  @Test
+  public void getScreenshotAs() {
+    final String DATA = "data";
+    WebDriver mockedDriver = mock(WebDriver.class, withSettings().extraInterfaces(TakesScreenshot.class));
+    WebDriverEventListener listener = mock(WebDriverEventListener.class);
+    EventFiringWebDriver testedDriver = new EventFiringWebDriver(mockedDriver).register(listener);
+    
+    Mockito.doReturn(DATA).when((TakesScreenshot)mockedDriver).getScreenshotAs(OutputType.BASE64);
+    
+    String screenshot = ((TakesScreenshot)testedDriver).getScreenshotAs(OutputType.BASE64);
+    assertTrue(screenshot.equals(DATA));
+    
+    InOrder order = Mockito.inOrder(mockedDriver, listener);
+    order.verify(listener).beforeGetScreenshotAs(OutputType.BASE64);
+    order.verify((TakesScreenshot)mockedDriver).getScreenshotAs(OutputType.BASE64);
+    order.verify(listener).afterGetScreenshotAs(OutputType.BASE64, screenshot);
+    verifyNoMoreInteractions(mockedDriver, listener);
+  
+  }
+
 }