diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java index 378c23540e4c2..69205b0f20061 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java @@ -148,6 +148,11 @@ protected void assertQuery(Session session, @Language("SQL") String actual, @Lan QueryAssertions.assertQuery(queryRunner, session, actual, expectedQueryRunner, expected, false, false); } + protected void assertQuery(Session actualSession, @Language("SQL") String actual, Session expectedSession, @Language("SQL") String expected) + { + QueryAssertions.assertQuery(queryRunner, actualSession, actual, expectedQueryRunner, expectedSession, expected, false, false); + } + protected void assertQuery(Session session, @Language("SQL") String sql, Consumer planAssertion) { checkArgument(queryRunner instanceof DistributedQueryRunner, "pattern assertion is only supported for DistributedQueryRunner"); diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java b/presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java index 8737594ac0c54..b5d9d5eb78198 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/QueryAssertions.java @@ -103,7 +103,20 @@ public static void assertQuery( boolean ensureOrdering, boolean compareUpdate) { - assertQuery(actualQueryRunner, session, actual, expectedQueryRunner, expected, ensureOrdering, compareUpdate, Optional.empty()); + assertQuery(actualQueryRunner, session, actual, expectedQueryRunner, session, expected, ensureOrdering, compareUpdate, Optional.empty()); + } + + public static void assertQuery( + QueryRunner actualQueryRunner, + Session actualSession, + @Language("SQL") String actual, + ExpectedQueryRunner expectedQueryRunner, + Session expectedSession, + @Language("SQL") String expected, + boolean ensureOrdering, + boolean compareUpdate) + { + assertQuery(actualQueryRunner, actualSession, actual, expectedQueryRunner, expectedSession, expected, ensureOrdering, compareUpdate, Optional.empty()); } public static void assertQuery( @@ -116,14 +129,29 @@ public static void assertQuery( boolean compareUpdate, Consumer planAssertion) { - assertQuery(actualQueryRunner, session, actual, expectedQueryRunner, expected, ensureOrdering, compareUpdate, Optional.of(planAssertion)); + assertQuery(actualQueryRunner, session, actual, expectedQueryRunner, session, expected, ensureOrdering, compareUpdate, Optional.of(planAssertion)); + } + + public static void assertQuery( + QueryRunner actualQueryRunner, + Session actualSession, + @Language("SQL") String actual, + ExpectedQueryRunner expectedQueryRunner, + Session expectedSession, + @Language("SQL") String expected, + boolean ensureOrdering, + boolean compareUpdate, + Consumer planAssertion) + { + assertQuery(actualQueryRunner, actualSession, actual, expectedQueryRunner, expectedSession, expected, ensureOrdering, compareUpdate, Optional.of(planAssertion)); } private static void assertQuery( QueryRunner actualQueryRunner, - Session session, + Session actualSession, @Language("SQL") String actual, ExpectedQueryRunner expectedQueryRunner, + Session expectedSession, @Language("SQL") String expected, boolean ensureOrdering, boolean compareUpdate, @@ -134,7 +162,7 @@ private static void assertQuery( Plan queryPlan = null; if (planAssertion.isPresent()) { try { - MaterializedResultWithPlan resultWithPlan = actualQueryRunner.executeWithPlan(session, actual, WarningCollector.NOOP); + MaterializedResultWithPlan resultWithPlan = actualQueryRunner.executeWithPlan(actualSession, actual, WarningCollector.NOOP); queryPlan = resultWithPlan.getQueryPlan(); actualResults = resultWithPlan.getMaterializedResult().toTestTypes(); } @@ -144,7 +172,7 @@ private static void assertQuery( } else { try { - actualResults = actualQueryRunner.execute(session, actual).toTestTypes(); + actualResults = actualQueryRunner.execute(actualSession, actual).toTestTypes(); } catch (RuntimeException ex) { fail("Execution of 'actual' query failed: " + actual, ex); @@ -158,7 +186,7 @@ private static void assertQuery( long expectedStart = System.nanoTime(); MaterializedResult expectedResults = null; try { - expectedResults = expectedQueryRunner.execute(session, expected, actualResults.getTypes()); + expectedResults = expectedQueryRunner.execute(expectedSession, expected, actualResults.getTypes()); } catch (RuntimeException ex) { fail("Execution of 'expected' query failed: " + expected, ex);