#16573 unwrap double-wrapped webdriver

Author: asolntsev

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

@@ -195,7 +195,7 @@ protected static class Definition {
 
     public Definition(Decorated<?> decorated) {
       this.decoratedClass = decorated.getClass();
-      this.originalClass = decorated.getOriginal().getClass();
+      this.originalClass = unwrapOriginal(decorated).getClass();
     }
 
     @Override
@@ -217,6 +217,17 @@ public int hashCode() {
     }
   }
 
+  private static Object unwrapOriginal(Decorated<?> decorated) {
+    return unwrapWebdriver(decorated.getOriginal());
+  }
+
+  private static Object unwrapWebdriver(Object webDriver) {
+    if (webDriver instanceof WebDriver && webDriver instanceof WrapsDriver) {
+      return unwrapWebdriver(((WrapsDriver) webDriver).getWrappedDriver());
+    }
+    return webDriver;
+  }
+
   public interface HasTarget<Z> {
     Decorated<Z> getTarget();
 
@@ -480,12 +491,15 @@ private static void extractInterfaces(final Set<Class<?>> collector, final Class
   private <Z> Map<Class<?>, Function<Z, InvocationHandler>> deriveAdditionalInterfaces(Z sample) {
     Map<Class<?>, Function<Z, InvocationHandler>> handlers = new HashMap<>();
 
-    if (sample instanceof WebDriver && !(sample instanceof WrapsDriver)) {
+    if (sample instanceof WebDriver) {
       handlers.put(
           WrapsDriver.class,
           (instance) ->
               (proxy, method, args) -> {
                 if ("getWrappedDriver".equals(method.getName())) {
+                  return unwrapWebdriver(instance);
+                }
+                if (sample instanceof WrapsDriver) {
                   return instance;
                 }
                 throw new UnsupportedOperationException(method.getName());

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

@@ -55,6 +55,28 @@ void shouldImplementWrapsDriverToProvideAccessToUnderlyingDriver() {
     RemoteWebDriver underlying =
         (RemoteWebDriver) ((WrapsDriver) decoratedDriver).getWrappedDriver();
     assertThat(underlying.getSessionId()).isEqualTo(sessionId);
+    assertThat(underlying).isSameAs(originalDriver);
+  }
+
+  @Test
+  void canWrapDriverMultipleTimes() {
+    SessionId sessionId = new SessionId(UUID.randomUUID());
+    RemoteWebDriver originalDriver = mock(RemoteWebDriver.class);
+    when(originalDriver.getSessionId()).thenReturn(sessionId);
+
+    RemoteWebDriver decorated1 =
+        new WebDriverDecorator<>(RemoteWebDriver.class).decorate(originalDriver);
+    RemoteWebDriver decorated2 =
+        new WebDriverDecorator<>(RemoteWebDriver.class).decorate(decorated1);
+    RemoteWebDriver decoratedDriver =
+        new WebDriverDecorator<>(RemoteWebDriver.class).decorate(decorated2);
+
+    assertThat(decoratedDriver.getSessionId()).isEqualTo(sessionId);
+
+    RemoteWebDriver underlying =
+        (RemoteWebDriver) ((WrapsDriver) decoratedDriver).getWrappedDriver();
+    assertThat(underlying.getSessionId()).isEqualTo(sessionId);
+    assertThat(underlying).isSameAs(originalDriver);
   }
 
   @Test