feat: support MongoDB replica set

docs/community/mongodb.rst

@@ -1,3 +1,4 @@
 .. autoclass:: testcontainers.community.mongodb.MongoDbContainer
+.. autoclass:: testcontainers.community.mongodb.MongoDbReplicaSetContainer
 .. autoclass:: testcontainers.community.mongodb.MongoDBAtlasLocalContainer
 .. title:: testcontainers.community.mongodb.MongoDbContainer

docs/modules/mongodb_example.py

@@ -3,7 +3,7 @@
 
 from pymongo import MongoClient
 
-from testcontainers.community.mongodb import MongoDbContainer
+from testcontainers.community.mongodb import MongoDbContainer, MongoDbReplicaSetContainer
 
 
 def basic_example():
@@ -81,5 +81,45 @@ def basic_example():
         print(json.dumps(stats, default=str, indent=2))
 
 
+def replica_set_example():
+    with MongoDbReplicaSetContainer() as mongodb:
+        client = mongodb.get_connection_client()
+        db = client.test_db
+
+        print("\nConnected to MongoDB replica set")
+        print(json.dumps(client.admin.command("hello"), default=str, indent=2))
+
+        # Run a multi-document transaction
+        accounts = db.accounts
+        with client.start_session() as session, session.start_transaction():
+            accounts.insert_one({"name": "checking", "balance": 100}, session=session)
+            accounts.update_one(
+                {"name": "checking"},
+                {"$inc": {"balance": -25}},
+                session=session,
+            )
+
+        print("\nTransaction result:")
+        print(json.dumps(accounts.find_one({"name": "checking"}), default=str, indent=2))
+
+        # Observe an insert through a change stream
+        events = db.events
+        with events.watch([{"$match": {"operationType": "insert"}}]) as changes:
+            events.insert_one({"kind": "created", "value": 42})
+            change = next(changes)
+
+        print("\nChange stream event:")
+        print(json.dumps(change, default=str, indent=2))
+
+
+def unauthenticated_replica_set_example():
+    with MongoDbReplicaSetContainer(auth_enabled=False) as mongodb:
+        client = mongodb.get_connection_client()
+        print("\nConnected to MongoDB replica set without authentication")
+        print(json.dumps(client.admin.command("hello"), default=str, indent=2))
+
+
 if __name__ == "__main__":
     basic_example()
+    replica_set_example()
+    unauthenticated_replica_set_example()

src/testcontainers/community/mongodb/__init__.py

@@ -10,16 +10,39 @@
 #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 #    License for the specific language governing permissions and limitations
 #    under the License.
+import base64
 import os
 import re
-from typing import Optional
+import shlex
+import time
+from typing import Any, Optional
+from urllib.parse import urlencode
 
 from pymongo import MongoClient
+from pymongo.errors import OperationFailure, PyMongoError
+from typing_extensions import Self
 
+from testcontainers.core.config import testcontainers_config
+from testcontainers.core.exceptions import ContainerStartException
 from testcontainers.core.generic import DbContainer
 from testcontainers.core.utils import raise_for_deprecated_parameter
 from testcontainers.core.wait_strategies import HealthcheckWaitStrategy, LogMessageWaitStrategy
 
