diff --git a/presto-main/src/test/java/com/facebook/presto/sql/planner/TestLogicalPlanner.java b/presto-main/src/test/java/com/facebook/presto/sql/planner/TestLogicalPlanner.java
index dc3cf505d411..4274d91d95b9 100644
--- a/presto-main/src/test/java/com/facebook/presto/sql/planner/TestLogicalPlanner.java
+++ b/presto-main/src/test/java/com/facebook/presto/sql/planner/TestLogicalPlanner.java
@@ -361,9 +361,8 @@ public void testCorrelatedSubqueries()
     public void testStreamingAggregationForCorrelatedSubquery()
     {
         // Use equi-clause to trigger hash partitioning of the join sources
-        assertPlanWithSession("SELECT name, (SELECT max(name) FROM region WHERE regionkey = nation.regionkey AND length(name) > length(nation.name)) FROM nation",
-                this.getQueryRunner().getDefaultSession(),
-                false,
+        assertDistributedPlan(
+                "SELECT name, (SELECT max(name) FROM region WHERE regionkey = nation.regionkey AND length(name) > length(nation.name)) FROM nation",
                 anyTree(
                         aggregation(
                                 ImmutableList.of(ImmutableList.of("n_name", "n_regionkey", "unique")),
@@ -381,9 +380,8 @@ public void testStreamingAggregationForCorrelatedSubquery()
                                                 tableScan("region", ImmutableMap.of("r_name", "name")))))));
 
         // Don't use equi-clauses to trigger replicated join
-        assertPlanWithSession("SELECT name, (SELECT max(name) FROM region WHERE regionkey > nation.regionkey) FROM nation",
-                this.getQueryRunner().getDefaultSession(),
-                false,
+        assertDistributedPlan(
+                "SELECT name, (SELECT max(name) FROM region WHERE regionkey > nation.regionkey) FROM nation",
                 anyTree(
                         aggregation(
                                 ImmutableList.of(ImmutableList.of("n_name", "n_regionkey", "unique")),
diff --git a/presto-main/src/test/java/com/facebook/presto/sql/planner/assertions/BasePlanTest.java b/presto-main/src/test/java/com/facebook/presto/sql/planner/assertions/BasePlanTest.java
index 24f4b9ea0186..f6dc8dfd8c72 100644
--- a/presto-main/src/test/java/com/facebook/presto/sql/planner/assertions/BasePlanTest.java
+++ b/presto-main/src/test/java/com/facebook/presto/sql/planner/assertions/BasePlanTest.java
@@ -138,6 +138,11 @@ protected void assertPlan(String sql, LogicalPlanner.Stage stage, PlanMatchPatte
         });
     }
 
+    protected void assertDistributedPlan(String sql, PlanMatchPattern pattern)
+    {
+        assertPlanWithSession(sql, getQueryRunner().getDefaultSession(), false, pattern);
+    }
+
     protected void assertMinimallyOptimizedPlan(@Language("SQL") String sql, PlanMatchPattern pattern)
     {
         List<PlanOptimizer> optimizers = ImmutableList.of(