From 04886f112f08c2838b9cb353a82470d0bd05f953 Mon Sep 17 00:00:00 2001
From: Gorkem Karaduman <gorkemkrdmn@gmail.com>
Date: Fri, 10 Nov 2023 18:01:54 +0000
Subject: [PATCH] Add username and password functions to mongodb

---
 modules/mongodb/mongodb.go      | 44 ++++++++++++++++++++++++++++++++-
 modules/mongodb/mongodb_test.go | 42 +++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/modules/mongodb/mongodb.go b/modules/mongodb/mongodb.go
index 8e3bb305d1..3dfa75066c 100644
--- a/modules/mongodb/mongodb.go
+++ b/modules/mongodb/mongodb.go
@@ -2,6 +2,7 @@ package mongodb
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/wait"
@@ -13,6 +14,8 @@ const defaultImage = "mongo:6"
 // MongoDBContainer represents the MongoDB container type used in the module
 type MongoDBContainer struct {
 	testcontainers.Container
+	username string
+	password string
 }
 
 // RunContainer creates an instance of the MongoDB container type
@@ -24,6 +27,7 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize
 			wait.ForLog("Waiting for connections"),
 			wait.ForListeningPort("27017/tcp"),
 		),
+		Env: map[string]string{},
 	}
 
 	genericContainerReq := testcontainers.GenericContainerRequest{
@@ -34,16 +38,54 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize
 	for _, opt := range opts {
 		opt.Customize(&genericContainerReq)
 	}
+	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 != "" {
+		return &MongoDBContainer{Container: container, username: username, password: password}, nil
+	}
 	return &MongoDBContainer{Container: container}, nil
 }
 
-// ConnectionString returns the connection string for the MongoDB container
+// WithUsername sets the initial username to be created when the container starts
+// It is used in conjunction with WithPassword to set a username and its password.
+// It will create the specified user with superuser power.
+func WithUsername(username string) testcontainers.CustomizeRequestOption {
+	return func(req *testcontainers.GenericContainerRequest) {
+		req.Env["MONGO_INITDB_ROOT_USERNAME"] = username
+	}
+}
+
+// WithPassword sets the initial password of the user to be created when the container starts
+// It is used in conjunction with WithUsername to set a username and its password.
+// This environment variable sets the superuser password for MongoDB.
+func WithPassword(password string) testcontainers.CustomizeRequestOption {
+	return func(req *testcontainers.GenericContainerRequest) {
+		req.Env["MONGO_INITDB_ROOT_PASSWORD"] = password
+	}
+}
+
+// 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) {
+	host, err := c.Host(ctx)
+	if err != nil {
+		return "", err
+	}
+	port, err := c.MappedPort(ctx, "27017/tcp")
+	if err != nil {
+		return "", err
+	}
+	if 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")
 }
diff --git a/modules/mongodb/mongodb_test.go b/modules/mongodb/mongodb_test.go
index 88b48ff812..22a908d70b 100644
--- a/modules/mongodb/mongodb_test.go
+++ b/modules/mongodb/mongodb_test.go
@@ -3,12 +3,14 @@ package mongodb_test
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/modules/mongodb"
+	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 func ExampleRunContainer() {
@@ -76,3 +78,43 @@ func ExampleRunContainer_connect() {
 	// Output:
 	// test
 }
+
+func ExampleRunContainer_withCredentials() {
+	ctx := context.Background()
+
+	container, err := mongodb.RunContainer(ctx,
+		testcontainers.WithImage("mongo:6"),
+		mongodb.WithUsername("root"),
+		mongodb.WithPassword("password"),
+		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
+}