+_REPLICA_SET_KEYFILE_PATH = "/tmp/testcontainers-mongodb-keyfile"
+_REPLICA_SET_ENTRYPOINT_PATH = "/tmp/testcontainers-mongodb-entrypoint.sh"
+_REPLICA_SET_ENTRYPOINT = f"""#!/bin/bash
+set -Eeuo pipefail
+chown mongodb:mongodb {_REPLICA_SET_KEYFILE_PATH}
+chmod 400 {_REPLICA_SET_KEYFILE_PATH}
+exec /usr/local/bin/docker-entrypoint.sh "$@"
+""".encode()
+
+
+def _is_root_container_user(user: Any) -> bool:
+    if user in (None, "", 0, "0", "root"):
+        return True
+    return isinstance(user, str) and user.partition(":")[0] in ("0", "root")
+
 
 class MongoDbContainer(DbContainer):
     """
@@ -89,6 +112,203 @@ def get_connection_client(self) -> MongoClient:
         return MongoClient(self.get_connection_url())
 
 
+class MongoDbReplicaSetContainer(MongoDbContainer):
+    """MongoDB container configured as a single-node replica set.
+
+    Authentication is enabled by default and uses an ephemeral keyfile for internal
+    replica-set authentication. Set ``auth_enabled=False`` to run without
+    authentication or a keyfile.
+
+    Automatic keyfile setup targets the official ``mongo`` image and requires the
+    container to start as root before the image drops privileges to ``mongodb``.
+
+    Example:
+
+        .. code-block:: python
+
+            from testcontainers.community.mongodb import MongoDbReplicaSetContainer
+
+            with MongoDbReplicaSetContainer("mongo:7.0.7") as mongo:
+                client = mongo.get_connection_client()
+                with client.start_session() as session, session.start_transaction():
+                    client.test.items.insert_one({"name": "example"}, session=session)
+
+            with MongoDbReplicaSetContainer(
+                "mongo:7.0.7",
+                auth_enabled=False,
+            ) as mongo:
+                client = mongo.get_connection_client()
+    """
+
+    def __init__(
+        self,
+        image: str = "mongo:latest",
+        port: int = 27017,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        dbname: Optional[str] = None,
+        replica_set: str = "docker-rs",
+        auth_enabled: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        if not replica_set:
+            raise ValueError("replica_set must not be empty")
+        if not auth_enabled and (username is not None or password is not None):
+            raise ValueError("username and password cannot be set when authentication is disabled")
+        if auth_enabled:
+            if "entrypoint" in kwargs:
+                raise ValueError("entrypoint cannot be overridden for an authenticated replica set")
+            if "user" in kwargs and not _is_root_container_user(kwargs["user"]):
+                raise ValueError("authenticated replica sets must start as the root container user")
+            kwargs["entrypoint"] = _REPLICA_SET_ENTRYPOINT_PATH
+
+        super().__init__(
+            image=image,
+            port=port,
+            username=username,
+            password=password,
+            dbname=dbname,
+            **kwargs,
+        )
+        self.auth_enabled = auth_enabled
+        self.replica_set = replica_set
+        if not auth_enabled:
+            self.username = ""
+            self.password = ""
+        else:
+            super().with_copy_into_container(_REPLICA_SET_ENTRYPOINT, _REPLICA_SET_ENTRYPOINT_PATH, mode=0o755)
+            keyfile = base64.b64encode(os.urandom(756))
+            super().with_copy_into_container(keyfile, _REPLICA_SET_KEYFILE_PATH, mode=0o400)
+        command = MongoDbReplicaSetContainer._replica_set_command(self, self._command)
+        super().with_command(command)
+
+    def _replica_set_command(self, command: Optional[str | list[str]]) -> list[str]:
+        command_parts = shlex.split(command) if isinstance(command, str) else list(command or [])
+        if command_parts and not (command_parts[0].startswith("-") or os.path.basename(command_parts[0]) == "mongod"):
+            raise ValueError("replica set commands must contain mongod options or start with mongod")
+        if any(
+            part in ("--replSet", "--keyFile") or part.startswith(("--replSet=", "--keyFile="))
+            for part in command_parts
+        ):
+            raise ValueError("replica set and keyfile options are managed by MongoDbReplicaSetContainer")
+
+        command_parts.extend(["--replSet", self.replica_set])
+        if self.auth_enabled:
+            command_parts.extend(["--keyFile", _REPLICA_SET_KEYFILE_PATH])
+        return command_parts
+
+    def with_command(self, command: str | list[str]) -> Self:
+        return super().with_command(self._replica_set_command(command))
+
+    def with_kwargs(self, **kwargs: Any) -> Self:
+        if self.auth_enabled:
+            if "entrypoint" in kwargs:
+                raise ValueError("entrypoint cannot be overridden for an authenticated replica set")
+            if "user" in kwargs and not _is_root_container_user(kwargs["user"]):
+                raise ValueError("authenticated replica sets must start as the root container user")
+            kwargs["entrypoint"] = _REPLICA_SET_ENTRYPOINT_PATH
+        return super().with_kwargs(**kwargs)
+
+    def _configure(self) -> None:
+        if self.auth_enabled:
+            assert self.username is not None
+            assert self.password is not None
+            self.with_env("MONGO_INITDB_ROOT_USERNAME", self.username)
+            self.with_env("MONGO_INITDB_ROOT_PASSWORD", self.password)
+        else:
+            self.env.pop("MONGO_INITDB_ROOT_USERNAME", None)
+            self.env.pop("MONGO_INITDB_ROOT_PASSWORD", None)
+        self.with_env("MONGO_DB", self.dbname)
+
+    def get_connection_url(self) -> str:
+        if self.auth_enabled:
+            url = super().get_connection_url()
+        else:
+            host = self.get_container_host_ip()
+            port = self.get_exposed_port(self.port)
+            url = f"mongodb://{host}:{port}"
+
+        return f"{url}/?{urlencode({'replicaSet': self.replica_set, 'directConnection': 'true'})}"
+
+    def _connect(self) -> None:
+        direct_url = self.get_connection_url().replace(
+            urlencode({"replicaSet": self.replica_set, "directConnection": "true"}),
+            urlencode({"directConnection": "true"}),
+        )
+        client: MongoClient[dict[str, Any]] = MongoClient(
+            direct_url,
+            serverSelectionTimeoutMS=1000,
+            connectTimeoutMS=1000,
+            socketTimeoutMS=1000,
+        )
+        deadline = time.monotonic() + testcontainers_config.timeout
+
+        try:
+            self._wait_for_mongodb(client, deadline)
+            self._wait_for_replica_set_primary(client, deadline)
+        finally:
+            client.close()
+
+    def _wait_for_mongodb(self, client: MongoClient[dict[str, Any]], deadline: float) -> None:
+        last_error: Optional[Exception] = None
+        while time.monotonic() < deadline:
+            self._raise_if_replica_set_container_stopped()
+            try:
+                client.admin.command("ping")
+                return
+            except PyMongoError as error:
+                last_error = error
+                time.sleep(testcontainers_config.sleep_time)
+
+        raise ContainerStartException("MongoDB did not become ready") from last_error
+
+    def _wait_for_replica_set_primary(self, client: MongoClient[dict[str, Any]], deadline: float) -> None:
+        last_error: Optional[Exception] = None
+        while time.monotonic() < deadline:
+            self._raise_if_replica_set_container_stopped()
+            try:
+                self._initiate_replica_set_if_needed(client)
+                if client.admin.command("hello").get("isWritablePrimary"):
+                    return
+            except OperationFailure:
+                raise
+            except PyMongoError as error:
+                last_error = error
+            time.sleep(testcontainers_config.sleep_time)
+
+        raise ContainerStartException("MongoDB replica set did not elect a primary") from last_error
+
+    def _initiate_replica_set_if_needed(self, client: MongoClient[dict[str, Any]]) -> None:
+        try:
+            client.admin.command("replSetGetStatus")
+            return
+        except OperationFailure as error:
+            if error.code != 94:  # NotYetInitialized
+                raise
+
+        try:
+            client.admin.command(
+                {
+                    "replSetInitiate": {
+                        "_id": self.replica_set,
+                        "members": [{"_id": 0, "host": f"localhost:{self.port}"}],
+                    }
+                }
+            )
+        except OperationFailure as error:
+            if error.code != 23:  # AlreadyInitialized
+                raise
+
+    def _raise_if_replica_set_container_stopped(self) -> None:
+        self.reload()
+        if self.status not in ("exited", "dead"):
+            return
+
+        stdout, stderr = self.get_logs()
+        logs = (stdout + stderr).decode(errors="replace")
+        raise ContainerStartException(f"MongoDB stopped while initializing its replica set:\n{logs}")
+
+
 class MongoDBAtlasLocalContainer(DbContainer):
     """
     MongoDB Atlas Local document-based database container.

