add connectToExisting and withWebSocketPort methods

Author: navin772

java/src/org/openqa/selenium/firefox/GeckoDriverService.java

@@ -147,6 +147,8 @@ public static class Builder
     private @Nullable FirefoxDriverLogLevel logLevel;
     private @Nullable Boolean logTruncate;
     private @Nullable File profileRoot;
+    private @Nullable Integer marionettePort;
+    private @Nullable Integer websocketPort;
 
     @Override
     public int score(Capabilities capabilities) {
@@ -204,6 +206,31 @@ public GeckoDriverService.Builder withProfileRoot(@Nullable File root) {
       return this;
     }
 
+    /**
+     * Configures geckodriver to connect to an existing Firefox instance via the specified
+     * Marionette port.
+     *
+     * @param marionettePort The port where Marionette is listening on the existing Firefox
+     *     instance.
+     * @return A self reference.
+     */
+    public GeckoDriverService.Builder connectToExisting(int marionettePort) {
+      this.marionettePort = marionettePort;
+      return this;
+    }
+
+    /**
+     * Configures the WebSocket port for BiDi. A value of 0 will automatically allocate a free port.
+     *
+     * @param websocketPort The port to use for WebSocket communication, or 0 for automatic
+     *     allocation.
+     * @return A self reference.
+     */
+    public GeckoDriverService.Builder withWebSocketPort(@Nullable Integer websocketPort) {
+      this.websocketPort = websocketPort;
+      return this;
+    }
+
     @Override
     protected void loadSystemProperties() {
       parseLogOutput(GECKO_DRIVER_LOG_PROPERTY);
@@ -229,27 +256,28 @@ protected List<String> createArgs() {
       List<String> args = new ArrayList<>();
       args.add(String.format(Locale.ROOT, "--port=%d", getPort()));
 
-      // Check if we're connecting to an existing Firefox instance
-      boolean connectExisting = false;
-      for (String arg : args) {
-        if (arg.contains("--connect-existing")) {
-          connectExisting = true;
-          break;
+      // Check if marionette port is specified via connectToExisting method
+      if (marionettePort != null) {
+        args.add("--connect-existing");
+        args.add("--marionette-port");
+        args.add(String.valueOf(marionettePort));
+      } else {
+        // Configure websocket port for BiDi communication
+        if (websocketPort != null) {
+          args.add("--websocket-port");
+          args.add(String.valueOf(websocketPort));
+
+          args.add("--allow-origins");
+          args.add(String.format("http://127.0.0.1:%d", websocketPort));
+          args.add(String.format("http://localhost:%d", websocketPort));
+          args.add(String.format("http://[::1]:%d", websocketPort));
+        } else {
+          // Use 0 to auto-allocate a free port
+          args.add("--websocket-port");
+          args.add("0");
         }
       }
 
-      // Only allocate a free port for the websocket when not connecting to an existing instance
-      // This avoids conflicts when multiple Firefox instances are started
-      if (!connectExisting) {
-        int wsPort = PortProber.findFreePort();
-        args.add(String.format("--websocket-port=%d", wsPort));
-
-        args.add("--allow-origins");
-        args.add(String.format("http://127.0.0.1:%d", wsPort));
-        args.add(String.format("http://localhost:%d", wsPort));
-        args.add(String.format("http://[::1]:%d", wsPort));
-      }
-
       if (logLevel != null) {
         args.add("--log");
         args.add(logLevel.toString());

java/test/org/openqa/selenium/firefox/FirefoxDriverConcurrentTest.java

@@ -168,7 +168,7 @@ void shouldBeAbleToUseTheSameProfileMoreThanOnce() {
 
   @Test
   void multipleFirefoxInstancesWithBiDiEnabledCanRunSimultaneously() {
-    // Create two Firefox instances with BiDi enabled
+    // Create two Firefox instances with BiDi enabled, should use different ports
     FirefoxOptions options1 = new FirefoxOptions().enableBiDi();
     FirefoxOptions options2 = new FirefoxOptions().enableBiDi();
 
@@ -206,4 +206,28 @@ void multipleFirefoxInstancesWithBiDiEnabledCanRunSimultaneously() {
       }
     }
   }
+
+  @Test
+  void geckoDriverServiceConnectToExistingFirefox() {
+    GeckoDriverService.Builder builder = new GeckoDriverService.Builder();
+
+    // Test connectToExisting method
+    builder.connectToExisting(2829);
+    GeckoDriverService service = builder.build();
+
+    assertThat(service).isNotNull();
+    service.stop();
+  }
+
+  @Test
+  void geckoDriverServiceCustomWebSocketPort() {
+    GeckoDriverService.Builder builder = new GeckoDriverService.Builder();
+
+    // Test withWebSocketPort method
+    builder.withWebSocketPort(9225);
+    GeckoDriverService service = builder.build();
+
+    assertThat(service).isNotNull();
+    service.stop();
+  }
 }