From 2d8633d53cdc81725ece31b8ff8d66b88336713d Mon Sep 17 00:00:00 2001 From: Poon Yat Sing Date: Thu, 9 Oct 2025 11:16:23 +0900 Subject: [PATCH] Refactor PostgreSQLContainer initialization in tests --- .../io/trino/gateway/TrinoGatewayRunner.java | 3 +- .../ha/TestGatewayHaMultipleBackend.java | 3 +- .../ha/TestGatewayHaSingleBackend.java | 3 +- ...atewayHaWithRoutingRulesSingleBackend.java | 3 +- .../io/trino/gateway/ha/TestNoXForwarded.java | 3 +- .../trino/gateway/ha/TestTrinoResource.java | 3 +- .../TestDatabaseMigrationsPostgreSql.java | 5 ++- ...TestExternalUrlQueryHistoryPostgreSql.java | 4 +- .../TestQueryHistoryManagerPostgresql.java | 5 ++- .../gateway/ha/router/TestRoutingAPI.java | 3 +- .../ha/security/TestAuthorization.java | 3 +- .../trino/gateway/ha/security/TestOIDC.java | 8 ++-- .../gateway/ha/util/TestcontainersUtils.java | 38 +++++++++++++++++++ .../proxyserver/TestProxyRequestHandler.java | 3 +- 14 files changed, 68 insertions(+), 19 deletions(-) create mode 100644 gateway-ha/src/test/java/io/trino/gateway/ha/util/TestcontainersUtils.java diff --git a/gateway-ha/src/test/java/io/trino/gateway/TrinoGatewayRunner.java b/gateway-ha/src/test/java/io/trino/gateway/TrinoGatewayRunner.java index 8c145bde8..adbfc8f1c 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/TrinoGatewayRunner.java +++ b/gateway-ha/src/test/java/io/trino/gateway/TrinoGatewayRunner.java @@ -22,6 +22,7 @@ import java.util.List; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static org.testcontainers.utility.MountableFile.forClasspathResource; public final class TrinoGatewayRunner @@ -43,7 +44,7 @@ public static void main(String[] args) trino2.withCopyFileToContainer(forClasspathResource("trino-config.properties"), "/etc/trino/config.properties"); trino2.start(); - PostgreSQLContainer postgres = new PostgreSQLContainer("postgres:17"); + PostgreSQLContainer postgres = createPostgreSqlContainer(); postgres.withUsername("trino_gateway_db_admin"); postgres.withPassword("P0stG&es"); postgres.withDatabaseName("trino_gateway_db"); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java index 931bd2fdf..b11ebe1d3 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java @@ -50,6 +50,7 @@ import static com.google.common.net.HttpHeaders.CONTENT_TYPE; import static com.google.common.net.MediaType.JSON_UTF_8; import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static org.assertj.core.api.Assertions.assertThat; import static org.testcontainers.utility.MountableFile.forClasspathResource; @@ -65,7 +66,7 @@ final class TestGatewayHaMultipleBackend private TrinoContainer adhocTrino; private TrinoContainer scheduledTrino; - private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:17"); + private final PostgreSQLContainer postgresql = createPostgreSqlContainer(); public static String oauthInitiatePath = OAuth2GatewayCookie.OAUTH2_PATH; public static String oauthCallbackPath = oauthInitiatePath + "/callback"; diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaSingleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaSingleBackend.java index 3f9b43659..f9ceeaae5 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaSingleBackend.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaSingleBackend.java @@ -34,6 +34,7 @@ import java.io.File; import java.util.List; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static org.assertj.core.api.Assertions.assertThat; import static org.testcontainers.utility.MountableFile.forClasspathResource; @@ -41,7 +42,7 @@ final class TestGatewayHaSingleBackend { private TrinoContainer trino; - private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:17"); + private final PostgreSQLContainer postgresql = createPostgreSqlContainer(); int routerPort = 21001 + (int) (Math.random() * 1000); @BeforeAll diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithRoutingRulesSingleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithRoutingRulesSingleBackend.java index 6c74a367f..43720018e 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithRoutingRulesSingleBackend.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaWithRoutingRulesSingleBackend.java @@ -29,6 +29,7 @@ import java.io.File; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static org.assertj.core.api.Assertions.assertThat; import static org.testcontainers.utility.MountableFile.forClasspathResource; @@ -37,7 +38,7 @@ final class TestGatewayHaWithRoutingRulesSingleBackend { private final OkHttpClient httpClient = new OkHttpClient(); private TrinoContainer trino; - private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:17"); + private final PostgreSQLContainer postgresql = createPostgreSqlContainer(); int routerPort = 21001 + (int) (Math.random() * 1000); @BeforeAll diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestNoXForwarded.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestNoXForwarded.java index 45d446eee..fd762605a 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestNoXForwarded.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestNoXForwarded.java @@ -30,6 +30,7 @@ import java.io.File; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static org.assertj.core.api.Assertions.assertThat; import static org.testcontainers.utility.MountableFile.forClasspathResource; @@ -38,7 +39,7 @@ final class TestNoXForwarded { private final OkHttpClient httpClient = new OkHttpClient(); private TrinoContainer trino; - private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:17"); + private final PostgreSQLContainer postgresql = createPostgreSqlContainer(); int routerPort = 21001 + (int) (Math.random() * 1000); int backendPort; diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java index ab9423442..9d3307ac3 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/TestTrinoResource.java @@ -39,6 +39,7 @@ import static io.trino.gateway.ha.router.ResourceGroupsManager.GlobalPropertiesDetail; import static io.trino.gateway.ha.router.ResourceGroupsManager.ResourceGroupsDetail; import static io.trino.gateway.ha.router.ResourceGroupsManager.SelectorsDetail; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.MethodOrderer.OrderAnnotation; @@ -48,7 +49,7 @@ final class TestTrinoResource { private static final OkHttpClient httpClient = new OkHttpClient(); private static final ObjectMapper objectMapper = new ObjectMapper(); - private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:17"); + private final PostgreSQLContainer postgresql = createPostgreSqlContainer(); int routerPort = 22000 + (int) (Math.random() * 1000); JdbcConnectionManager connectionManager; diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgreSql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgreSql.java index 7bd51488f..3f413035d 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgreSql.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/persistence/TestDatabaseMigrationsPostgreSql.java @@ -14,14 +14,15 @@ package io.trino.gateway.ha.persistence; import org.jdbi.v3.core.Handle; -import org.testcontainers.containers.PostgreSQLContainer; + +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; final class TestDatabaseMigrationsPostgreSql extends BaseTestDatabaseMigrations { public TestDatabaseMigrationsPostgreSql() { - super(new PostgreSQLContainer<>("postgres:11"), "public"); + super(createPostgreSqlContainer(), "public"); } @Override diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalUrlQueryHistoryPostgreSql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalUrlQueryHistoryPostgreSql.java index 69312c680..5dc6b8ef6 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalUrlQueryHistoryPostgreSql.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestExternalUrlQueryHistoryPostgreSql.java @@ -13,13 +13,13 @@ */ package io.trino.gateway.ha.router; -import org.testcontainers.containers.PostgreSQLContainer; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; public class TestExternalUrlQueryHistoryPostgreSql extends BaseExternalUrlQueryHistoryTest { public TestExternalUrlQueryHistoryPostgreSql() { - super(new PostgreSQLContainer<>("postgres:15.5")); + super(createPostgreSqlContainer()); } } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestQueryHistoryManagerPostgresql.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestQueryHistoryManagerPostgresql.java index 458c5343d..a9b65c796 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestQueryHistoryManagerPostgresql.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestQueryHistoryManagerPostgresql.java @@ -14,7 +14,8 @@ package io.trino.gateway.ha.router; import org.testcontainers.containers.JdbcDatabaseContainer; -import org.testcontainers.containers.PostgreSQLContainer; + +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; public class TestQueryHistoryManagerPostgresql extends BaseTestQueryHistoryManager @@ -22,7 +23,7 @@ public class TestQueryHistoryManagerPostgresql @Override protected final JdbcDatabaseContainer startContainer() { - JdbcDatabaseContainer container = new PostgreSQLContainer<>("postgres:17"); + JdbcDatabaseContainer container = createPostgreSqlContainer(); container.start(); return container; } diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java index 248dc1e71..ae9cfdac9 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/router/TestRoutingAPI.java @@ -34,6 +34,7 @@ import java.io.File; import java.util.List; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static org.assertj.core.api.Assertions.assertThat; import static org.testcontainers.utility.MountableFile.forClasspathResource; @@ -43,7 +44,7 @@ final class TestRoutingAPI public static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private final OkHttpClient httpClient = new OkHttpClient(); private TrinoContainer trino; - private final PostgreSQLContainer postgresql = new PostgreSQLContainer<>("postgres:17"); + private final PostgreSQLContainer postgresql = createPostgreSqlContainer(); int routerPort = 21001 + (int) (Math.random() * 1000); int backendPort; diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestAuthorization.java b/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestAuthorization.java index 52e8819c7..e922f5523 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestAuthorization.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestAuthorization.java @@ -38,6 +38,7 @@ import java.util.List; import java.util.Optional; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static org.assertj.core.api.Assertions.assertThat; @TestInstance(Lifecycle.PER_CLASS) @@ -45,7 +46,7 @@ final class TestAuthorization { private static final OkHttpClient httpClient = new OkHttpClient(); private static final ObjectMapper objectMapper = new ObjectMapper(); - private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:17"); + private final PostgreSQLContainer postgresql = createPostgreSqlContainer(); int routerPort = 22000 + (int) (Math.random() * 1000); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestOIDC.java b/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestOIDC.java index 6ff2abaf2..2ba71405e 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestOIDC.java +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/security/TestOIDC.java @@ -55,6 +55,7 @@ import static io.trino.gateway.ha.HaGatewayTestUtils.buildPostgresVars; import static io.trino.gateway.ha.security.OidcCookie.OIDC_COOKIE; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static java.lang.String.format; import static org.assertj.core.api.Assertions.assertThat; import static org.testcontainers.utility.MountableFile.forClasspathResource; @@ -75,13 +76,12 @@ void setup() { Network network = Network.newNetwork(); - PostgreSQLContainer databaseContainer = new PostgreSQLContainer<>("postgres:17") + PostgreSQLContainer databaseContainer = createPostgreSqlContainer() .withNetwork(network) .withNetworkAliases("hydra-db") .withUsername("hydra") .withPassword("mysecretpassword") - .withDatabaseName("hydra") - .waitingFor(Wait.forLogMessage(".*ready to accept connections.*", 1)); + .withDatabaseName("hydra"); databaseContainer.start(); GenericContainer migrationContainer = new GenericContainer(HYDRA_IMAGE) @@ -146,7 +146,7 @@ void setup() "--callbacks", callbackUrl); clientCreatingContainer.start(); - PostgreSQLContainer gatewayBackendDatabase = new PostgreSQLContainer("postgres:17"); + PostgreSQLContainer gatewayBackendDatabase = createPostgreSqlContainer(); gatewayBackendDatabase.start(); URL resource = HaGatewayTestUtils.class.getClassLoader().getResource("auth/localhost.jks"); diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/util/TestcontainersUtils.java b/gateway-ha/src/test/java/io/trino/gateway/ha/util/TestcontainersUtils.java new file mode 100644 index 000000000..2f5ea238e --- /dev/null +++ b/gateway-ha/src/test/java/io/trino/gateway/ha/util/TestcontainersUtils.java @@ -0,0 +1,38 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.gateway.ha.util; + +import org.testcontainers.containers.PostgreSQLContainer; +import org.testcontainers.containers.wait.strategy.LogMessageWaitStrategy; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.containers.wait.strategy.WaitAllStrategy; + +import java.time.Duration; + +public final class TestcontainersUtils +{ + private TestcontainersUtils() {} + + public static PostgreSQLContainer createPostgreSqlContainer() + { + //noinspection resource + return new PostgreSQLContainer<>("postgres:17") + .waitingFor(new WaitAllStrategy() + .withStrategy(Wait.forListeningPort()) + .withStrategy(new LogMessageWaitStrategy() + .withRegEx(".*database system is ready to accept connections.*\\s") + .withTimes(2) + .withStartupTimeout(Duration.ofMinutes(1)))); + } +} diff --git a/gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java b/gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java index 8311c2ead..fece434d4 100644 --- a/gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java +++ b/gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyRequestHandler.java @@ -43,6 +43,7 @@ import static io.trino.gateway.ha.HaGatewayTestUtils.setUpBackend; import static io.trino.gateway.ha.handler.HttpUtils.V1_STATEMENT_PATH; import static io.trino.gateway.ha.handler.ProxyUtils.SOURCE_HEADER; +import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer; import static java.nio.charset.StandardCharsets.UTF_8; import static org.assertj.core.api.Assertions.assertThat; import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; @@ -52,7 +53,7 @@ final class TestProxyRequestHandler { private final OkHttpClient httpClient = new OkHttpClient(); private final MockWebServer mockTrinoServer = new MockWebServer(); - private final PostgreSQLContainer postgresql = new PostgreSQLContainer("postgres:17"); + private final PostgreSQLContainer postgresql = createPostgreSqlContainer(); private final int routerPort = 21001 + (int) (Math.random() * 1000); private final int customBackendPort = 21000 + (int) (Math.random() * 1000);