From 4e6248ce6edc08ea7e1a6eeb09054fb34b3dd2c7 Mon Sep 17 00:00:00 2001
From: Heitor Danilo <heitor.danilo@icloud.com>
Date: Tue, 7 May 2024 10:31:52 -0300
Subject: [PATCH] feat(mongodb): add WithReplicaSet option

The `WithReplicaSet` option configures the MongoDB container to run a
single-node replica set named "rs". The container will wait until the
replica set is ready.
---
 modules/mongodb/mongodb.go      | 35 ++++++++++++++++++++++++++
 modules/mongodb/mongodb_test.go | 44 ++++++++++++++++++++++++---------
 2 files changed, 67 insertions(+), 12 deletions(-)

diff --git a/modules/mongodb/mongodb.go b/modules/mongodb/mongodb.go
index 565e8bc466b..ac294148251 100644
--- a/modules/mongodb/mongodb.go
+++ b/modules/mongodb/mongodb.go
@@ -79,6 +79,29 @@ func WithPassword(password string) testcontainers.CustomizeRequestOption {
 	}
 }
 
+// WithReplicaSet configures the container to run a single-node MongoDB replica set named "rs".
+// It will wait until the replica set is ready.
+func WithReplicaSet() testcontainers.CustomizeRequestOption {
+	return func(req *testcontainers.GenericContainerRequest) error {
+		req.Cmd = append(req.Cmd, "--replSet", "rs")
+		req.LifecycleHooks = append(req.LifecycleHooks, testcontainers.ContainerLifecycleHooks{
+			PostStarts: []testcontainers.ContainerHook{
+				func(ctx context.Context, c testcontainers.Container) error {
+					ip, err := c.ContainerIP(ctx)
+					if err != nil {
+						return err
+					}
+
+					cmd := eval("rs.initiate({ _id: 'rs', members: [ { _id: 0, host: '%s:27017' } ] })", ip)
+					return wait.ForExec(cmd).WaitUntilReady(ctx, c)
+				},
+			},
+		})
+
+		return nil
+	}
+}
+
 // ConnectionString returns the connection string for the MongoDB container.
 // If you provide a username and a password, the connection string will also include them.
 func (c *MongoDBContainer) ConnectionString(ctx context.Context) (string, error) {
@@ -95,3 +118,15 @@ func (c *MongoDBContainer) ConnectionString(ctx context.Context) (string, error)
 	}
 	return c.Endpoint(ctx, "mongodb")
 }
+
+// eval builds an mongosh|mongo eval command.
+func eval(command string, args ...any) []string {
+	command = "\"" + fmt.Sprintf(command, args...) + "\""
+
+	return []string{
+		"sh",
+		"-c",
+		// In previous versions, the binary "mongosh" was named "mongo".
+		"mongosh --quiet --eval " + command + " || mongo --quiet --eval " + command,
+	}
+}
diff --git a/modules/mongodb/mongodb_test.go b/modules/mongodb/mongodb_test.go
index 994b4c448cb..8d588026db7 100644
--- a/modules/mongodb/mongodb_test.go
+++ b/modules/mongodb/mongodb_test.go
@@ -14,32 +14,52 @@ import (
 
 func TestMongoDB(t *testing.T) {
 	type tests struct {
-		name  string
-		image string
+		name string
+		opts []testcontainers.ContainerCustomizer
 	}
 	testCases := []tests{
 		{
-			name:  "From Docker Hub",
-			image: "mongo:6",
+			name: "From Docker Hub",
+			opts: []testcontainers.ContainerCustomizer{
+				testcontainers.WithImage("mongo:6"),
+			},
 		},
 		{
-			name:  "Community Server",
-			image: "mongodb/mongodb-community-server:7.0.2-ubi8",
+			name: "Community Server",
+			opts: []testcontainers.ContainerCustomizer{
+				testcontainers.WithImage("mongodb/mongodb-community-server:7.0.2-ubi8"),
+			},
 		},
 		{
-			name:  "Enterprise Server",
-			image: "mongodb/mongodb-enterprise-server:7.0.0-ubi8",
+			name: "Enterprise Server",
+			opts: []testcontainers.ContainerCustomizer{
+				testcontainers.WithImage("mongodb/mongodb-enterprise-server:7.0.0-ubi8"),
+			},
+		},
+		{
+			name: "With Replica set and mongo:4",
+			opts: []testcontainers.ContainerCustomizer{
+				testcontainers.WithImage("mongo:4"),
+				mongodb.WithReplicaSet(),
+			},
+		},
+		{
+			name: "With Replica set and mongo:6",
+			opts: []testcontainers.ContainerCustomizer{
+				testcontainers.WithImage("mongo:6"),
+				mongodb.WithReplicaSet(),
+			},
 		},
 	}
 
-	for _, tc := range testCases {
-		image := tc.image
-		t.Run(image, func(t *testing.T) {
+	for _, tt := range testCases {
+		tc := tt
+		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 
 			ctx := context.Background()
 
-			mongodbContainer, err := mongodb.RunContainer(ctx, testcontainers.WithImage(image))
+			mongodbContainer, err := mongodb.RunContainer(ctx, tc.opts...)
 			if err != nil {
 				t.Fatalf("failed to start container: %s", err)
 			}