Wrapping result of executeScript in EventFiringWebDriver. Fixes #3380

Author: barancev

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

@@ -58,6 +58,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 /**
  * A wrapper around an arbitrary {@link WebDriver} instance which supports registering of a
@@ -216,7 +217,7 @@ public Object executeScript(String script, Object... args) {
       Object[] usedArgs = unpackWrappedArgs(args);
       Object result = ((JavascriptExecutor) driver).executeScript(script, usedArgs);
       dispatcher.afterScript(script, driver);
-      return result;
+      return wrapResult(result);
     }
     throw new UnsupportedOperationException(
         "Underlying driver instance does not support executing javascript");
@@ -265,6 +266,22 @@ private Object unpackWrappedElement(Object arg) {
     }
   }
 
+  private Object wrapResult(Object result) {
+    if (result instanceof WebElement) {
+      return new EventFiringWebElement((WebElement) result);
+    }
+    if (result instanceof List) {
+      return ((List) result).stream().map(this::wrapResult).collect(Collectors.toList());
+    }
+    if (result instanceof Map) {
+      return ((Map<String, Object>) result).entrySet().stream().collect(Collectors.toMap(
+          e -> e.getKey(), e -> wrapResult(e.getValue())
+      ));
+    }
+
+    return result;
+  }
+
   public <X> X getScreenshotAs(OutputType<X> target) throws WebDriverException {
     if (driver instanceof TakesScreenshot) {
       return ((TakesScreenshot) driver).getScreenshotAs(target);

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

@@ -17,8 +17,10 @@
 
 package org.openqa.selenium.support.events;
 
+import static org.hamcrest.CoreMatchers.*;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
@@ -32,6 +34,7 @@
 
 import com.google.common.collect.ImmutableMap;
 
+import org.hamcrest.CoreMatchers;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -254,6 +257,24 @@ public void shouldUnpackElementArgsWhenCallingScripts() {
     verify((JavascriptExecutor) mockedDriver).executeScript("foo", element);
   }
 
+  @Test
+  public void shouldWrapElementFoundWhenCallingScripts() {
+    final WebDriver mockedDriver = mock(WebDriver.class,
+                                        withSettings().extraInterfaces(JavascriptExecutor.class));
+    final WebElement stubbedElement = mock(WebElement.class);
+
+    when(((JavascriptExecutor) mockedDriver).executeScript("foo"))
+        .thenReturn(stubbedElement);
+
+    EventFiringWebDriver testedDriver = new EventFiringWebDriver(mockedDriver);
+
+    Object res = testedDriver.executeScript("foo");
+    verify((JavascriptExecutor) mockedDriver).executeScript("foo");
+    assertThat(res, instanceOf(WebElement.class));
+    assertThat(res, instanceOf(WrapsElement.class));
+    assertSame(stubbedElement, ((WrapsElement) res).getWrappedElement());
+  }
+
   @Test
   public void testShouldUnpackListOfElementArgsWhenCallingScripts() {
     final WebDriver mockedDriver = mock(WebDriver.class,
@@ -276,6 +297,27 @@ public void testShouldUnpackListOfElementArgsWhenCallingScripts() {
     verify((JavascriptExecutor) mockedDriver).executeScript("foo", args);
   }
 
+  @Test
+  public void shouldWrapMultipleElementsFoundWhenCallingScripts() {
+    final WebDriver mockedDriver = mock(WebDriver.class,
+                                        withSettings().extraInterfaces(JavascriptExecutor.class));
+    final WebElement stubbedElement1 = mock(WebElement.class);
+    final WebElement stubbedElement2 = mock(WebElement.class);
+
+    when(((JavascriptExecutor) mockedDriver).executeScript("foo"))
+        .thenReturn(Arrays.asList(stubbedElement1, stubbedElement2));
+
+    EventFiringWebDriver testedDriver = new EventFiringWebDriver(mockedDriver);
+
+    Object res = testedDriver.executeScript("foo");
+    verify((JavascriptExecutor) mockedDriver).executeScript("foo");
+    assertThat(res, instanceOf(List.class));
+    List<Object> resList = (List<Object>) res;
+    resList.stream().forEach(el -> assertTrue(el instanceof WrapsElement));
+    assertSame(stubbedElement1, ((WrapsElement) resList.get(0)).getWrappedElement());
+    assertSame(stubbedElement2, ((WrapsElement) resList.get(1)).getWrappedElement());
+  }
+
   @Test
   public void testShouldUnpackMapOfElementArgsWhenCallingScripts() {
     final WebDriver mockedDriver = mock(WebDriver.class,