tests/community/mongodb/test_mongodb.py

@@ -1,8 +1,9 @@
 import time
+from urllib.parse import parse_qs, urlsplit
 
 import pytest
 
-from testcontainers.community.mongodb import MongoDBAtlasLocalContainer, MongoDbContainer
+from testcontainers.community.mongodb import MongoDBAtlasLocalContainer, MongoDbContainer, MongoDbReplicaSetContainer
 
 
 @pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"])
@@ -28,6 +29,80 @@ def test_docker_run_mongodb(version: str):
         assert cursor.next()["restaurant_id"] == doc["restaurant_id"]
 
 
+@pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"])
+def test_docker_run_mongodb_replica_set_without_authentication(version: str):
+    with MongoDbReplicaSetContainer(f"mongo:{version}", replica_set="docker-rs", auth_enabled=False) as mongo:
+        connection_url = urlsplit(mongo.get_connection_url())
+        assert connection_url.username is None
+        assert parse_qs(connection_url.query) == {
+            "directConnection": ["true"],
+            "replicaSet": ["docker-rs"],
+        }
+
+        client = mongo.get_connection_client()
+
+        hello = client.admin.command("hello")
+        assert hello["setName"] == "docker-rs"
+        assert hello["isWritablePrimary"] is True
+
+        collection = client.test.transactions
+        with client.start_session() as session, session.start_transaction():
+            result = collection.insert_one({"message": "replica sets support transactions"}, session=session)
+
+        assert collection.find_one({"_id": result.inserted_id}) is not None
+
+
+@pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"])
+def test_docker_run_authenticated_mongodb_replica_set(version: str):
+    container = MongoDbReplicaSetContainer(
+        f"mongo:{version}",
+        username="replica-user",
+        password="replica-password",
+        replica_set="secure-rs",
+        command=["--profile", "1"],
+    ).with_kwargs(labels={"testcontainers.mongodb.replica-set": "true"})
+
+    with container as mongo:
+        connection_url = urlsplit(mongo.get_connection_url())
+        assert connection_url.username == "replica-user"
+        assert connection_url.password == "replica-password"
+        assert parse_qs(connection_url.query) == {
+            "directConnection": ["true"],
+            "replicaSet": ["secure-rs"],
+        }
+
+        client = mongo.get_connection_client()
+
+        hello = client.admin.command("hello")
+        assert hello["setName"] == "secure-rs"
+        assert hello["isWritablePrimary"] is True
+
+        result = client.test.documents.insert_one({"authenticated": True})
+        assert client.test.documents.find_one({"_id": result.inserted_id}) == {
+            "_id": result.inserted_id,
+            "authenticated": True,
+        }
+
+
+@pytest.mark.parametrize("username,password", [("user", None), (None, "password")])
+def test_mongodb_rejects_credentials_when_authentication_is_disabled(
+    username: str | None,
+    password: str | None,
+):
+    with pytest.raises(ValueError, match="authentication is disabled"):
+        MongoDbReplicaSetContainer(auth_enabled=False, username=username, password=password)
+
+
+def test_authenticated_mongodb_replica_set_rejects_non_root_container_user():
+    with pytest.raises(ValueError, match="root container user"):
+        MongoDbReplicaSetContainer(replica_set="secure-rs", user="mongodb")
+
+
+def test_mongodb_replica_set_rejects_non_mongod_command():
+    with pytest.raises(ValueError, match="mongod options"):
+        MongoDbReplicaSetContainer(replica_set="secure-rs", command=["bash", "-lc", "exec mongod"])
+
+
 @pytest.mark.parametrize("version", ["8.0.13", "7.0.23"])
 def test_docker_run_mongodb_atlas_local(version: str):
     with MongoDBAtlasLocalContainer(f"mongodb/mongodb-atlas-local:{version}") as mongo_atlas: