From 35ad5c164b924f30b579d126a6ab10036acd642e Mon Sep 17 00:00:00 2001
From: Piotr Findeisen <piotr.findeisen@gmail.com>
Date: Thu, 21 Dec 2023 17:15:39 +0100
Subject: [PATCH] Set up only one concurrent query runner in TestQueryManager

`TestQueryManager` defined a default `queryRunner` used by one test
method, but other test methods created additional runners. This commit
moves the default one into the only test method that used it. In the
result, there is at most 1 query runner at the same time created during
test execution.
---
 .../java/io/trino/tests/TestQueryManager.java | 77 +++++++------------
 1 file changed, 29 insertions(+), 48 deletions(-)

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