Add BasePlanTest#assertPlanWithSession that accepts planValidator
@@ -33,6 +33,7 @@ import org.testng.annotations.BeforeClass;
import java.util.List;
import java.util.Map;
+import java.util.function.Consumer;
import java.util.function.Predicate;
import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
@@ -154,6 +155,16 @@ public class BasePlanTest
});
}
+ protected void assertPlanWithSession(@Language("SQL") String sql, Session session, boolean forceSingleNode, PlanMatchPattern pattern, Consumer<Plan> planValidator)
+ {
+ queryRunner.inTransaction(session, transactionSession -> {
+ Plan actualPlan = queryRunner.createPlan(transactionSession, sql, LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED, forceSingleNode);
+ PlanAssert.assertPlan(transactionSession, queryRunner.getMetadata(), queryRunner.getStatsCalculator(), actualPlan, pattern);
+ planValidator.accept(actualPlan);
+ return null;
+ });
+ }
+
protected Plan plan(String sql)
{
return plan(sql, LogicalPlanner.Stage.OPTIMIZED_AND_VALIDATED);
