Skip to content

Commit 2b92d57

Browse files
committed
feat: support MongoDB replica set
1 parent f7d3887 commit 2b92d57

4 files changed

Lines changed: 339 additions & 3 deletions

File tree

docs/community/mongodb.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
.. autoclass:: testcontainers.community.mongodb.MongoDbContainer
2+
.. autoclass:: testcontainers.community.mongodb.MongoDbReplicaSetContainer
23
.. autoclass:: testcontainers.community.mongodb.MongoDBAtlasLocalContainer
34
.. title:: testcontainers.community.mongodb.MongoDbContainer

docs/modules/mongodb_example.py

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
from pymongo import MongoClient
55

6-
from testcontainers.community.mongodb import MongoDbContainer
6+
from testcontainers.community.mongodb import MongoDbContainer, MongoDbReplicaSetContainer
77

88

99
def basic_example():
@@ -81,5 +81,45 @@ def basic_example():
8181
print(json.dumps(stats, default=str, indent=2))
8282

8383

84+
def replica_set_example():
85+
with MongoDbReplicaSetContainer() as mongodb:
86+
client = mongodb.get_connection_client()
87+
db = client.test_db
88+
89+
print("\nConnected to MongoDB replica set")
90+
print(json.dumps(client.admin.command("hello"), default=str, indent=2))
91+
92+
# Run a multi-document transaction
93+
accounts = db.accounts
94+
with client.start_session() as session, session.start_transaction():
95+
accounts.insert_one({"name": "checking", "balance": 100}, session=session)
96+
accounts.update_one(
97+
{"name": "checking"},
98+
{"$inc": {"balance": -25}},
99+
session=session,
100+
)
101+
102+
print("\nTransaction result:")
103+
print(json.dumps(accounts.find_one({"name": "checking"}), default=str, indent=2))
104+
105+
# Observe an insert through a change stream
106+
events = db.events
107+
with events.watch([{"$match": {"operationType": "insert"}}]) as changes:
108+
events.insert_one({"kind": "created", "value": 42})
109+
change = next(changes)
110+
111+
print("\nChange stream event:")
112+
print(json.dumps(change, default=str, indent=2))
113+
114+
115+
def unauthenticated_replica_set_example():
116+
with MongoDbReplicaSetContainer(auth_enabled=False) as mongodb:
117+
client = mongodb.get_connection_client()
118+
print("\nConnected to MongoDB replica set without authentication")
119+
print(json.dumps(client.admin.command("hello"), default=str, indent=2))
120+
121+
84122
if __name__ == "__main__":
85123
basic_example()
124+
replica_set_example()
125+
unauthenticated_replica_set_example()

src/testcontainers/community/mongodb/__init__.py

Lines changed: 221 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,16 +10,39 @@
1010
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
1111
# License for the specific language governing permissions and limitations
1212
# under the License.
13+
import base64
1314
import os
1415
import re
15-
from typing import Optional
16+
import shlex
17+
import time
18+
from typing import Any, Optional
19+
from urllib.parse import urlencode
1620

1721
from pymongo import MongoClient
22+
from pymongo.errors import OperationFailure, PyMongoError
23+
from typing_extensions import Self
1824

25+
from testcontainers.core.config import testcontainers_config
26+
from testcontainers.core.exceptions import ContainerStartException
1927
from testcontainers.core.generic import DbContainer
2028
from testcontainers.core.utils import raise_for_deprecated_parameter
2129
from testcontainers.core.wait_strategies import HealthcheckWaitStrategy, LogMessageWaitStrategy
2230

31+
_REPLICA_SET_KEYFILE_PATH = "/tmp/testcontainers-mongodb-keyfile"
32+
_REPLICA_SET_ENTRYPOINT_PATH = "/tmp/testcontainers-mongodb-entrypoint.sh"
33+
_REPLICA_SET_ENTRYPOINT = f"""#!/bin/bash
34+
set -Eeuo pipefail
35+
chown mongodb:mongodb {_REPLICA_SET_KEYFILE_PATH}
36+
chmod 400 {_REPLICA_SET_KEYFILE_PATH}
37+
exec /usr/local/bin/docker-entrypoint.sh "$@"
38+
""".encode()
39+
40+
41+
def _is_root_container_user(user: Any) -> bool:
42+
if user in (None, "", 0, "0", "root"):
43+
return True
44+
return isinstance(user, str) and user.partition(":")[0] in ("0", "root")
45+
2346

