Refactor TestQueryRunnerUtil

testing/trino-tests/src/test/java/io/trino/execution/TestQueryRunnerUtil.java → testing/trino-tests/src/test/java/io/trino/execution/QueryRunnerUtil.java

@@ -13,28 +13,24 @@
  */
 package io.trino.execution;
 
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import io.trino.Session;
 import io.trino.dispatcher.DispatchManager;
-import io.trino.plugin.tpch.TpchPlugin;
 import io.trino.server.BasicQueryInfo;
 import io.trino.server.protocol.Slug;
 import io.trino.spi.QueryId;
 import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.TestingSessionContext;
 
-import java.util.Map;
 import java.util.Set;
 
 import static io.airlift.concurrent.MoreFutures.getFutureValue;
 import static io.trino.execution.QueryState.RUNNING;
-import static io.trino.testing.TestingSession.testSessionBuilder;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
-public final class TestQueryRunnerUtil
+public final class QueryRunnerUtil
 {
-    private TestQueryRunnerUtil() {}
+    private QueryRunnerUtil() {}
 
     public static QueryId createQuery(DistributedQueryRunner queryRunner, Session session, String sql)
     {
@@ -69,29 +65,4 @@ public static void waitForQueryState(DistributedQueryRunner queryRunner, QueryId
         }
         while (!expectedQueryStates.contains(dispatchManager.getQueryInfo(queryId).getState()));
     }
-
-    public static DistributedQueryRunner createQueryRunner()
-            throws Exception
-    {
-        return createQueryRunner(ImmutableMap.of());
-    }
-
-    public static DistributedQueryRunner createQueryRunner(Map<String, String> extraProperties)
-            throws Exception
-    {
-        DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(testSessionBuilder().build())
-                .setExtraProperties(extraProperties)
-                .setNodeCount(2)
-                .build();
-
-        try {
-            queryRunner.installPlugin(new TpchPlugin());
-            queryRunner.createCatalog("tpch", "tpch");
-            return queryRunner;
-        }
-        catch (Exception e) {
-            queryRunner.close();
-            throw e;
-        }
-    }
 }

testing/trino-tests/src/test/java/io/trino/execution/TestDeprecatedFunctionWarning.java

@@ -14,8 +14,6 @@
 package io.trino.execution;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import io.trino.Session;
 import io.trino.client.Warning;
 import io.trino.metadata.InternalFunctionBundle;
 import io.trino.operator.aggregation.state.LongAndDoubleState;
@@ -34,6 +32,7 @@
 import io.trino.spi.function.TypeParameter;
 import io.trino.spi.function.WindowFunctionSignature;
 import io.trino.spi.type.StandardTypes;
+import io.trino.testing.DistributedQueryRunner;
 import io.trino.testing.QueryRunner;
 import org.intellij.lang.annotations.Language;
 import org.testng.annotations.AfterClass;
@@ -44,7 +43,6 @@
 import java.util.Set;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static io.trino.execution.TestQueryRunnerUtil.createQueryRunner;
 import static io.trino.spi.type.DoubleType.DOUBLE;
 import static io.trino.testing.TestingSession.testSessionBuilder;
 import static org.testng.Assert.fail;
