[java][cdp] Use devtools script pinning mechanism by default (#11622)

This helps achieve feature parity in script pinning among different Selenium language bindings. Co-authored-by: Diego Molina <diemol@users.noreply.github.com>
SeleniumHQ · Aug 11, 2023 · a4beba5 · a4beba5
1 parent 3eab322
commit a4beba5
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 3 deletions.
diff --git a/java/src/org/openqa/selenium/UnpinnedScriptKey.java b/java/src/org/openqa/selenium/UnpinnedScriptKey.java
@@ -21,13 +21,16 @@
 import java.util.HashSet;
 import java.util.Objects;
 import java.util.Set;
+import java.util.UUID;
 import java.util.WeakHashMap;
 
-class UnpinnedScriptKey extends ScriptKey {
+public class UnpinnedScriptKey extends ScriptKey {
 
   private static final WeakHashMap<JavascriptExecutor, Set<UnpinnedScriptKey>> pinnedScripts =
       new WeakHashMap<>();
   private final String script;
+  private String scriptId;
+  private final String scriptHandle;
 
   static UnpinnedScriptKey pin(JavascriptExecutor executor, String script) {
     UnpinnedScriptKey toReturn = new UnpinnedScriptKey(script);
@@ -51,16 +54,41 @@ static Set<UnpinnedScriptKey> getPinnedScripts(JavascriptExecutor executor) {
     return Collections.unmodifiableSet(toReturn);
   }
 
-  private UnpinnedScriptKey(String script) {
+  public UnpinnedScriptKey(String script) {
     super(script);
 
+    this.scriptHandle = UUID.randomUUID().toString().replace("-", "");
     this.script = script;
   }
 
-  String getScript() {
+  public void setScriptId(String id) {
+    this.scriptId = id;
+  }
+
+  public String getScriptId() {
+    return this.scriptId;
+  }
+
+  public String getScript() {
     return script;
   }
 
+  public String getScriptHandle() {
+    return scriptHandle;
+  }
+
+  public String creationScript() {
+    return String.format("function __webdriver_%s(arguments) { %s }", this.scriptHandle, this.script);
+  }
+
+  public String executionScript() {
+    return String.format("return __webdriver_%s(arguments)", this.scriptHandle);
+  }
+
+  public String removalScript() {
+    return String.format("__webdriver_%s = undefined", this.scriptHandle);
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

diff --git a/java/src/org/openqa/selenium/remote/RemoteWebDriver.java b/java/src/org/openqa/selenium/remote/RemoteWebDriver.java
@@ -52,6 +52,7 @@
 import org.openqa.selenium.Dimension;
 import org.openqa.selenium.HasCapabilities;
 import org.openqa.selenium.ImmutableCapabilities;
+import org.openqa.selenium.JavascriptException;
 import org.openqa.selenium.JavascriptExecutor;
 import org.openqa.selenium.MutableCapabilities;
 import org.openqa.selenium.NoAlertPresentException;
@@ -62,9 +63,11 @@
 import org.openqa.selenium.Platform;
 import org.openqa.selenium.Point;
 import org.openqa.selenium.PrintsPage;
+import org.openqa.selenium.ScriptKey;
 import org.openqa.selenium.SearchContext;
 import org.openqa.selenium.SessionNotCreatedException;
 import org.openqa.selenium.TakesScreenshot;
+import org.openqa.selenium.UnpinnedScriptKey;
 import org.openqa.selenium.WebDriver;
 import org.openqa.selenium.WebDriverException;
 import org.openqa.selenium.WebElement;
@@ -79,6 +82,7 @@
 import org.openqa.selenium.interactions.Sequence;
 import org.openqa.selenium.internal.Debug;
 import org.openqa.selenium.internal.Require;
+import org.openqa.selenium.json.TypeToken;
 import org.openqa.selenium.logging.LocalLogs;
 import org.openqa.selenium.logging.LoggingHandler;
 import org.openqa.selenium.logging.Logs;
@@ -471,6 +475,76 @@ public Object executeAsyncScript(String script, Object... args) {
     return execute(DriverCommand.EXECUTE_ASYNC_SCRIPT(script, convertedArgs)).getValue();
   }
 
+  @Override
+  public ScriptKey pin(String script) {
+    UnpinnedScriptKey key = (UnpinnedScriptKey) JavascriptExecutor.super.pin(script);
+    String browserName = getCapabilities().getBrowserName().toLowerCase();
+    if ((browserName.equals("chrome") ||
+         browserName.equals("msedge") ||
+         browserName.equals("microsoftedge")) && this instanceof HasDevTools) {
+
+      ((HasDevTools) this).maybeGetDevTools().ifPresent(devTools -> {
+        devTools.createSessionIfThereIsNotOne();
+        devTools.send(new org.openqa.selenium.devtools.Command<>("Page.enable",
+                                                                 ImmutableMap.of()));
+        devTools.send(new org.openqa.selenium.devtools.Command<>("Runtime.evaluate",
+                                                                 ImmutableMap.of("expression",
+                                                                                 key.creationScript())));
+        Map<String, Object> result = devTools.send(new org.openqa.selenium.devtools.Command<>(
+          "Page.addScriptToEvaluateOnNewDocument",
+          ImmutableMap.of("source", key.creationScript()),
+          new TypeToken<Map<String, Object>>() {
+          }.getType()));
+        key.setScriptId((String) result.get("identifier"));
+      });
+    }
+    return key;
+  }
+
+  @Override
+  public void unpin(ScriptKey scriptKey) {
+    UnpinnedScriptKey key = (UnpinnedScriptKey) scriptKey;
+
+    JavascriptExecutor.super.unpin(key);
+
+    String browserName = getCapabilities().getBrowserName().toLowerCase();
+    if ((browserName.equals("chrome") ||
+         browserName.equals("msedge") ||
+         browserName.equals("microsoftedge")) && this instanceof HasDevTools) {
+      ((HasDevTools) this).maybeGetDevTools().ifPresent(devTools -> {
+        devTools.send(new org.openqa.selenium.devtools.Command<>("Page.enable",
+                                                                 ImmutableMap.of()));
+        devTools.send(new org.openqa.selenium.devtools.Command<>(
+          "Runtime.evaluate",
+          ImmutableMap.of("expression", key.removalScript())));
+        devTools.send(new org.openqa.selenium.devtools.Command<>(
+          "Page.removeScriptToEvaluateOnLoad",
+          ImmutableMap.of("identifier", key.getScriptId())));
+      });
+    }
+  }
+
+  @Override
+  public Object executeScript(ScriptKey key, Object... args) {
+    Require.stateCondition(
+      key instanceof UnpinnedScriptKey,
+      "Script key should have been generated by this driver");
+
+    if (!getPinnedScripts().contains(key)) {
+      throw new JavascriptException("Script is unpinned");
+    }
+
+    String browserName = getCapabilities().getBrowserName().toLowerCase();
+
+    if ((browserName.equals("chrome") ||
+         browserName.equals("msedge") ||
+         browserName.equals("microsoftedge")) && this instanceof HasDevTools) {
+      return executeScript(((UnpinnedScriptKey) key).executionScript(), args);
+    }
+
+    return executeScript(((UnpinnedScriptKey) key).getScript(), args);
+  }
+
   @Override
   public TargetLocator switchTo() {
     return new RemoteTargetLocator();