Refactor PostgreSQLContainer initialization in tests

oneonestar · ebyhr · commit 535802e9a731 · 2025-10-10T13:38:03.000+09:00
diff --git 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
@@ -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
@@ -34,14 +34,15 @@
 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;
 
 @TestInstance(Lifecycle.PER_CLASS)
 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -14,15 +14,16 @@
 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
 {
     @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
@@ -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
@@ -38,14 +38,15 @@
 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)
 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
@@ -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
@@ -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
@@ -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);

Original file line number	Diff line number	Diff line change
`@@ -14,14 +14,15 @@`
`14`	`14`	`package io.trino.gateway.ha.persistence;`
`15`	`15`
`16`	`16`	`import org.jdbi.v3.core.Handle;`
`17`		`-import org.testcontainers.containers.PostgreSQLContainer;`
	`17`	`+`
	`18`	`+import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer;`
`18`	`19`
`19`	`20`	`final class TestDatabaseMigrationsPostgreSql`
`20`	`21`	`extends BaseTestDatabaseMigrations`
`21`	`22`	`{`
`22`	`23`	`public TestDatabaseMigrationsPostgreSql()`
`23`	`24`	`{`
`24`		`- super(new PostgreSQLContainer<>("postgres:11"), "public");`
	`25`	`+ super(createPostgreSqlContainer(), "public");`
`25`	`26`	`}`
`26`	`27`
`27`	`28`	`@Override`
Original file line number	Diff line number	Diff line change
`@@ -13,13 +13,13 @@`
`13`	`13`	`*/`
`14`	`14`	`package io.trino.gateway.ha.router;`
`15`	`15`
`16`		`-import org.testcontainers.containers.PostgreSQLContainer;`
	`16`	`+import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer;`
`17`	`17`
`18`	`18`	`public class TestExternalUrlQueryHistoryPostgreSql`
`19`	`19`	`extends BaseExternalUrlQueryHistoryTest`
`20`	`20`	`{`
`21`	`21`	`public TestExternalUrlQueryHistoryPostgreSql()`
`22`	`22`	`{`
`23`		`- super(new PostgreSQLContainer<>("postgres:15.5"));`
	`23`	`+ super(createPostgreSqlContainer());`
`24`	`24`	`}`
`25`	`25`	`}`
Original file line number	Diff line number	Diff line change
`@@ -14,15 +14,16 @@`
`14`	`14`	`package io.trino.gateway.ha.router;`
`15`	`15`
`16`	`16`	`import org.testcontainers.containers.JdbcDatabaseContainer;`
`17`		`-import org.testcontainers.containers.PostgreSQLContainer;`
	`17`	`+`
	`18`	`+import static io.trino.gateway.ha.util.TestcontainersUtils.createPostgreSqlContainer;`
`18`	`19`
`19`	`20`	`public class TestQueryHistoryManagerPostgresql`
`20`	`21`	`extends BaseTestQueryHistoryManager`
`21`	`22`	`{`
`22`	`23`	`@Override`
`23`	`24`	`protected final JdbcDatabaseContainer<?> startContainer()`
`24`	`25`	`{`
`25`		`- JdbcDatabaseContainer<?> container = new PostgreSQLContainer<>("postgres:17");`
	`26`	`+ JdbcDatabaseContainer<?> container = createPostgreSqlContainer();`
`26`	`27`	`container.start();`
`27`	`28`	`return container;`
`28`	`29`	`}`