diff --git a/presto-main/src/main/java/io/prestosql/testing/DataProviders.java b/presto-main/src/main/java/io/prestosql/testing/DataProviders.java index 472b8e514b72..00e4dd3ec61a 100644 --- a/presto-main/src/main/java/io/prestosql/testing/DataProviders.java +++ b/presto-main/src/main/java/io/prestosql/testing/DataProviders.java @@ -15,23 +15,23 @@ import org.testng.annotations.DataProvider; -import java.util.ArrayList; import java.util.stream.Collector; +import static java.util.stream.Collectors.collectingAndThen; +import static java.util.stream.Collectors.mapping; +import static java.util.stream.Collectors.toList; + public final class DataProviders { private DataProviders() {} public static Collector toDataProvider() { - return Collector.of( - ArrayList::new, - (builder, entry) -> builder.add(new Object[] {entry}), - (left, right) -> { - left.addAll(right); - return left; - }, - builder -> builder.toArray(new Object[][] {})); + return collectingAndThen( + mapping( + value -> new Object[] {value}, + toList()), + list -> list.toArray(new Object[][] {})); } @DataProvider