Address comments

Author: pujagani

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

@@ -31,7 +31,12 @@
 import org.openqa.selenium.remote.RemoteWebElement;
 
 import java.lang.reflect.Method;
-import java.util.*;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.WeakHashMap;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 

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

@@ -26,9 +26,9 @@
 import org.slf4j.LoggerFactory;
 
 import java.lang.reflect.Method;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.Callable;
+import java.util.stream.Collectors;
 
 import static io.appium.java_client.proxy.MethodCallListener.UNSET;
 
@@ -115,21 +115,25 @@ public static Object intercept(
         }
 
         if (result instanceof RemoteWebElement) {
-            result = Helpers.wrapElement((RemoteWebElement) result, (HasMethodCallListeners) self, listeners);
+            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;
+
+            if (!originalList.isEmpty()) {
+                result = originalList.stream()
+                        .map(item -> {
+                            if (item instanceof RemoteWebElement) {
+                                return Helpers.wrapElement(
+                                        (RemoteWebElement) item,
+                                        (HasMethodCallListeners) self,
+                                        listeners);
+                            }
+                            return item;
+                        })
+                        .collect(Collectors.toList());
             }
         }