2447
class MongoDbContainer(DbContainer):
2548
"""
@@ -89,6 +112,203 @@ def get_connection_client(self) -> MongoClient:
89112
return MongoClient(self.get_connection_url())
90113

91114

115+
class MongoDbReplicaSetContainer(MongoDbContainer):
116+
"""MongoDB container configured as a single-node replica set.
117+
118+
Authentication is enabled by default and uses an ephemeral keyfile for internal
119+
replica-set authentication. Set ``auth_enabled=False`` to run without
120+
authentication or a keyfile.
121+
122+
Automatic keyfile setup targets the official ``mongo`` image and requires the
123+
container to start as root before the image drops privileges to ``mongodb``.
124+
125+
Example:
126+
127+
.. code-block:: python
128+
129+
from testcontainers.community.mongodb import MongoDbReplicaSetContainer
130+
131+
with MongoDbReplicaSetContainer("mongo:7.0.7") as mongo:
132+
client = mongo.get_connection_client()
133+
with client.start_session() as session, session.start_transaction():
134+
client.test.items.insert_one({"name": "example"}, session=session)
135+
136+
with MongoDbReplicaSetContainer(
137+
"mongo:7.0.7",
138+
auth_enabled=False,
139+
) as mongo:
140+
client = mongo.get_connection_client()
141+
"""
142+
143+
def __init__(
144+
self,
145+
image: str = "mongo:latest",
146+
port: int = 27017,
147+
username: Optional[str] = None,
148+
password: Optional[str] = None,
149+
dbname: Optional[str] = None,
150+
replica_set: str = "docker-rs",
151+
auth_enabled: bool = True,
152+
**kwargs: Any,
153+
) -> None:
154+
if not replica_set:
155+
raise ValueError("replica_set must not be empty")
156+
if not auth_enabled and (username is not None or password is not None):
157+
raise ValueError("username and password cannot be set when authentication is disabled")
158+
if auth_enabled:
159+
if "entrypoint" in kwargs:
160+
raise ValueError("entrypoint cannot be overridden for an authenticated replica set")
161+
if "user" in kwargs and not _is_root_container_user(kwargs["user"]):
162+
raise ValueError("authenticated replica sets must start as the root container user")
163+
kwargs["entrypoint"] = _REPLICA_SET_ENTRYPOINT_PATH
164+
165+
super().__init__(
166+
image=image,
167+
port=port,
168+
username=username,
169+
password=password,
170+
dbname=dbname,
171+
**kwargs,
172+
)
173+
self.auth_enabled = auth_enabled
174+
self.replica_set = replica_set
175+
if not auth_enabled:
176+
self.username = ""
177+
self.password = ""
178+
else:
179+
super().with_copy_into_container(_REPLICA_SET_ENTRYPOINT, _REPLICA_SET_ENTRYPOINT_PATH, mode=0o755)
180+
keyfile = base64.b64encode(os.urandom(756))
181+
super().with_copy_into_container(keyfile, _REPLICA_SET_KEYFILE_PATH, mode=0o400)
182+
command = MongoDbReplicaSetContainer._replica_set_command(self, self._command)
183+
super().with_command(command)
184+
185+
def _replica_set_command(self, command: Optional[str | list[str]]) -> list[str]:
186+
command_parts = shlex.split(command) if isinstance(command, str) else list(command or [])
187+
if command_parts and not (command_parts[0].startswith("-") or os.path.basename(command_parts[0]) == "mongod"):
188+
raise ValueError("replica set commands must contain mongod options or start with mongod")
189+
if any(
190+
part in ("--replSet", "--keyFile") or part.startswith(("--replSet=", "--keyFile="))
191+
for part in command_parts
192+
):
193+
raise ValueError("replica set and keyfile options are managed by MongoDbReplicaSetContainer")
194+
195+
command_parts.extend(["--replSet", self.replica_set])
196+
if self.auth_enabled:
197+
command_parts.extend(["--keyFile", _REPLICA_SET_KEYFILE_PATH])
198+
return command_parts
199+
200+
def with_command(self, command: str | list[str]) -> Self:
201+
return super().with_command(self._replica_set_command(command))
202+
203+
def with_kwargs(self, **kwargs: Any) -> Self:
204+
if self.auth_enabled:
205+
if "entrypoint" in kwargs:
206+
raise ValueError("entrypoint cannot be overridden for an authenticated replica set")
207+
if "user" in kwargs and not _is_root_container_user(kwargs["user"]):
208+
raise ValueError("authenticated replica sets must start as the root container user")
209+
kwargs["entrypoint"] = _REPLICA_SET_ENTRYPOINT_PATH
210+
return super().with_kwargs(**kwargs)
211+
212+
def _configure(self) -> None:
213+
if self.auth_enabled:
214+
assert self.username is not None
215+
assert self.password is not None
216+
self.with_env("MONGO_INITDB_ROOT_USERNAME", self.username)
217+
self.with_env("MONGO_INITDB_ROOT_PASSWORD", self.password)
218+
else:
219+
self.env.pop("MONGO_INITDB_ROOT_USERNAME", None)
220+
self.env.pop("MONGO_INITDB_ROOT_PASSWORD", None)
221+
self.with_env("MONGO_DB", self.dbname)
222+
223+
def get_connection_url(self) -> str:
224+
if self.auth_enabled:
225+
url = super().get_connection_url()
226+
else:
227+
host = self.get_container_host_ip()
228+
port = self.get_exposed_port(self.port)
229+
url = f"mongodb://{host}:{port}"
230+
231+
return f"{url}/?{urlencode({'replicaSet': self.replica_set, 'directConnection': 'true'})}"
232+
233+
def _connect(self) -> None:
234+
direct_url = self.get_connection_url().replace(
235+
urlencode({"replicaSet": self.replica_set, "directConnection": "true"}),
236+
urlencode({"directConnection": "true"}),
237+
)
238+
client: MongoClient[dict[str, Any]] = MongoClient(
239+
direct_url,
240+
serverSelectionTimeoutMS=1000,
241+
connectTimeoutMS=1000,
242+
socketTimeoutMS=1000,
243+
)
244+
deadline = time.monotonic() + testcontainers_config.timeout
245+
246+
try:
247+
self._wait_for_mongodb(client, deadline)
248+
self._wait_for_replica_set_primary(client, deadline)
249+
finally:
250+
client.close()
251+
252+
def _wait_for_mongodb(self, client: MongoClient[dict[str, Any]], deadline: float) -> None:
253+
last_error: Optional[Exception] = None
254+
while time.monotonic() < deadline:
255+
self._raise_if_replica_set_container_stopped()
256+
try:
257+
client.admin.command("ping")
258+
return
259+
except PyMongoError as error:
260+
last_error = error
261+
time.sleep(testcontainers_config.sleep_time)
262+
263+
raise ContainerStartException("MongoDB did not become ready") from last_error
264+
265+
def _wait_for_replica_set_primary(self, client: MongoClient[dict[str, Any]], deadline: float) -> None:
266+
last_error: Optional[Exception] = None
267+
while time.monotonic() < deadline:
268+
self._raise_if_replica_set_container_stopped()
269+
try:
270+
self._initiate_replica_set_if_needed(client)
271+
if client.admin.command("hello").get("isWritablePrimary"):
272+
return
273+
except OperationFailure:
274+
raise
275+
except PyMongoError as error:
276+
last_error = error
277+
time.sleep(testcontainers_config.sleep_time)
278+
279+
raise ContainerStartException("MongoDB replica set did not elect a primary") from last_error
280+
281+
def _initiate_replica_set_if_needed(self, client: MongoClient[dict[str, Any]]) -> None:
282+
try:
283+
client.admin.command("replSetGetStatus")
284+
return
285+
except OperationFailure as error:
286+
if error.code != 94: # NotYetInitialized
287+
raise
288+
289+
try:
290+
client.admin.command(
291+
{
292+
"replSetInitiate": {
293+
"_id": self.replica_set,
294+
"members": [{"_id": 0, "host": f"localhost:{self.port}"}],
295+
}
296+
}
297+
)
298+
except OperationFailure as error:
299+
if error.code != 23: # AlreadyInitialized
300+
raise
301+
302+
def _raise_if_replica_set_container_stopped(self) -> None:
303+
self.reload()
304+
if self.status not in ("exited", "dead"):
305+
return
306+
307+
stdout, stderr = self.get_logs()
308+
logs = (stdout + stderr).decode(errors="replace")
309+
raise ContainerStartException(f"MongoDB stopped while initializing its replica set:\n{logs}")
310+
311+
92312
class MongoDBAtlasLocalContainer(DbContainer):
93313
"""
94314
MongoDB Atlas Local document-based database container.

