diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestQueryRunnerUtil.java b/testing/trino-tests/src/test/java/io/trino/execution/QueryRunnerUtil.java similarity index 72% rename from testing/trino-tests/src/test/java/io/trino/execution/TestQueryRunnerUtil.java rename to testing/trino-tests/src/test/java/io/trino/execution/QueryRunnerUtil.java index 8df5c721dea8..87ddb86786a3 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/TestQueryRunnerUtil.java +++ b/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 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; - } - } } diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestDeprecatedFunctionWarning.java b/testing/trino-tests/src/test/java/io/trino/execution/TestDeprecatedFunctionWarning.java index f1e3116c0e61..9994780ec88e 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/TestDeprecatedFunctionWarning.java +++ b/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 warnings = queryRunner.execute(session, sql).getWarnings(); + List 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 expectedWarnings) { - Session session = testSessionBuilder() - .setCatalog("tpch") - .setSchema("tiny") - .build(); - Set warnings = queryRunner.execute(session, sql).getWarnings().stream() + Set warnings = queryRunner.execute(sql).getWarnings().stream() .map(Warning::getWarningCode) .map(Warning.Code::getCode) .collect(toImmutableSet()); diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestExecutionJmxMetrics.java b/testing/trino-tests/src/test/java/io/trino/execution/TestExecutionJmxMetrics.java index 89dd64c22e48..129473fc8463 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/TestExecutionJmxMetrics.java +++ b/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; diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestPendingStageState.java b/testing/trino-tests/src/test/java/io/trino/execution/TestPendingStageState.java index 97b0545cd501..245e3fd4b76f 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/TestPendingStageState.java +++ b/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; diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestQueues.java b/testing/trino-tests/src/test/java/io/trino/execution/TestQueues.java index 60a9fd313c7e..9f7d4787b13e 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/TestQueues.java +++ b/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(); diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestWarnings.java b/testing/trino-tests/src/test/java/io/trino/execution/TestWarnings.java index 47269a85cecc..f5ab18d75e74 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/TestWarnings.java +++ b/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 expectedWarnings) + private static void assertWarnings(QueryRunner queryRunner, @Language("SQL") String sql, List expectedWarnings) { - Set warnings = queryRunner.execute(session, sql).getWarnings().stream() + Set warnings = queryRunner.execute(sql).getWarnings().stream() .map(Warning::getWarningCode) .map(Warning.Code::getCode) .collect(toImmutableSet()); diff --git a/testing/trino-tests/src/test/java/io/trino/execution/resourcegroups/db/TestEnvironments.java b/testing/trino-tests/src/test/java/io/trino/execution/resourcegroups/db/TestEnvironments.java index 8080528025fd..196cb1f75972 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/resourcegroups/db/TestEnvironments.java +++ b/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; diff --git a/testing/trino-tests/src/test/java/io/trino/execution/resourcegroups/db/TestQueuesDb.java b/testing/trino-tests/src/test/java/io/trino/execution/resourcegroups/db/TestQueuesDb.java index ed4bb35641eb..8a9a503a1b57 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/resourcegroups/db/TestQueuesDb.java +++ b/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; diff --git a/testing/trino-tests/src/test/java/io/trino/tests/TestQueryManager.java b/testing/trino-tests/src/test/java/io/trino/tests/TestQueryManager.java index 930af42529b8..5c919b934bd6 100644 --- a/testing/trino-tests/src/test/java/io/trino/tests/TestQueryManager.java +++ b/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;