[java] Add ability to decorate child classes of WebDriver

java/src/org/openqa/selenium/support/decorators/DefaultDecorated.java

@@ -23,9 +23,9 @@
 public class DefaultDecorated<T> implements Decorated<T> {
 
   private final T original;
-  private final WebDriverDecorator decorator;
+  private final WebDriverDecorator<?> decorator;
 
-  public DefaultDecorated(final T original, final WebDriverDecorator decorator) {
+  public DefaultDecorated(final T original, final WebDriverDecorator<?> decorator) {
     this.original = original;
     this.decorator = decorator;
   }
@@ -34,7 +34,7 @@ public final T getOriginal() {
     return original;
   }
 
-  public final WebDriverDecorator getDecorator() {
+  public final WebDriverDecorator<?> getDecorator() {
     return decorator;
   }
 

java/src/org/openqa/selenium/support/decorators/WebDriverDecorator.java

@@ -179,22 +179,33 @@
  * </code></pre>
  */
 @Beta
-public class WebDriverDecorator {
+public class WebDriverDecorator<T extends WebDriver> {
 
-  private Decorated<WebDriver> decorated;
+  private final Class<T> targetWebDriverClass;
 
-  public final WebDriver decorate(WebDriver original) {
+  private Decorated<T> decorated;
+
+  @SuppressWarnings("unchecked")
+  public WebDriverDecorator() {
+    this((Class<T>) WebDriver.class);
+  }
+
+  public WebDriverDecorator(Class<T> targetClass) {
+    this.targetWebDriverClass = targetClass;
+  }
+
+  public final T decorate(T original) {
     Require.nonNull("WebDriver", original);
 
     decorated = createDecorated(original);
-    return createProxy(decorated, WebDriver.class);
+    return createProxy(decorated, targetWebDriverClass);
   }
 
-  public Decorated<WebDriver> getDecoratedDriver() {
+  public Decorated<T> getDecoratedDriver() {
     return decorated;
   }
 
-  public Decorated<WebDriver> createDecorated(WebDriver driver) {
+  public Decorated<T> createDecorated(T driver) {
     return new DefaultDecorated<>(driver, this);
   }
 
@@ -248,7 +259,7 @@ public Object onError(
 
   private Object decorateResult(Object toDecorate) {
     if (toDecorate instanceof WebDriver) {
-      return createProxy(getDecoratedDriver(), WebDriver.class);
+      return createProxy(getDecoratedDriver(), targetWebDriverClass);
     }
     if (toDecorate instanceof WebElement) {
       return createProxy(createDecorated((WebElement) toDecorate), WebElement.class);
@@ -316,7 +327,7 @@ protected final <Z> Z createProxy(final Decorated<Z> decorated, Class<Z> clazz)
     Class<?>[] allInterfacesArray = allInterfaces.toArray(new Class<?>[0]);
 
     Class<? extends Z> proxy = new ByteBuddy()
-      .subclass(Object.class)
+      .subclass(clazz.isInterface() ? Object.class : clazz)
       .implement(allInterfacesArray)
       .method(ElementMatchers.any())
       .intercept(InvocationHandlerAdapter.of(handler))

java/src/org/openqa/selenium/support/events/EventFiringDecorator.java

@@ -155,7 +155,7 @@
  * extending {@link WebDriverDecorator}, not by creating sophisticated listeners.
  */
 @Beta
-public class EventFiringDecorator extends WebDriverDecorator  {
+public class EventFiringDecorator<T extends WebDriver> extends WebDriverDecorator<T>  {
 
   private static final Logger logger = Logger.getLogger(EventFiringDecorator.class.getName());
 

java/test/org/openqa/selenium/remote/AugmenterTest.java

@@ -440,7 +440,7 @@ public Capabilities getCapabilities() {
     }
   }
 
-  private static class ModifyTitleWebDriverDecorator extends WebDriverDecorator {
+  private static class ModifyTitleWebDriverDecorator extends WebDriverDecorator<WebDriver> {
 
     @Override
     public Object call(Decorated<?> target, Method method, Object[] args) throws Throwable {

java/test/org/openqa/selenium/support/decorators/DecoratedAlertTest.java

@@ -49,7 +49,7 @@ public Fixture() {
       originalDriver = mock(WebDriver.class);
       when(originalSwitch.alert()).thenReturn(original);
       when(originalDriver.switchTo()).thenReturn(originalSwitch);
-      decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+      decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
       decorated = decoratedDriver.switchTo().alert();
     }
   }

java/test/org/openqa/selenium/support/decorators/DecoratedNavigationTest.java

@@ -45,7 +45,7 @@ public Fixture() {
       original = mock(WebDriver.Navigation.class);
       originalDriver = mock(WebDriver.class);
       when(originalDriver.navigate()).thenReturn(original);
-      decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+      decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
       decorated = decoratedDriver.navigate();
     }
   }

java/test/org/openqa/selenium/support/decorators/DecoratedOptionsTest.java

@@ -50,7 +50,7 @@ public Fixture() {
       original = mock(WebDriver.Options.class);
       originalDriver = mock(WebDriver.class);
       when(originalDriver.manage()).thenReturn(original);
-      decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+      decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
       decorated = decoratedDriver.manage();
     }
   }

java/test/org/openqa/selenium/support/decorators/DecoratedRemoteWebDriverTest.java

@@ -48,20 +48,19 @@ void shouldImplementWrapsDriverToProvideAccessToUnderlyingDriver() {
     RemoteWebDriver originalDriver = mock(RemoteWebDriver.class);
     when(originalDriver.getSessionId()).thenReturn(sessionId);
 
-    WebDriver decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+    RemoteWebDriver decoratedDriver = new WebDriverDecorator<>(RemoteWebDriver.class).decorate(originalDriver);
 
-    RemoteWebDriver
-      underlying =
-      (RemoteWebDriver) ((WrapsDriver) decoratedDriver).getWrappedDriver();
+    assertThat(decoratedDriver.getSessionId()).isEqualTo(sessionId);
 
+    RemoteWebDriver underlying = (RemoteWebDriver) ((WrapsDriver) decoratedDriver).getWrappedDriver();
     assertThat(underlying.getSessionId()).isEqualTo(sessionId);
   }
 
   @Test
   void cannotConvertDecoratedToRemoteWebDriver() {
     RemoteWebDriver originalDriver = mock(RemoteWebDriver.class);
 
-    WebDriver decorated = new WebDriverDecorator().decorate(originalDriver);
+    WebDriver decorated = new WebDriverDecorator<>().decorate(originalDriver);
 
     assertThat(decorated).isNotInstanceOf(RemoteWebDriver.class);
   }
@@ -70,7 +69,7 @@ void cannotConvertDecoratedToRemoteWebDriver() {
   void decoratedDriversShouldImplementWrapsDriver() {
     RemoteWebDriver originalDriver = mock(RemoteWebDriver.class);
 
-    WebDriver decorated = new WebDriverDecorator().decorate(originalDriver);
+    WebDriver decorated = new WebDriverDecorator<>().decorate(originalDriver);
 
     assertThat(decorated).isInstanceOf(WrapsDriver.class);
   }
@@ -85,7 +84,7 @@ void decoratedElementsShouldImplementWrapsElement() {
 
     when(originalDriver.findElement(any())).thenReturn(originalElement);
 
-    WebDriver decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+    WebDriver decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
     WebElement element = decoratedDriver.findElement(By.id("test"));
 
     assertThat(element).isInstanceOf(WrapsElement.class);
@@ -101,7 +100,7 @@ void canConvertDecoratedRemoteWebElementToJson() {
 
     when(originalDriver.findElement(any())).thenReturn(originalElement);
 
-    WebDriver decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+    WebDriver decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
 
     WebElement element = decoratedDriver.findElement(By.id("test"));
 

java/test/org/openqa/selenium/support/decorators/DecoratedSwitchToTest.java

@@ -48,7 +48,7 @@ public Fixture() {
       original = mock(WebDriver.TargetLocator.class);
       originalDriver = mock(WebDriver.class);
       when(originalDriver.switchTo()).thenReturn(original);
-      decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+      decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
       decorated = decoratedDriver.switchTo();
     }
   }

java/test/org/openqa/selenium/support/decorators/DecoratedTimeoutsTest.java

@@ -48,7 +48,7 @@ public Fixture() {
       originalDriver = mock(WebDriver.class);
       when(originalOptions.timeouts()).thenReturn(original);
       when(originalDriver.manage()).thenReturn(originalOptions);
-      decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+      decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
       decorated = decoratedDriver.manage().timeouts();
     }
   }

java/test/org/openqa/selenium/support/decorators/DecoratedWebDriverTest.java

@@ -86,9 +86,9 @@ void canCompareDecorated() {
     WebDriver original1 = mock(WebDriver.class);
     WebDriver original2 = mock(WebDriver.class);
 
-    WebDriver decorated1 = new WebDriverDecorator().decorate(original1);
-    WebDriver decorated2 = new WebDriverDecorator().decorate(original1);
-    WebDriver decorated3 = new WebDriverDecorator().decorate(original2);
+    WebDriver decorated1 = new WebDriverDecorator<>().decorate(original1);
+    WebDriver decorated2 = new WebDriverDecorator<>().decorate(original1);
+    WebDriver decorated3 = new WebDriverDecorator<>().decorate(original2);
     assertThat(decorated1).isEqualTo(decorated2);
     assertThat(decorated1).isNotEqualTo(decorated3);
 
@@ -101,7 +101,7 @@ void canCompareDecorated() {
   @Test
   void testHashCode() {
     WebDriver original = mock(WebDriver.class);
-    WebDriver decorated = new WebDriverDecorator().decorate(original);
+    WebDriver decorated = new WebDriverDecorator<>().decorate(original);
     assertThat(decorated.hashCode()).isEqualTo(original.hashCode());
   }
 

java/test/org/openqa/selenium/support/decorators/DecoratedWebElementTest.java

@@ -56,7 +56,7 @@ public Fixture() {
       original = mock(WebElement.class);
       originalDriver = mock(WebDriver.class);
       when(originalDriver.findElement(any())).thenReturn(original);
-      decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+      decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
       decorated = decoratedDriver.findElement(By.id("test"));
     }
   }

java/test/org/openqa/selenium/support/decorators/DecoratedWindowTest.java

@@ -50,7 +50,7 @@ public Fixture() {
       originalDriver = mock(WebDriver.class);
       when(originalOptions.window()).thenReturn(original);
       when(originalDriver.manage()).thenReturn(originalOptions);
-      decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+      decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
       decorated = decoratedDriver.manage().window();
     }
   }

java/test/org/openqa/selenium/support/decorators/IntegrationTest.java

@@ -34,7 +34,7 @@
 @Tag("UnitTests")
 class IntegrationTest {
 
-  static class CountCalls extends WebDriverDecorator {
+  static class CountCalls extends WebDriverDecorator<WebDriver> {
 
     int counterBefore = 0;
     int counterAfter = 0;

java/test/org/openqa/selenium/support/decorators/InterfacesTest.java

@@ -36,7 +36,7 @@ void shouldNotAddInterfacesNotAvailableInTheOriginalDriver() {
     WebDriver driver = mock(WebDriver.class);
     assertThat(driver).isNotInstanceOf(SomeOtherInterface.class);
 
-    WebDriver decorated = new WebDriverDecorator().decorate(driver);
+    WebDriver decorated = new WebDriverDecorator<>().decorate(driver);
     assertThat(decorated).isNotInstanceOf(SomeOtherInterface.class);
   }
 
@@ -45,7 +45,7 @@ void shouldRespectInterfacesAvailableInTheOriginalDriver() {
     WebDriver driver = mock(ExtendedDriver.class);
     assertThat(driver).isInstanceOf(SomeOtherInterface.class);
 
-    WebDriver decorated = new WebDriverDecorator().decorate(driver);
+    WebDriver decorated = new WebDriverDecorator<>().decorate(driver);
     assertThat(decorated).isInstanceOf(SomeOtherInterface.class);
   }
 }