tests/community/mongodb/test_mongodb.py

Lines changed: 76 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import time
2+
from urllib.parse import parse_qs, urlsplit
23

34
import pytest
45

5-
from testcontainers.community.mongodb import MongoDBAtlasLocalContainer, MongoDbContainer
6+
from testcontainers.community.mongodb import MongoDBAtlasLocalContainer, MongoDbContainer, MongoDbReplicaSetContainer
67

78

89
@pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"])
@@ -28,6 +29,80 @@ def test_docker_run_mongodb(version: str):
2829
assert cursor.next()["restaurant_id"] == doc["restaurant_id"]
2930

3031

32+
@pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"])
33+
def test_docker_run_mongodb_replica_set_without_authentication(version: str):
34+
with MongoDbReplicaSetContainer(f"mongo:{version}", replica_set="docker-rs", auth_enabled=False) as mongo:
35+
connection_url = urlsplit(mongo.get_connection_url())
36+
assert connection_url.username is None
37+
assert parse_qs(connection_url.query) == {
38+
"directConnection": ["true"],
39+
"replicaSet": ["docker-rs"],
40+
}
41+
42+
client = mongo.get_connection_client()
43+
44+
hello = client.admin.command("hello")
45+
assert hello["setName"] == "docker-rs"
46+
assert hello["isWritablePrimary"] is True
47+
48+
collection = client.test.transactions
49+
with client.start_session() as session, session.start_transaction():
50+
result = collection.insert_one({"message": "replica sets support transactions"}, session=session)
51+
52+
assert collection.find_one({"_id": result.inserted_id}) is not None
53+
54+
55+
@pytest.mark.parametrize("version", ["7.0.7", "6.0.14", "5.0.26"])
56+
def test_docker_run_authenticated_mongodb_replica_set(version: str):
57+
container = MongoDbReplicaSetContainer(
58+
f"mongo:{version}",
59+
username="replica-user",
60+
password="replica-password",
61+
replica_set="secure-rs",
62+
command=["--profile", "1"],
63+
).with_kwargs(labels={"testcontainers.mongodb.replica-set": "true"})
64+
65+
with container as mongo:
66+
connection_url = urlsplit(mongo.get_connection_url())
67+
assert connection_url.username == "replica-user"
68+
assert connection_url.password == "replica-password"
69+
assert parse_qs(connection_url.query) == {
70+
"directConnection": ["true"],
71+
"replicaSet": ["secure-rs"],
72+
}
73+
74+
client = mongo.get_connection_client()
75+
76+
hello = client.admin.command("hello")
77+
assert hello["setName"] == "secure-rs"
78+
assert hello["isWritablePrimary"] is True
79+
80+
result = client.test.documents.insert_one({"authenticated": True})
81+
assert client.test.documents.find_one({"_id": result.inserted_id}) == {
82+
"_id": result.inserted_id,
83+
"authenticated": True,
84+
}
85+
86+
87+
@pytest.mark.parametrize("username,password", [("user", None), (None, "password")])
88+
def test_mongodb_rejects_credentials_when_authentication_is_disabled(
89+
username: str | None,
90+
password: str | None,
91+
):
92+
with pytest.raises(ValueError, match="authentication is disabled"):
93+
MongoDbReplicaSetContainer(auth_enabled=False, username=username, password=password)
94+
95+
96+
def test_authenticated_mongodb_replica_set_rejects_non_root_container_user():
97+
with pytest.raises(ValueError, match="root container user"):
98+
MongoDbReplicaSetContainer(replica_set="secure-rs", user="mongodb")
99+
100+
101+
def test_mongodb_replica_set_rejects_non_mongod_command():
102+
with pytest.raises(ValueError, match="mongod options"):
103+
MongoDbReplicaSetContainer(replica_set="secure-rs", command=["bash", "-lc", "exec mongod"])
104+
105+
31106
@pytest.mark.parametrize("version", ["8.0.13", "7.0.23"])
32107
def test_docker_run_mongodb_atlas_local(version: str):
33108
with MongoDBAtlasLocalContainer(f"mongodb/mongodb-atlas-local:{version}") as mongo_atlas:

0 commit comments

Comments
 (0)