Fix cockroachdb wait strategy with version >= 22.1.0

Starting with version 22.1.0, wait strategy can rely on `/cockroach/init_success` file created when scripts are executed. Fixes #8555
testcontainers · Nov 7, 2024 · 0c0a868 · 0c0a868
1 parent be7251b
commit 0c0a868
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 12 deletions.
diff --git a/modules/cockroachdb/src/main/java/org/testcontainers/containers/CockroachContainer.java b/modules/cockroachdb/src/main/java/org/testcontainers/containers/CockroachContainer.java
@@ -1,6 +1,7 @@
 package org.testcontainers.containers;
 
-import org.testcontainers.containers.wait.strategy.HttpWaitStrategy;
+import org.testcontainers.containers.wait.strategy.Wait;
+import org.testcontainers.containers.wait.strategy.WaitAllStrategy;
 import org.testcontainers.utility.ComparableVersion;
 import org.testcontainers.utility.DockerImageName;
 
@@ -66,19 +67,31 @@ public CockroachContainer(final String dockerImageName) {
     public CockroachContainer(final DockerImageName dockerImageName) {
         super(dockerImageName);
         dockerImageName.assertCompatibleWith(DEFAULT_IMAGE_NAME);
-        isVersionGreaterThanOrEqualTo221 = isVersionGreaterThanOrEqualTo221(dockerImageName);
+        this.isVersionGreaterThanOrEqualTo221 = isVersionGreaterThanOrEqualTo221(dockerImageName);
 
-        withExposedPorts(REST_API_PORT, DB_PORT);
-        waitingFor(
-            new HttpWaitStrategy()
-                .forPath("/health")
-                .forPort(REST_API_PORT)
-                .forStatusCode(200)
-                .withStartupTimeout(Duration.ofMinutes(1))
+        WaitAllStrategy waitStrategy = new WaitAllStrategy();
+        waitStrategy.withStrategy(
+            Wait.forHttp("/health").forPort(REST_API_PORT).forStatusCode(200).withStartupTimeout(Duration.ofMinutes(1))
         );
+        if (this.isVersionGreaterThanOrEqualTo221) {
+            waitStrategy.withStrategy(Wait.forSuccessfulCommand("[ -f ./init_success ] || { exit 1; }"));
+        }
+
+        withExposedPorts(REST_API_PORT, DB_PORT);
+        waitingFor(waitStrategy);
         withCommand("start-single-node --insecure");
     }
 
+    @Override
+    protected void configure() {
+        withEnv("COCKROACH_USER", this.username);
+        withEnv("COCKROACH_PASSWORD", this.password);
+        if (this.password != null && !this.password.isEmpty()) {
+            withCommand("start-single-node");
+        }
+        withEnv("COCKROACH_DATABASE", this.databaseName);
+    }
+
     @Override
     public String getDriverClassName() {
         return JDBC_DRIVER_CLASS_NAME;
@@ -123,21 +136,21 @@ public String getTestQueryString() {
     public CockroachContainer withUsername(String username) {
         validateIfVersionSupportsUsernameOrPasswordOrDatabase("username");
         this.username = username;
-        return withEnv("COCKROACH_USER", username);
+        return this;
     }
 
     @Override
     public CockroachContainer withPassword(String password) {
         validateIfVersionSupportsUsernameOrPasswordOrDatabase("password");
         this.password = password;
-        return withEnv("COCKROACH_PASSWORD", password).withCommand("start-single-node");
+        return this;
     }
 
     @Override
     public CockroachContainer withDatabaseName(final String databaseName) {
         validateIfVersionSupportsUsernameOrPasswordOrDatabase("databaseName");
         this.databaseName = databaseName;
-        return withEnv("COCKROACH_DATABASE", databaseName);
+        return this;
     }
 
     private boolean isVersionGreaterThanOrEqualTo221(DockerImageName dockerImageName) {
@@ -152,4 +165,9 @@ private void validateIfVersionSupportsUsernameOrPasswordOrDatabase(String parame
             );
         }
     }
+
+    @Override
+    protected void waitUntilContainerStarted() {
+        getWaitStrategy().waitUntilReady(this);
+    }
 }
diff --git a/...cockroachdb/src/test/java/org/testcontainers/junit/cockroachdb/SimpleCockroachDBTest.java b/...cockroachdb/src/test/java/org/testcontainers/junit/cockroachdb/SimpleCockroachDBTest.java
@@ -4,6 +4,7 @@
 import org.testcontainers.CockroachDBTestImages;
 import org.testcontainers.containers.CockroachContainer;
 import org.testcontainers.db.AbstractContainerDatabaseTest;
+import org.testcontainers.images.builder.Transferable;
 
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -105,4 +106,25 @@ public void testAnExceptionIsThrownWhenImageDoesNotSupportEnvVars() {
             .isInstanceOf(UnsupportedOperationException.class)
             .withFailMessage("Setting a databaseName in not supported in the versions below 22.1.0");
     }
+
+    @Test
+    public void testInitializationScript() throws SQLException {
+        String sql =
+            "USE postgres; \n" +
+            "CREATE TABLE bar (foo VARCHAR(255)); \n" +
+            "INSERT INTO bar (foo) VALUES ('hello world');";
+
+        try (
+            CockroachContainer cockroach = new CockroachContainer(CockroachDBTestImages.COCKROACHDB_IMAGE)
+                .withCopyToContainer(Transferable.of(sql), "/docker-entrypoint-initdb.d/init.sql")
+                .withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
+        ) { // CockroachDB is expected to be compatible with Postgres
+            cockroach.start();
+
+            ResultSet resultSet = performQuery(cockroach, "SELECT foo FROM bar");
+
+            String firstColumnValue = resultSet.getString(1);
+            assertThat(firstColumnValue).as("Value from init script should equal real value").isEqualTo("hello world");
+        }
+    }
 }