Add functionality to wait till mongo container becomes available

testcontainers/core/waiting_utils.py

@@ -17,9 +17,12 @@
 
 import blindspin
 import crayons
+import requests
 import wrapt
+from requests import RequestException
 
 from testcontainers.core import config
+from testcontainers.core.container import DockerContainer
 from testcontainers.core.exceptions import TimeoutException
 
 
@@ -90,3 +93,101 @@ def wait_for_logs(container, predicate, timeout=None, interval=1):
             raise TimeoutError("container did not emit logs satisfying predicate in %.3f seconds"
                                % timeout)
         time.sleep(interval)
+
+
+def wait_for_port(container: DockerContainer, port: int, timeout=None, interval=1):
+    """
+    Wait for the container port to be available.
+
+    Parameters
+    ----------
+    container : DockerContainer
+        Container whose port to wait for.
+    port : int
+        Port to check against.
+    timeout : float or None
+        Number of seconds to wait for the port to be open. Defaults to wait indefinitely.
+    interval : float
+        Interval at which to poll the port.
+
+    Returns
+    -------
+    duration : float
+        Number of seconds until the check passed.
+
+    """
+    start = time.time()
+    while True:
+        duration = time.time() - start
+
+        # build command
+        parts = ["/bin/sh", "-c", "\"true", "&&", "(",
+                 "cat", "/proc/net/tcp*", "|", "awk", "'{print $2}'", "|", "grep", "-i", "':0*%x'" % port,
+                 "||", "nc", "-vz", "-w", "1", "localhost", "%d" % port,
+                 "||", "/bin/bash", "-c", "'</dev/tcp/localhost/%d'" % port, ")\""]
+
+        cmd = ' '.join(parts)
+
+        res = container._container.exec_run(cmd)
+
+        if res.exit_code == 0:
+            return duration
+        if timeout and duration > timeout:
+            raise TimeoutError("container did not start listening on port %d in %.3f seconds"
+                               % (port, timeout))
+        time.sleep(interval)
+
+
+def wait_for_http_code(container: DockerContainer, status_code: int, port: int = 80, path: str = '/',
+                       scheme: str = 'http', timeout=None, interval=1, request_kwargs: dict = None):
+    """
+    Wait for a specific http status code.
+
+    Parameters
+    ----------
+    container : DockerContainer
+        Container which is queried for a specific status code.
+    status_code : int
+        Status code to wait for.
+    port : int
+        Port to query request on.
+    path : str
+        Path to use for request. Default is '/'
+    scheme : str
+        Scheme to use in request query. Default is 'http'
+    timeout : float or None
+        Number of seconds to wait for the port to be open. Defaults to wait indefinitely.
+    interval : float
+        Interval at which to poll the port.
+    request_kwargs: dict
+        kwargs to pass into the request, e.g.: {'verify': False}
+
+    Returns
+    -------
+    duration : float
+        Number of seconds until the check passed.
+    """
+    if request_kwargs is None:
+        request_kwargs = {'timeout': 1.0}
+    elif 'timeout' not in request_kwargs:
+        request_kwargs['timeout'] = 1.0
+
+    start = time.time()
+    # wait for port to open before continuing with http check
+    wait_for_port(container, port, timeout, interval)
+    while True:
+        duration = time.time() - start
+        dest = "%s://%s:%d%s" % (scheme,
+                                 container.get_container_host_ip(),
+                                 int(container.get_exposed_port(port)),
+                                 path)
+        res = None
+        try:
+            res = requests.get(dest, **request_kwargs)
+        except RequestException:
+            pass
+        if res and res.status_code == status_code:
+            return duration
+        if timeout and duration > timeout:
+            raise TimeoutError("container did not respond with %d listening on port %d in %.3f seconds"
+                               % (status_code, port, timeout))

testcontainers/mongodb.py

@@ -13,6 +13,7 @@
 import os
 
 from testcontainers.core.generic import DockerContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready
 
 
 class MongoDbContainer(DockerContainer):
@@ -69,3 +70,14 @@ def get_connection_client(self):
         from pymongo import MongoClient
         return MongoClient("mongodb://{}:{}".format(self.get_container_host_ip(),
                                                     self.get_exposed_port(self.port_to_expose)))
+
+    @wait_container_is_ready()
+    def _connect(self):
+        client = self.get_connection_client()
+        # will raise pymongo.errors.ServerSelectionTimeoutError if no connection is established
+        client.admin.command('ismaster')
+
+    def start(self):
+        super().start()
+        self._connect()
+        return self

tests/test_core.py

@@ -1,7 +1,7 @@
 import pytest
 
 from testcontainers.core.container import DockerContainer
-from testcontainers.core.waiting_utils import wait_for_logs
+from testcontainers.core.waiting_utils import wait_for_logs, wait_for_port, wait_for_http_code
 
 
 def test_raise_timeout():
@@ -13,3 +13,31 @@ def test_raise_timeout():
 def test_wait_for_hello():
     with DockerContainer("hello-world") as container:
         wait_for_logs(container, "Hello from Docker!")
+
+
+def test_raise_timeout_port():
+    with pytest.raises(TimeoutError):
+        with DockerContainer("alpine").with_command("nc -l -p 81").with_exposed_ports(80) as container:
+            wait_for_port(container, 80, timeout=1)
+
+
+def test_wait_for_port():
+    with DockerContainer("alpine").with_command("nc -l -p 80").with_exposed_ports(80) as container:
+        wait_for_port(container, 80)
+
+
+def test_wait_for_http_code():
+    with DockerContainer("nginx").with_exposed_ports(80) as container:
+        wait_for_http_code(container, 200)
+
+
+def test_raise_timeout_for_http_code():
+    with pytest.raises(TimeoutError):
+        with DockerContainer("nginx").with_exposed_ports(81) as container:
+            wait_for_http_code(container, 200, port=81, timeout=2)
+
+
+def test_raise_timeout_invalid_path():
+    with pytest.raises(TimeoutError):
+        with DockerContainer("nginx").with_exposed_ports(80) as container:
+            wait_for_http_code(container, 200, path="/fail", port=80, timeout=2)