From 5dd385c9522a0ece8f8ff58d96382d949176d89f Mon Sep 17 00:00:00 2001
From: Puja Jagani <puja.jagani93@gmail.com>
Date: Thu, 25 Jul 2024 18:37:46 +0530
Subject: [PATCH] [bidi] [java] Add script pinning methods (#14305)

---
 .../openqa/selenium/remote/RemoteScript.java  | 10 +++++
 .../org/openqa/selenium/remote/Script.java    |  4 ++
 .../org/openqa/selenium/WebScriptTest.java    | 43 +++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/java/src/org/openqa/selenium/remote/RemoteScript.java b/java/src/org/openqa/selenium/remote/RemoteScript.java
index b070f71ada309..df47192e87ace 100644
--- a/java/src/org/openqa/selenium/remote/RemoteScript.java
+++ b/java/src/org/openqa/selenium/remote/RemoteScript.java
@@ -121,4 +121,14 @@ public long addDomMutationHandler(Consumer<DomMutation> consumer) {
   public void removeDomMutationHandler(long id) {
     this.biDi.removeListener(id);
   }
+
+  @Override
+  public String pin(String script) {
+    return this.script.addPreloadScript(script);
+  }
+
+  @Override
+  public void unpin(String id) {
+    this.script.removePreloadScript(id);
+  }
 }
diff --git a/java/src/org/openqa/selenium/remote/Script.java b/java/src/org/openqa/selenium/remote/Script.java
index e02527010a076..e8ea6d154d128 100644
--- a/java/src/org/openqa/selenium/remote/Script.java
+++ b/java/src/org/openqa/selenium/remote/Script.java
@@ -36,4 +36,8 @@ public interface Script {
   long addDomMutationHandler(Consumer<DomMutation> event);
 
   void removeDomMutationHandler(long id);
+
+  String pin(String script);
+
+  void unpin(String id);
 }
diff --git a/java/test/org/openqa/selenium/WebScriptTest.java b/java/test/org/openqa/selenium/WebScriptTest.java
index c05d630ffb9f4..bf7bcae431b66 100644
--- a/java/test/org/openqa/selenium/WebScriptTest.java
+++ b/java/test/org/openqa/selenium/WebScriptTest.java
@@ -245,4 +245,47 @@ void canRemoveDomMutationHandler() throws InterruptedException {
 
     Assertions.assertThat(latch.await(10, SECONDS)).isFalse();
   }
+
+  @Test
+  void canPinScript() throws ExecutionException, InterruptedException, TimeoutException {
+    CompletableFuture<ConsoleLogEntry> future = new CompletableFuture<>();
+
+    ((RemoteWebDriver) driver).script().pin("() => { console.log('Hello!'); }");
+
+    long id = ((RemoteWebDriver) driver).script().addConsoleMessageHandler(future::complete);
+
+    page = server.whereIs("/bidi/logEntryAdded.html");
+    driver.get(page);
+
+    ConsoleLogEntry logEntry = future.get(5, TimeUnit.SECONDS);
+
+    assertThat(logEntry.getText()).isEqualTo("Hello!");
+
+    ((RemoteWebDriver) driver).script().removeConsoleMessageHandler(id);
+  }
+
+  @Test
+  void canUnpinScript() throws ExecutionException, InterruptedException, TimeoutException {
+    CountDownLatch latch = new CountDownLatch(2);
+
+    String pinnedScript =
+        ((RemoteWebDriver) driver).script().pin("() => { console.log('Hello!'); }");
+
+    long id =
+        ((RemoteWebDriver) driver)
+            .script()
+            .addConsoleMessageHandler(consoleLogEntry -> latch.countDown());
+
+    page = server.whereIs("/bidi/logEntryAdded.html");
+
+    driver.get(page);
+
+    ((RemoteWebDriver) driver).script().unpin(pinnedScript);
+
+    driver.get(page);
+
+    assertThat(latch.getCount()).isEqualTo(1L);
+
+    ((RemoteWebDriver) driver).script().removeConsoleMessageHandler(id);
+  }
 }