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 4b41a0f303bf..f7a7a91a4dbb 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 @@ -27,10 +27,7 @@ import io.trino.testing.DistributedQueryRunner; import io.trino.testing.TestingSessionContext; import io.trino.tests.tpch.TpchQueryRunnerBuilder; -import org.junit.jupiter.api.AfterAll; -import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.TestInstance; import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.parallel.Execution; @@ -48,64 +45,48 @@ import static java.util.Arrays.stream; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Fail.fail; -import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; import static org.junit.jupiter.api.parallel.ExecutionMode.SAME_THREAD; -@TestInstance(PER_CLASS) @Execution(SAME_THREAD) // run single threaded to avoid creating multiple query runners at once public class TestQueryManager { - private DistributedQueryRunner queryRunner; - - @BeforeAll - public void setUp() - throws Exception - { - queryRunner = TpchQueryRunnerBuilder.builder().build(); - } - - @AfterAll - public void tearDown() - { - queryRunner.close(); - queryRunner = null; - } - @Test @Timeout(60) public void testFailQuery() throws Exception { - DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); - QueryId queryId = dispatchManager.createQueryId(); - dispatchManager.createQuery( - queryId, - Span.getInvalid(), - Slug.createNew(), - TestingSessionContext.fromSession(TEST_SESSION), - "SELECT * FROM lineitem") - .get(); + try (DistributedQueryRunner queryRunner = TpchQueryRunnerBuilder.builder().build()) { + DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); + QueryId queryId = dispatchManager.createQueryId(); + dispatchManager.createQuery( + queryId, + Span.getInvalid(), + Slug.createNew(), + TestingSessionContext.fromSession(TEST_SESSION), + "SELECT * FROM lineitem") + .get(); - // wait until query starts running - while (true) { - QueryState state = dispatchManager.getQueryInfo(queryId).getState(); - if (state.isDone()) { - fail("unexpected query state: " + state); + // wait until query starts running + while (true) { + QueryState state = dispatchManager.getQueryInfo(queryId).getState(); + if (state.isDone()) { + fail("unexpected query state: " + state); + } + if (state == RUNNING) { + break; + } + Thread.sleep(100); } - if (state == RUNNING) { - break; - } - Thread.sleep(100); - } - // cancel query - QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); - queryManager.failQuery(queryId, new TrinoException(GENERIC_INTERNAL_ERROR, "mock exception")); - QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); - assertThat(queryInfo.getState()).isEqualTo(FAILED); - assertThat(queryInfo.getErrorCode()).isEqualTo(GENERIC_INTERNAL_ERROR.toErrorCode()); - assertThat(queryInfo.getFailureInfo()).isNotNull(); - assertThat(queryInfo.getFailureInfo().getMessage()).isEqualTo("mock exception"); + // cancel query + QueryManager queryManager = queryRunner.getCoordinator().getQueryManager(); + queryManager.failQuery(queryId, new TrinoException(GENERIC_INTERNAL_ERROR, "mock exception")); + QueryInfo queryInfo = queryManager.getFullQueryInfo(queryId); + assertThat(queryInfo.getState()).isEqualTo(FAILED); + assertThat(queryInfo.getErrorCode()).isEqualTo(GENERIC_INTERNAL_ERROR.toErrorCode()); + assertThat(queryInfo.getFailureInfo()).isNotNull(); + assertThat(queryInfo.getFailureInfo().getMessage()).isEqualTo("mock exception"); + } } @Test