feat(nats): Add new module

.devcontainer/commands/post-create-command.sh

@@ -2,4 +2,5 @@ echo "Running post-create-command.sh"
 
 curl -sSL https://install.python-poetry.org | python3 -
 
+poetry lock --no-update
 poetry install --all-extras

INDEX.rst

@@ -27,6 +27,7 @@ testcontainers-python facilitates the use of Docker containers for functional an
     modules/mongodb/README
     modules/mssql/README
     modules/mysql/README
+    modules/nats/README
     modules/neo4j/README
     modules/nginx/README
     modules/opensearch/README

conf.py

@@ -52,7 +52,7 @@
 
 # General information about the project.
 project = "testcontainers"
-copyright = "2017, Sergey Pirogov"  # noqa: A001
+copyright = "2017-2024, Sergey Pirogov and Testcontainers Python contributors"  # noqa: A001
 author = "Sergey Pirogov"
 
 # The version info for the project you're documenting, acts as replacement for

core/testcontainers/core/container.py

@@ -76,6 +76,7 @@ def start(self) -> "DockerContainer":
 
     def stop(self, force=True, delete_volume=True) -> None:
         self._container.remove(force=force, v=delete_volume)
+        self.get_docker_client().client.close()
 
     def __enter__(self) -> "DockerContainer":
         return self.start()

core/testcontainers/core/generic.py

@@ -36,7 +36,10 @@ def _connect(self) -> None:
         import sqlalchemy
 
         engine = sqlalchemy.create_engine(self.get_connection_url())
-        engine.connect()
+        try:
+            engine.connect()
+        finally:
+            engine.dispose()
 
     def get_connection_url(self) -> str:
         raise NotImplementedError

modules/clickhouse/testcontainers/clickhouse/__init__.py

@@ -12,9 +12,8 @@
 #    under the License.
 import os
 from typing import Optional
-
-import clickhouse_driver
-from clickhouse_driver.errors import Error
+from urllib.error import HTTPError, URLError
+from urllib.request import urlopen
 
 from testcontainers.core.generic import DbContainer
 from testcontainers.core.utils import raise_for_deprecated_parameter
@@ -48,7 +47,7 @@ def __init__(
         username: Optional[str] = None,
         password: Optional[str] = None,
         dbname: Optional[str] = None,
-        **kwargs
+        **kwargs,
     ) -> None:
         raise_for_deprecated_parameter(kwargs, "user", "username")
         super().__init__(image=image, **kwargs)
@@ -57,11 +56,14 @@ def __init__(
         self.dbname = dbname or os.environ.get("CLICKHOUSE_DB", "test")
         self.port = port
         self.with_exposed_ports(self.port)
+        self.with_exposed_ports(8123)
 
-    @wait_container_is_ready(Error, EOFError)
+    @wait_container_is_ready(HTTPError, URLError)
     def _connect(self) -> None:
-        with clickhouse_driver.Client.from_url(self.get_connection_url()) as client:
-            client.execute("SELECT version()")
+        # noinspection HttpUrlsUsage
+        url = f"http://{self.get_container_host_ip()}:{self.get_exposed_port(8123)}"
+        with urlopen(url) as r:
+            assert b"Ok" in r.read()
 
     def _configure(self) -> None:
         self.with_env("CLICKHOUSE_USER", self.username)

modules/elasticsearch/tests/test_elasticsearch.py

@@ -6,8 +6,8 @@
 from testcontainers.elasticsearch import ElasticSearchContainer
 
 
-# The versions below were the current supported versions at time of writing (2022-08-11)
-@pytest.mark.parametrize("version", ["6.8.23", "7.17.5", "8.3.3"])
+# The versions below should reflect the latest stable releases
+@pytest.mark.parametrize("version", ["7.17.18", "8.12.2"])
 def test_docker_run_elasticsearch(version):
     with ElasticSearchContainer(f"elasticsearch:{version}", mem_limit="3G") as es:
         resp = urllib.request.urlopen(es.get_url())

modules/mongodb/testcontainers/mongodb/__init__.py

@@ -17,7 +17,7 @@
 
 from testcontainers.core.generic import DbContainer
 from testcontainers.core.utils import raise_for_deprecated_parameter
-from testcontainers.core.waiting_utils import wait_container_is_ready
+from testcontainers.core.waiting_utils import wait_for_logs
 
 
 class MongoDbContainer(DbContainer):
@@ -81,9 +81,8 @@ def get_connection_url(self) -> str:
             port=self.port,
         )
 
-    @wait_container_is_ready()
-    def _connect(self) -> MongoClient:
-        return MongoClient(self.get_connection_url())
+    def _connect(self) -> None:
+        wait_for_logs(self, "Waiting for connections")
 
     def get_connection_client(self) -> MongoClient:
-        return self._connect()
+        return MongoClient(self.get_connection_url())

modules/nats/README.rst

@@ -0,0 +1 @@
+.. autoclass:: testcontainers.nats.NatsContainer

modules/nats/testcontainers/nats/__init__.py

