fix: Proxy commands issues via RemoteWebElement

Fixes #2239

Author: pujagani

src/main/java/io/appium/java_client/proxy/Helpers.java

@@ -27,13 +27,11 @@
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
 import org.jspecify.annotations.Nullable;
+import org.openqa.selenium.remote.RemoteWebDriver;
+import org.openqa.selenium.remote.RemoteWebElement;
 
 import java.lang.reflect.Method;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.Map;
-import java.util.Set;
-import java.util.WeakHashMap;
+import java.util.*;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -222,4 +220,28 @@ private static class ProxyClassSignature {
         Class<?>[] constructorArgTypes;
         ElementMatcher<MethodDescription> extraMethodMatcher;
     }
+
+    public static RemoteWebElement wrapElement(
+            RemoteWebElement original,
+            HasMethodCallListeners parent,
+            MethodCallListener[] listeners
+    ) {
+        RemoteWebElement proxy = createProxy(
+                RemoteWebElement.class,
+                new Object[]{},
+                new Class[]{},
+                List.of(listeners),
+                ElementMatchers.not(
+                        namedOneOf(
+                        OBJECT_METHOD_NAMES.toArray(new String[0]))
+                                .or(ElementMatchers.named("setId").or(ElementMatchers.named("setParent")))
+                )
+        );
+
+        proxy.setId(original.getId());
+
+        proxy.setParent((RemoteWebDriver) parent);
+
+        return proxy;
+    }
 }

src/main/java/io/appium/java_client/proxy/Interceptor.java

@@ -21,10 +21,13 @@
 import net.bytebuddy.implementation.bind.annotation.RuntimeType;
 import net.bytebuddy.implementation.bind.annotation.SuperCall;
 import net.bytebuddy.implementation.bind.annotation.This;
+import org.openqa.selenium.remote.RemoteWebElement;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.Callable;
 
 import static io.appium.java_client.proxy.MethodCallListener.UNSET;
@@ -111,6 +114,25 @@ public static Object intercept(
             }
         }
 
+        if (result instanceof RemoteWebElement) {
+            result = Helpers.wrapElement((RemoteWebElement) result, (HasMethodCallListeners) self, listeners);
+        } else if (result instanceof List) {
+            List<?> originalList = (List<?>) result;
+            if (!originalList.isEmpty() && originalList.get(0) instanceof RemoteWebElement) {
+                List<Object> wrappedList = new ArrayList<>(originalList.size());
+                for (Object item : originalList) {
+                    if (item instanceof RemoteWebElement) {
+                        wrappedList.add(Helpers.wrapElement(
+                                (RemoteWebElement) item,
+                                (HasMethodCallListeners) self, listeners));
+                    } else {
+                        wrappedList.add(item);
+                    }
+                }
+                result = wrappedList;
+            }
+        }
+
         final Object endResult = result == UNSET ? null : result;
         for (var listener : listeners) {
             try {

src/test/java/io/appium/java_client/proxy/ProxyHelpersTest.java

@@ -19,14 +19,20 @@
 import io.appium.java_client.ios.IOSDriver;
 import io.appium.java_client.ios.options.XCUITestOptions;
 import org.junit.jupiter.api.Test;
+import org.openqa.selenium.By;
 import org.openqa.selenium.Capabilities;
+import org.openqa.selenium.NoSuchSessionException;
+import org.openqa.selenium.WebElement;
 import org.openqa.selenium.remote.RemoteWebDriver;
+import org.openqa.selenium.remote.RemoteWebElement;
 import org.openqa.selenium.remote.UnreachableBrowserException;
 
 import java.lang.reflect.Method;
 import java.net.MalformedURLException;
 import java.net.URL;
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.Callable;
 
 import static io.appium.java_client.proxy.Helpers.createProxy;
@@ -45,6 +51,31 @@ public FakeIOSDriver(URL url, Capabilities caps) {
         @Override
         protected void startSession(Capabilities capabilities) {
         }
+
+        @Override
+        public WebElement findElement(By locator) {
+            RemoteWebElement webElement = new RemoteWebElement();
+            webElement.setId(locator.toString());
+            webElement.setParent(this);
+            return webElement;
+        }
+
+        @Override
+        public List<WebElement> findElements(By locator) {
+            List<WebElement> webElements = new ArrayList<>();
+
+            RemoteWebElement webElement1 = new RemoteWebElement();
+            webElement1.setId("1234");
+            webElement1.setParent(this);
+            webElements.add(webElement1);
+
+            RemoteWebElement webElement2 = new RemoteWebElement();
+            webElement2.setId("5678");
+            webElement2.setParent(this);
+            webElements.add(webElement2);
+
+            return webElements;
+        }
     }
 
     @Test
@@ -133,4 +164,52 @@ public Object onError(Object obj, Method method, Object[] args, Throwable e) thr
                         "onError get")
         )));
     }
+
+
+    @Test
+    void shouldFireEventsForRemoteWebElement() throws MalformedURLException {
+        final StringBuilder acc = new StringBuilder();
+        MethodCallListener listener = new MethodCallListener() {
+            @Override
+            public void beforeCall(Object target, Method method, Object[] args) {
+                acc.append("beforeCall ").append(method.getName()).append("\n");
+            }
+        };
+
+        FakeIOSDriver driver = createProxy(
+                FakeIOSDriver.class,
+                new Object[] {new URL("http://localhost:4723/"), new XCUITestOptions()},
+                new Class[] {URL.class, Capabilities.class},
+                listener
+        );
+
+        WebElement element = driver.findElement(By.id("button"));
+
+        assertThrows(
+                NoSuchSessionException.class,
+                element::click
+        );
+
+        List<WebElement> elements = driver.findElements(By.id("button"));
+
+        assertThrows(
+                NoSuchSessionException.class,
+                () -> elements.get(1).isSelected()
+        );
+
+        assertThat(acc.toString().trim(), is(equalTo(
+                String.join("\n",
+                        "beforeCall findElement",
+                        "beforeCall click",
+                        "beforeCall getSessionId",
+                        "beforeCall getCapabilities",
+                        "beforeCall getCapabilities",
+                        "beforeCall findElements",
+                        "beforeCall isSelected",
+                        "beforeCall getSessionId",
+                        "beforeCall getCapabilities",
+                        "beforeCall getCapabilities"
+                )
+        )));
+    }
 }