@@ -60,7 +58,7 @@ public class TestDeprecatedFunctionWarning
     public void setUp()
             throws Exception
     {
-        queryRunner = createQueryRunner(ImmutableMap.of());
+        queryRunner = DistributedQueryRunner.builder(testSessionBuilder().build()).build();
         ImmutableList.of(
                 TestScalaFunction.class,
                 TestDeprecatedParametericScalaFunction.class,
@@ -101,11 +99,7 @@ public void testDeprecatedScalaFunction()
     public void testDeprecatedDescription()
     {
         String sql = "SELECT deprecated_scalar()";
-        Session session = testSessionBuilder()
-                .setCatalog("tpch")
-                .setSchema("tiny")
-                .build();
-        List<Warning> warnings = queryRunner.execute(session, sql).getWarnings();
+        List<Warning> warnings = queryRunner.execute(sql).getWarnings();
         if (warnings.size() != 1 || !warnings.get(0).getMessage().contains(EXPECTED_WARNING_MSG)) {
             fail("Expected warning: " + EXPECTED_WARNING_MSG);
         }
@@ -163,11 +157,7 @@ public void testDeprecatedWindowFunction()
 
     private static void assertPlannerWarnings(QueryRunner queryRunner, @Language("SQL") String sql, List<WarningCode> expectedWarnings)
     {
-        Session session = testSessionBuilder()
-                .setCatalog("tpch")
-                .setSchema("tiny")
-                .build();
-        Set<Integer> warnings = queryRunner.execute(session, sql).getWarnings().stream()
+        Set<Integer> warnings = queryRunner.execute(sql).getWarnings().stream()
                 .map(Warning::getWarningCode)
                 .map(Warning.Code::getCode)
                 .collect(toImmutableSet());

testing/trino-tests/src/test/java/io/trino/execution/TestExecutionJmxMetrics.java

@@ -26,12 +26,12 @@
 import javax.management.MBeanServer;
 import javax.management.ObjectName;
 
+import static io.trino.execution.QueryRunnerUtil.cancelQuery;
+import static io.trino.execution.QueryRunnerUtil.createQuery;
+import static io.trino.execution.QueryRunnerUtil.waitForQueryState;
 import static io.trino.execution.QueryState.FAILED;
 import static io.trino.execution.QueryState.QUEUED;
 import static io.trino.execution.QueryState.RUNNING;
-import static io.trino.execution.TestQueryRunnerUtil.cancelQuery;
-import static io.trino.execution.TestQueryRunnerUtil.createQuery;
-import static io.trino.execution.TestQueryRunnerUtil.waitForQueryState;
 import static io.trino.testing.TestingSession.testSessionBuilder;
 import static org.testng.Assert.assertEquals;
 

testing/trino-tests/src/test/java/io/trino/execution/TestPendingStageState.java

@@ -23,9 +23,9 @@
 import org.testng.annotations.Test;
 
 import static io.trino.SessionTestUtils.TEST_SESSION;
+import static io.trino.execution.QueryRunnerUtil.createQuery;
+import static io.trino.execution.QueryRunnerUtil.waitForQueryState;
 import static io.trino.execution.QueryState.RUNNING;
-import static io.trino.execution.TestQueryRunnerUtil.createQuery;
-import static io.trino.execution.TestQueryRunnerUtil.waitForQueryState;
 import static io.trino.plugin.tpch.TpchConnectorFactory.TPCH_SPLITS_PER_NODE;
 import static io.trino.testing.assertions.Assert.assertEventually;
 import static java.util.concurrent.TimeUnit.SECONDS;

testing/trino-tests/src/test/java/io/trino/execution/TestQueues.java

@@ -20,6 +20,7 @@
 import io.trino.Session;
 import io.trino.dispatcher.DispatchManager;
 import io.trino.plugin.resourcegroups.ResourceGroupManagerPlugin;
+import io.trino.plugin.tpch.TpchPlugin;
 import io.trino.spi.QueryId;
 import io.trino.spi.resourcegroups.ResourceGroupId;
 import io.trino.spi.session.ResourceEstimates;
@@ -33,15 +34,14 @@
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
 import static io.airlift.units.DataSize.Unit.TERABYTE;
 import static io.trino.SystemSessionProperties.MAX_HASH_PARTITION_COUNT;
+import static io.trino.execution.QueryRunnerUtil.cancelQuery;
+import static io.trino.execution.QueryRunnerUtil.createQuery;
+import static io.trino.execution.QueryRunnerUtil.waitForQueryState;
 import static io.trino.execution.QueryState.FAILED;
 import static io.trino.execution.QueryState.FINISHED;
 import static io.trino.execution.QueryState.FINISHING;
 import static io.trino.execution.QueryState.QUEUED;
 import static io.trino.execution.QueryState.RUNNING;
-import static io.trino.execution.TestQueryRunnerUtil.cancelQuery;
-import static io.trino.execution.TestQueryRunnerUtil.createQuery;
-import static io.trino.execution.TestQueryRunnerUtil.createQueryRunner;
-import static io.trino.execution.TestQueryRunnerUtil.waitForQueryState;
 import static io.trino.spi.StandardErrorCode.QUERY_REJECTED;
 import static io.trino.testing.TestingSession.testSessionBuilder;
 import static java.lang.String.format;
@@ -312,6 +312,24 @@ private void testRejection()
         }
     }
 
+    private static DistributedQueryRunner createQueryRunner()
+            throws Exception
+    {
+        DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(testSessionBuilder().build())
+                .setNodeCount(2)
+                .build();
+
+        try {
+            queryRunner.installPlugin(new TpchPlugin());
+            queryRunner.createCatalog("tpch", "tpch");
+            return queryRunner;
+        }
+        catch (Exception e) {
+            queryRunner.close();
+            throw e;
+        }
+    }
+
     private String getResourceFilePath(String fileName)
     {
         return this.getClass().getClassLoader().getResource(fileName).getPath();

testing/trino-tests/src/test/java/io/trino/execution/TestWarnings.java

@@ -14,11 +14,10 @@
 package io.trino.execution;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
-import io.trino.Session;
 import io.trino.client.Warning;
 import io.trino.spi.WarningCode;
 import io.trino.testing.QueryRunner;
+import io.trino.tests.tpch.TpchQueryRunnerBuilder;
 import org.intellij.lang.annotations.Language;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -28,9 +27,7 @@
 import java.util.Set;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
-import static io.trino.execution.TestQueryRunnerUtil.createQueryRunner;
 import static io.trino.spi.connector.StandardWarningCode.TOO_MANY_STAGES;
-import static io.trino.testing.TestingSession.testSessionBuilder;
 import static org.testng.Assert.fail;
 
 public class TestWarnings
@@ -42,7 +39,9 @@ public class TestWarnings
     public void setUp()
             throws Exception
     {
-        queryRunner = createQueryRunner(ImmutableMap.of("query.stage-count-warning-threshold", String.valueOf(STAGE_COUNT_WARNING_THRESHOLD)));
+        queryRunner = TpchQueryRunnerBuilder.builder()
+                .addExtraProperty("query.stage-count-warning-threshold", String.valueOf(STAGE_COUNT_WARNING_THRESHOLD))
+                .build();
     }
 
     @AfterClass(alwaysRun = true)
@@ -63,17 +62,13 @@ public void testStageCountWarningThreshold()
                     .append(stageIndex);
         }
         String query = queryBuilder.toString();
-        Session session = testSessionBuilder()
-                .setCatalog("tpch")
-                .setSchema("tiny")
-                .build();
-        assertWarnings(queryRunner, session, query, ImmutableList.of(TOO_MANY_STAGES.toWarningCode()));
-        assertWarnings(queryRunner, session, noWarningsQuery, ImmutableList.of());
+        assertWarnings(queryRunner, query, ImmutableList.of(TOO_MANY_STAGES.toWarningCode()));
+        assertWarnings(queryRunner, noWarningsQuery, ImmutableList.of());
     }
 
-    private static void assertWarnings(QueryRunner queryRunner, Session session, @Language("SQL") String sql, List<WarningCode> expectedWarnings)
+    private static void assertWarnings(QueryRunner queryRunner, @Language("SQL") String sql, List<WarningCode> expectedWarnings)
     {
-        Set<Integer> warnings = queryRunner.execute(session, sql).getWarnings().stream()
+        Set<Integer> warnings = queryRunner.execute(sql).getWarnings().stream()
                 .map(Warning::getWarningCode)
                 .map(Warning.Code::getCode)
                 .collect(toImmutableSet());

testing/trino-tests/src/test/java/io/trino/execution/resourcegroups/db/TestEnvironments.java

@@ -18,10 +18,10 @@
 import io.trino.testing.DistributedQueryRunner;
 import org.testng.annotations.Test;
 
+import static io.trino.execution.QueryRunnerUtil.createQuery;
+import static io.trino.execution.QueryRunnerUtil.waitForQueryState;
 import static io.trino.execution.QueryState.FAILED;
 import static io.trino.execution.QueryState.RUNNING;
-import static io.trino.execution.TestQueryRunnerUtil.createQuery;
-import static io.trino.execution.TestQueryRunnerUtil.waitForQueryState;
 import static io.trino.execution.resourcegroups.db.H2TestUtil.TEST_ENVIRONMENT;
 import static io.trino.execution.resourcegroups.db.H2TestUtil.TEST_ENVIRONMENT_2;
 import static io.trino.execution.resourcegroups.db.H2TestUtil.adhocSession;

testing/trino-tests/src/test/java/io/trino/execution/resourcegroups/db/TestQueuesDb.java

@@ -37,13 +37,13 @@
 
 import static io.airlift.testing.Assertions.assertContains;
 import static io.trino.SystemSessionProperties.QUERY_MAX_EXECUTION_TIME;
+import static io.trino.execution.QueryRunnerUtil.cancelQuery;
+import static io.trino.execution.QueryRunnerUtil.createQuery;
+import static io.trino.execution.QueryRunnerUtil.waitForQueryState;
 import static io.trino.execution.QueryState.FAILED;
 import static io.trino.execution.QueryState.FINISHED;
 import static io.trino.execution.QueryState.QUEUED;
 import static io.trino.execution.QueryState.RUNNING;
-import static io.trino.execution.TestQueryRunnerUtil.cancelQuery;
-import static io.trino.execution.TestQueryRunnerUtil.createQuery;
-import static io.trino.execution.TestQueryRunnerUtil.waitForQueryState;
 import static io.trino.execution.TestQueues.createResourceGroupId;
 import static io.trino.execution.resourcegroups.db.H2TestUtil.TEST_ENVIRONMENT;
 import static io.trino.execution.resourcegroups.db.H2TestUtil.adhocSession;

testing/trino-tests/src/test/java/io/trino/tests/TestQueryManager.java

@@ -32,10 +32,10 @@
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static io.trino.SessionTestUtils.TEST_SESSION;
+import static io.trino.execution.QueryRunnerUtil.createQuery;
+import static io.trino.execution.QueryRunnerUtil.waitForQueryState;
 import static io.trino.execution.QueryState.FAILED;
 import static io.trino.execution.QueryState.RUNNING;
-import static io.trino.execution.TestQueryRunnerUtil.createQuery;
-import static io.trino.execution.TestQueryRunnerUtil.waitForQueryState;
 import static io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME;
 import static io.trino.spi.StandardErrorCode.EXCEEDED_CPU_LIMIT;
 import static io.trino.spi.StandardErrorCode.EXCEEDED_SCAN_LIMIT;