@@ -0,0 +1,65 @@
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
+
+
+class NatsContainer(DockerContainer):
+    """
+    Nats container.
+
+    Example:
+
+        .. doctest::
+
+            >>> from testcontainers.nats import NatsContainer
+
+            >>> with NatsContainer() as nats_container:
+            ...     nc = nats_container.get_client()
+    """
+
+    def __init__(
+        self,
+        image: str = "nats:latest",
+        client_port: int = 4222,
+        management_port: int = 8222,
+        expected_ready_log: str = "Server is ready",
+        ready_timeout_secs: int = 120,
+        **kwargs,
+    ) -> None:
+        super().__init__(image, **kwargs)
+        self.client_port = client_port
+        self.management_port = management_port
+        self._expected_ready_log = expected_ready_log
+        self._ready_timeout_secs = max(ready_timeout_secs, 0)
+        self.with_exposed_ports(self.client_port, self.management_port)
+
+    @wait_container_is_ready()
+    def _healthcheck(self) -> None:
+        wait_for_logs(self, self._expected_ready_log, timeout=self._ready_timeout_secs)
+
+    def nats_uri(self) -> str:
+        return f"nats://{self.get_container_host_ip()}:{self.get_exposed_port(self.client_port)}"
+
+    def nats_host_and_port(self) -> tuple[str, int]:
+        return self.get_container_host_ip(), self.get_exposed_port(self.client_port)
+
+    def nats_management_uri(self) -> str:
+        return f"nats://{self.get_container_host_ip()}:{self.get_exposed_port(self.management_port)}"
+
+    def start(self) -> "NatsContainer":
+        super().start()
+        self._healthcheck()
+        return self

modules/nats/tests/test_nats.py

@@ -0,0 +1,101 @@
+from testcontainers.nats import NatsContainer
+from uuid import uuid4
+import pytest
+
+
+"""
+If you are developing this and you want to test more advanced scenarios using a client
+Activate your poetry shell.
+pip install nats-py
+This will get nats-py into your environment but keep it out of the project
+
+
+"""
+
+
+NO_NATS_CLIENT = True
+try:
+    from nats import connect as nats_connect
+    from nats.aio.client import Client as NATSClient
+
+    NO_NATS_CLIENT = False
+except ImportError:
+    pass
+
+
+async def get_client(container: NatsContainer) -> "NATSClient":
+    """
+    Get a nats client.
+
+    Returns:
+        client: Nats client to connect to the container.
+    """
+    conn_string = container.nats_uri()
+    client = await nats_connect(conn_string)
+    return client
+
+
+def test_basic_container_ops():
+    with NatsContainer() as container:
+        # Not sure how to get type information without doing this
+        container: NatsContainer = container
+        h, p = container.nats_host_and_port()
+        assert h == "localhost"
+        uri = container.nats_uri()
+        management_uri = container.nats_management_uri()
+
+        assert uri != management_uri
+
+
+pytest.mark.usefixtures("anyio_backend")
+
+
+@pytest.mark.skipif(NO_NATS_CLIENT, reason="No NATS Client Available")
+@pytest.mark.parametrize("anyio_backend", ["asyncio"])
+async def test_pubsub(anyio_backend):
+    with NatsContainer() as container:
+        nc: NATSClient = await get_client(container)
+
+        topic = str(uuid4())
+
+        sub = await nc.subscribe(topic)
+        sent_message = b"Test-Containers"
+        await nc.publish(topic, b"Test-Containers")
+        received_msg = await sub.next_msg()
+        print("Received:", received_msg)
+        assert sent_message == received_msg.data
+        await nc.flush()
+        await nc.close()
+
+
+pytest.mark.usefixtures("anyio_backend")
+
+
+@pytest.mark.parametrize("anyio_backend", ["asyncio"])
+@pytest.mark.skipif(NO_NATS_CLIENT, reason="No NATS Client Available")
+async def test_more_complex_example(anyio_backend):
+    with NatsContainer() as container:
+        nc: NATSClient = await get_client(container)
+
+        await nc.publish("greet.joe", b"hello")
+
+        sub = await nc.subscribe("greet.*")
+
+        try:
+            await sub.next_msg(timeout=0.1)
+        except TimeoutError:
+            pass
+
+        await nc.publish("greet.joe", b"hello.joe")
+        await nc.publish("greet.pam", b"hello.pam")
+
+        first = await sub.next_msg(timeout=0.1)
+        assert b"hello.joe" == first.data
+
+        second = await sub.next_msg(timeout=0.1)
+        assert b"hello.pam" == second.data
+
+        await nc.publish("greet.bob", b"hello")
+
+        await sub.unsubscribe()
+        await nc.drain()

pyproject.toml

@@ -41,6 +41,7 @@ packages = [
     { include = "testcontainers", from = "modules/mongodb" },
     { include = "testcontainers", from = "modules/mssql" },
     { include = "testcontainers", from = "modules/mysql" },
+    { include = "testcontainers", from = "modules/nats" },
     { include = "testcontainers", from = "modules/neo4j" },
     { include = "testcontainers", from = "modules/nginx" },
     { include = "testcontainers", from = "modules/opensearch" },
@@ -83,6 +84,7 @@ psycopg2-binary = { version = "*", optional = true }
 pika = { version = "*", optional = true }
 redis = { version = "*", optional = true }
 selenium = { version = "*", optional = true }
+nats-py = { version = "*", optional = true }
 
 [tool.poetry.extras]
 arangodb = ["python-arango"]
@@ -98,6 +100,7 @@ minio = ["minio"]
 mongodb = ["pymongo"]
 mssql = ["sqlalchemy", "pymssql"]
 mysql = ["sqlalchemy", "pymysql"]
+nats = []
 neo4j = ["neo4j"]
 nginx = []
 opensearch = ["opensearch-py"]
@@ -116,6 +119,7 @@ pytest-cov = "4.1.0"
 sphinx = "^7.2.6"
 twine = "^4.0.2"
 anyio = "^4.3.0"
+pytest-asyncio = "^0.23.5"
 
 [[tool.poetry.source]]
 name = "PyPI"
@@ -222,6 +226,7 @@ mypy_path = [
 #    "modules/mongodb",
 #    "modules/mssql",
 #    "modules/mysql",
+#    "modules/nats",
 #    "modules/neo4j",
 #    "modules/nginx",
 #    "modules/opensearch",