diff --git a/java/src/org/openqa/selenium/support/decorators/WebDriverDecorator.java b/java/src/org/openqa/selenium/support/decorators/WebDriverDecorator.java index f500774138277..7ee1156b9c908 100644 --- a/java/src/org/openqa/selenium/support/decorators/WebDriverDecorator.java +++ b/java/src/org/openqa/selenium/support/decorators/WebDriverDecorator.java @@ -195,7 +195,15 @@ protected static class Definition { public Definition(Decorated decorated) { this.decoratedClass = decorated.getClass(); - this.originalClass = decorated.getOriginal().getClass(); + this.originalClass = unwrapOriginal(decorated).getClass(); + } + + private static Object unwrapOriginal(Decorated decorated) { + Object original = decorated.getOriginal(); + while (original instanceof WrapsDriver) { + original = ((WrapsDriver) original).getWrappedDriver(); + } + return original; } @Override @@ -480,12 +488,12 @@ private static void extractInterfaces(final Set> collector, final Class private Map, Function> deriveAdditionalInterfaces(Z sample) { Map, Function> 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())) { + if (sample instanceof WrapsDriver || "getWrappedDriver".equals(method.getName())) { return instance; } throw new UnsupportedOperationException(method.getName());