[java] Add ability to decorate child classes of WebDriver

SeleniumHQ · Jun 6, 2022 · 81e80f3 · 81e80f3
1 parent 307673b
commit 81e80f3
Show file tree

Hide file tree

Showing 16 changed files with 66 additions and 36 deletions.
diff --git a/java/src/org/openqa/selenium/support/decorators/DefaultDecorated.java b/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;
   }
 

diff --git a/java/src/org/openqa/selenium/support/decorators/WebDriverDecorator.java b/java/src/org/openqa/selenium/support/decorators/WebDriverDecorator.java
@@ -34,12 +34,15 @@
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * This class helps to create decorators for instances of {@link WebDriver} and
@@ -179,22 +182,35 @@
  * </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 +264,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);
@@ -293,7 +309,8 @@ protected final <Z> Z createProxy(final Decorated<Z> decorated, Class<Z> clazz)
             || decoratedInterfaces.contains(method.getDeclaringClass())) {
           return method.invoke(decorated, args);
         }
-        if (originalInterfaces.contains(method.getDeclaringClass())) {
+        if (originalInterfaces.contains(method.getDeclaringClass()) || findInterfaceByMethod(originalInterfaces,
+          method).isPresent()) {
           decorated.beforeCall(method, args);
           Object result = decorated.call(method, args);
           decorated.afterCall(method, result, args);
@@ -316,7 +333,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))
@@ -332,6 +349,18 @@ protected final <Z> Z createProxy(final Decorated<Z> decorated, Class<Z> clazz)
     }
   }
 
+  private Optional<Class<?>> findInterfaceByMethod(Set<Class<?>> interfaces, Method method) {
+    String methodName = method.getName();
+    Class<?>[] methodParameterTypes = method.getParameterTypes();
+    return interfaces.stream().filter(i -> doesMethodBelongToClass(i, methodName, methodParameterTypes)).findFirst();
+  }
+
+  private boolean doesMethodBelongToClass(Class<?> clazz, String methodName, Class<?>[] methodParameterTypes) {
+    return Stream.of(clazz.getMethods()).anyMatch(
+      m -> m.getName().equals(methodName) && Arrays.equals(m.getParameterTypes(), methodParameterTypes)
+    );
+  }
+
   static Set<Class<?>> extractInterfaces(final Object object) {
     return extractInterfaces(object.getClass());
   }

diff --git a/java/src/org/openqa/selenium/support/events/EventFiringDecorator.java b/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());
 

diff --git a/java/test/org/openqa/selenium/remote/AugmenterTest.java b/java/test/org/openqa/selenium/remote/AugmenterTest.java
@@ -429,7 +429,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 {

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedAlertTest.java b/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();
     }
   }

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedNavigationTest.java b/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();
     }
   }

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedOptionsTest.java b/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();
     }
   }

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedRemoteWebDriverTest.java b/java/test/org/openqa/selenium/support/decorators/DecoratedRemoteWebDriverTest.java
@@ -49,18 +49,19 @@ public 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
   public void cannotConvertDecoratedToRemoteWebDriver() {
     RemoteWebDriver originalDriver = mock(RemoteWebDriver.class);
 
-    WebDriver decorated = new WebDriverDecorator().decorate(originalDriver);
+    WebDriver decorated = new WebDriverDecorator<>().decorate(originalDriver);
 
     assertThat(decorated).isNotInstanceOf(RemoteWebDriver.class);
   }
@@ -69,7 +70,7 @@ public void cannotConvertDecoratedToRemoteWebDriver() {
   public void decoratedDriversShouldImplementWrapsDriver() {
     RemoteWebDriver originalDriver = mock(RemoteWebDriver.class);
 
-    WebDriver decorated = new WebDriverDecorator().decorate(originalDriver);
+    WebDriver decorated = new WebDriverDecorator<>().decorate(originalDriver);
 
     assertThat(decorated).isInstanceOf(WrapsDriver.class);
   }
@@ -84,7 +85,7 @@ public 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);
@@ -100,7 +101,7 @@ public 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"));
 

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedSwitchToTest.java b/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();
     }
   }

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedTimeoutsTest.java b/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();
     }
   }

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedVirtualAuthenticatorTest.java b/java/test/org/openqa/selenium/support/decorators/DecoratedVirtualAuthenticatorTest.java
@@ -53,7 +53,7 @@ public Fixture() {
       originalDriver = mock(
         WebDriver.class, withSettings().extraInterfaces(HasVirtualAuthenticator.class));
       when(((HasVirtualAuthenticator) originalDriver).addVirtualAuthenticator(any())).thenReturn(original);
-      decoratedDriver = new WebDriverDecorator().decorate(originalDriver);
+      decoratedDriver = new WebDriverDecorator<>().decorate(originalDriver);
       decorated = ((HasVirtualAuthenticator) decoratedDriver)
         .addVirtualAuthenticator(new VirtualAuthenticatorOptions());
     }

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedWebDriverTest.java b/java/test/org/openqa/selenium/support/decorators/DecoratedWebDriverTest.java
@@ -62,7 +62,7 @@ public Fixture() {
         .extraInterfaces(JavascriptExecutor.class, TakesScreenshot.class,
                          Interactive.class, HasVirtualAuthenticator.class));
       originalAuth = mock(VirtualAuthenticator.class);
-      decorated = new WebDriverDecorator().decorate(original);
+      decorated = new WebDriverDecorator<>().decorate(original);
       when(((HasVirtualAuthenticator) original).addVirtualAuthenticator(any())).thenReturn(originalAuth);
     }
   }
@@ -85,9 +85,9 @@ public 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);
 
@@ -100,7 +100,7 @@ public void canCompareDecorated() {
   @Test
   public 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());
   }
 

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedWebElementTest.java b/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"));
     }
   }

diff --git a/java/test/org/openqa/selenium/support/decorators/DecoratedWindowTest.java b/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();
     }
   }

diff --git a/java/test/org/openqa/selenium/support/decorators/IntegrationTest.java b/java/test/org/openqa/selenium/support/decorators/IntegrationTest.java
@@ -35,7 +35,7 @@
 @Category(UnitTests.class)
 public class IntegrationTest {
 
-  static class CountCalls extends WebDriverDecorator {
+  static class CountCalls extends WebDriverDecorator<WebDriver> {
 
     int counterBefore = 0;
     int counterAfter = 0;

diff --git a/java/test/org/openqa/selenium/support/decorators/InterfacesTest.java b/java/test/org/openqa/selenium/support/decorators/InterfacesTest.java
@@ -37,7 +37,7 @@ public 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);
   }
 
@@ -46,7 +46,7 @@ public 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);
   }
 }