Refactor PostgreSQLContainer initialization in tests

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");

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";

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

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

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;
 

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;

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

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());
     }
 }

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;
     }

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;
 

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);
 

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");

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))));
+    }
+}

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);