feat(modules.mongodb): setup mongodb replicaset

modules/mongodb/mongodb.go

@@ -3,6 +3,8 @@ package mongodb
 import (
 	"context"
 	"fmt"
+	"io"
+	"strings"
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/wait"
@@ -14,8 +16,9 @@ const defaultImage = "mongo:6"
 // MongoDBContainer represents the MongoDB container type used in the module
 type MongoDBContainer struct {
 	testcontainers.Container
-	username string
-	password string
+	username   string
+	password   string
+	replicaSet string
 }
 
 // RunContainer creates an instance of the MongoDB container type
@@ -38,21 +41,30 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize
 	for _, opt := range opts {
 		opt.Customize(&genericContainerReq)
 	}
+
+	container, err := testcontainers.GenericContainer(ctx, genericContainerReq)
+	if err != nil {
+		return nil, err
+	}
+	mongoContainer := &MongoDBContainer{Container: container}
+
 	username := req.Env["MONGO_INITDB_ROOT_USERNAME"]
 	password := req.Env["MONGO_INITDB_ROOT_PASSWORD"]
 	if username != "" && password == "" || username == "" && password != "" {
 		return nil, fmt.Errorf("if you specify username or password, you must provide both of them")
 	}
 
-	container, err := testcontainers.GenericContainer(ctx, genericContainerReq)
-	if err != nil {
-		return nil, err
+	if username != "" && password != "" {
+		mongoContainer.username = username
+		mongoContainer.password = password
 	}
 
-	if username != "" && password != "" {
-		return &MongoDBContainer{Container: container, username: username, password: password}, nil
+	replicaSet := req.Env["MONGO_REPLICASET_NAME"]
+	if replicaSet != "" {
+		mongoContainer.replicaSet = replicaSet
 	}
-	return &MongoDBContainer{Container: container}, nil
+
+	return mongoContainer, nil
 }
 
 // WithUsername sets the initial username to be created when the container starts
@@ -73,6 +85,15 @@ func WithPassword(password string) testcontainers.CustomizeRequestOption {
 	}
 }
 
+// WithReplicaSet TODO: help me fill this func comment
+func WithReplicaSet(rsName string) testcontainers.CustomizeRequestOption {
+	return func(req *testcontainers.GenericContainerRequest) {
+		req.Env["MONGO_REPLICASET_NAME"] = rsName
+		req.Cmd = append(req.Cmd, "--replSet", rsName)
+		req.LifecycleHooks = append(req.LifecycleHooks, replicaSetLifecycleHooks())
+	}
+}
+
 // 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) {
@@ -84,8 +105,66 @@ func (c *MongoDBContainer) ConnectionString(ctx context.Context) (string, error)
 	if err != nil {
 		return "", err
 	}
+
+	connStr := strings.Builder{}
+	connStr.WriteString("mongodb://")
+
 	if c.username != "" && c.password != "" {
+		connStr.WriteString(fmt.Sprintf("%s:%s@", c.username, c.password))
 		return fmt.Sprintf("mongodb://%s:%s@%s:%s", c.username, c.password, host, port.Port()), nil
 	}
-	return c.Endpoint(ctx, "mongodb")
+
+	connStr.WriteString(fmt.Sprintf("%s:%s", host, port.Port()))
+
+	if c.replicaSet != "" {
+		connStr.WriteString(fmt.Sprintf("/test?replicaSet=%s&directConnection=true", c.replicaSet))
+	}
+
+	return connStr.String(), nil
+}
+
+func replicaSetLifecycleHooks() testcontainers.ContainerLifecycleHooks {
+	return testcontainers.ContainerLifecycleHooks{
+		PostStarts: []testcontainers.ContainerHook{
+			replicaSetPostStart,
+		},
+	}
+}
+
+func replicaSetPostStart(ctx context.Context, container testcontainers.Container) error {
+	const (
+		waitScript = `var attempt = 0;
+		while (%s) {
+			if (attempt > %d) {
+				quit(1);
+			}
+
+			print('%s ' + attempt);
+			sleep(100);
+			attempt++;
+		}`
+		waitCond        = "db.runCommand( { isMaster: 1 } ).ismaster==false"
+		maxWaitAttempts = 60
+	)
+
+	exitCode, r, err := container.Exec(ctx, []string{"mongosh", "--eval", "rs.initiate();"})
+	if err != nil {
+		return err
+	}
+	if exitCode != 0 {
+		dataVjp, _ := io.ReadAll(r)
+		return fmt.Errorf("non-zero exit code to initiate replica: %d, %s", exitCode, string(dataVjp))
+	}
+
+	waitCmd := fmt.Sprintf(waitScript, waitCond, maxWaitAttempts, "An attempt to await for a single node replica set initialization:")
+	exitCode, _, err = container.Exec(ctx, []string{"mongosh", "--eval", waitCmd})
+	if err != nil {
+		return err
+	}
+
+	if exitCode != 0 {
+		return fmt.Errorf("none-zero exit code when await replica initiate")
+	}
+
+	return nil
 }

modules/mongodb/mongodb_test.go

@@ -28,7 +28,6 @@ func ExampleRunContainer() {
 			panic(err)
 		}
 	}()
-	// }
 
 	state, err := mongodbContainer.State(ctx)
 	if err != nil {
@@ -118,3 +117,44 @@ func ExampleRunContainer_withCredentials() {
 	// Output:
 	// mongodb://root:password
 }
+
+func ExampleRunContainer_withReplicaSet() {
+	ctx := context.Background()
+
+	container, err := mongodb.RunContainer(ctx,
+		testcontainers.WithImage("mongo:6"),
+		mongodb.WithUsername("root"),
+		mongodb.WithPassword("password"),
+		mongodb.WithReplicaSet("rs0"),
+		testcontainers.WithWaitStrategy(wait.ForLog("Waiting for connections")),
+	)
+	if err != nil {
+		panic(err)
+	}
+
+	// Clean up the container
+	defer func() {
+		if err := container.Terminate(ctx); err != nil {
+			panic(err)
+		}
+	}()
+
+	connStr, err := container.ConnectionString(ctx)
+	if err != nil {
+		panic(err)
+	}
+
+	mongoClient, err := mongo.Connect(ctx, options.Client().ApplyURI(connStr))
+	if err != nil {
+		panic(err)
+	}
+
+	err = mongoClient.Ping(ctx, nil)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Println(strings.Split(connStr, "@")[0])
+
+	// Output:
+	// mongodb://root:password
+}