Fix withInitScript compatibility with ReplicaSet and Sharding modes

Author: taole33

modules/mongodb/src/main/java/org/testcontainers/mongodb/MongoDBContainer.java

@@ -41,10 +41,16 @@ public class MongoDBContainer extends GenericContainer<MongoDBContainer> {
 
     private static final String STARTER_SCRIPT = "/testcontainers_start.sh";
 
+    private static final String SCRIPT_DESTINATION_DEFAULT = "/docker-entrypoint-initdb.d/init.js";
+
+    private static final String SCRIPT_DESTINATION_MANUAL = "/tmp/init.js";
+
     private boolean shardingEnabled;
 
     private boolean rsEnabled;
 
+    private String initScriptPath;
+
     public MongoDBContainer(@NonNull String dockerImageName) {
         this(DockerImageName.parse(dockerImageName));
     }
@@ -68,6 +74,26 @@ protected void containerIsStarted(InspectContainerResponse containerInfo, boolea
         if (this.rsEnabled) {
             initReplicaSet(reused);
         }
+
+        boolean isClusterMode = this.shardingEnabled || this.rsEnabled;
+
+        if (isClusterMode && this.initScriptPath != null) {
+            executeInitScriptInContainer();
+        }
+    }
+
+    @Override
+    protected void configure() {
+        super.configure();
+        boolean isClusterMode = this.shardingEnabled || this.rsEnabled;
+        if (this.initScriptPath != null) {
+            String destination = isClusterMode ? SCRIPT_DESTINATION_MANUAL : SCRIPT_DESTINATION_DEFAULT;
+            withCopyFileToContainer(MountableFile.forClasspathResource(this.initScriptPath), destination);
+        }
+
+        if (this.initScriptPath != null && !isClusterMode) {
+            this.waitStrategy = Wait.forLogMessage("(?i).*waiting for connections.*", 2);
+        }
     }
 
     private String[] buildMongoEvalCommand(String command) {
@@ -208,19 +234,42 @@ public String getReplicaSetUrl(String databaseName) {
     /**
      * Executes a MongoDB initialization script from the classpath during startup.
      * <p>
-     * The script will be copied to {@code /docker-entrypoint-initdb.d/init.js}.
-     * This method also adjusts the {@link org.testcontainers.containers.wait.strategy.WaitStrategy}
-     * to expect the "waiting for connections" log message twice, as the execution of an init script
-     * causes MongoDB to restart.
+     * In standalone mode, the script will be copied to {@code /docker-entrypoint-initdb.d/init.js},
+     * and the {@link org.testcontainers.containers.wait.strategy.WaitStrategy} is adjusted
+     * to expect the "waiting for connections" log message twice.
+     * <p>
+     * In Replica Set or Sharding mode, the script is copied to a temporary location and executed
+     * manually after the cluster is initialized.
      *
      * @param scriptPath the path to the init script file on the classpath
      * @return this container instance
      */
     public MongoDBContainer withInitScript(String scriptPath) {
-        withCopyFileToContainer(MountableFile.forClasspathResource(scriptPath), "/docker-entrypoint-initdb.d/init.js");
-
-        this.waitStrategy = Wait.forLogMessage("(?i).*waiting for connections.*", 2);
-
+        this.initScriptPath = scriptPath;
         return this;
     }
+
+    @SneakyThrows
+    private void executeInitScriptInContainer() {
+        String cmd =
+            "mongosh " +
+            MONGODB_DATABASE_NAME_DEFAULT +
+            " " +
+            SCRIPT_DESTINATION_MANUAL +
+            " || mongo " +
+            MONGODB_DATABASE_NAME_DEFAULT +
+            " " +
+            SCRIPT_DESTINATION_MANUAL;
+
+        ExecResult result = execInContainer("sh", "-c", cmd);
+        if (result.getExitCode() != CONTAINER_EXIT_CODE_OK) {
+            throw new IllegalStateException(
+                String.format(
+                    "Failed to execute init script.\nStdout: %s\nStderr: %s",
+                    result.getStdout(),
+                    result.getStderr()
+                )
+            );
+        }
+    }
 }

modules/mongodb/src/test/java/org/testcontainers/mongodb/MongoDBContainerTest.java

@@ -42,15 +42,80 @@ void shouldTestDatabaseName() {
     }
 
     @Test
-    void testWithInitScript() {
+    void shouldExecuteInitScript() {
         try (
             MongoDBContainer mongoDB = new MongoDBContainer("mongo:4.0.10")
                 .withInitScript("init.js")
                 .withStartupTimeout(Duration.ofSeconds(30))
         ) {
             mongoDB.start();
-
             assertThat(mongoDB.isRunning()).isTrue();
         }
     }
+
+    @Test
+    void shouldExecuteInitScriptWithEdgeCases() {
+        try (
+            MongoDBContainer mongoDB = new MongoDBContainer("mongo:4.0.10")
+                .withInitScript("initEdgeCase.js")
+                .withEnv("LANG", "C.UTF-8")
+                .withEnv("LC_ALL", "C.UTF-8")
+                .withStartupTimeout(Duration.ofSeconds(30))
+        ) {
+            mongoDB.start();
+
+            try (
+                com.mongodb.client.MongoClient client = com.mongodb.client.MongoClients.create(
+                    mongoDB.getReplicaSetUrl()
+                )
+            ) {
+                String expectedComplexName = "test_col_\"_with_specials_!@#%^&*()";
+                String expectedJapaneseName = "日本語 コレクション ﾃｽﾄ";
+
+                assertThat(client.getDatabase("test").listCollectionNames())
+                    .as("Check if init script created the collections with special chars and Japanese")
+                    .contains(expectedComplexName, expectedJapaneseName);
+            }
+        }
+    }
+
+    @Test
+    void shouldExecuteInitScriptWithReplicaSet() {
+        try (MongoDBContainer mongo = new MongoDBContainer("mongo:7.0.0").withInitScript("init.js").withReplicaSet()) {
+            mongo.start();
+            assertInitScriptExecuted(mongo);
+        }
+    }
+
+    @Test
+    void shouldExecuteInitScriptWithReplicaSetConfiguredFirst() {
+        try (MongoDBContainer mongo = new MongoDBContainer("mongo:7.0.0").withReplicaSet().withInitScript("init.js")) {
+            mongo.start();
+            assertInitScriptExecuted(mongo);
+        }
+    }
+
+    @Test
+    void shouldExecuteInitScriptWithSharding() {
+        try (MongoDBContainer mongo = new MongoDBContainer("mongo:7.0.0").withInitScript("init.js").withSharding()) {
+            mongo.start();
+            assertInitScriptExecuted(mongo);
+        }
+    }
+
+    @Test
+    void shouldExecuteInitScriptWithShardingConfiguredFirst() {
+        try (MongoDBContainer mongo = new MongoDBContainer("mongo:7.0.0").withSharding().withInitScript("init.js")) {
+            mongo.start();
+            assertInitScriptExecuted(mongo);
+        }
+    }
+
+    private void assertInitScriptExecuted(MongoDBContainer mongo) {
+        try (com.mongodb.client.MongoClient client = com.mongodb.client.MongoClients.create(mongo.getReplicaSetUrl())) {
+            assertThat(client.getDatabase("test").listCollectionNames())
+                .as("Check if init.js created the collection")
+                .contains("test_collection");
+        }
+    }
 }

modules/mongodb/src/test/resources/initEdgeCase.js

@@ -0,0 +1,16 @@
+var complexCollectionName = 'test_col_"_with_specials_!@#%^&*()';
+
+db.createCollection(complexCollectionName);
+
+var japaneseCollectionName = "日本語 コレクション ﾃｽﾄ";
+
+db.createCollection(japaneseCollectionName);
+
+db.getCollection(complexCollectionName).insertOne({
+    "_id": 1,
+    "key_with_quotes": "This is a \"double quoted\" string",
+    "key_with_json_chars": "{ } [ ] : ,",
+    "description": "特殊記号を含むコレクションへの挿入テスト"
+});
+
+print("Initialization completed: " + complexCollectionName);