Skip to content
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -56,26 +56,24 @@ public Object call(Object obj, Method method, Object[] args, Callable<?> origina
Object result = original.call();

if (result instanceof RemoteWebElement) {
result = wrapElement(
return wrapElement(
(RemoteWebElement) result,
parent,
listeners);
} else if (result instanceof List) {
List<?> originalList = (List<?>) result;

if (!originalList.isEmpty()) {
result = originalList.stream()
.map(item -> {
if (item instanceof RemoteWebElement) {
return wrapElement(
(RemoteWebElement) item,
parent,
listeners);
}
return item;
})
.collect(Collectors.toList());
}
}

if (result instanceof List) {
return ((List<?>) result).stream()
.map(item -> {
if (item instanceof RemoteWebElement) {
return wrapElement(
(RemoteWebElement) item,
parent,
listeners);
}
return item;
})
.collect(Collectors.toList());
}

return result;
Expand Down
Loading