Enhance assertQuery to support actual and expected session parameters.

prestodb · Apr 12, 2022 · 4bf475f · 4bf475f
1 parent 23b1e85
commit 4bf475f
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 6 deletions.
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
@@ -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<Plan> 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
@@ -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<Plan> 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<Plan> 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);