diff --git a/INDEX.rst b/INDEX.rst index be5e3d1cd..dfaa34b8f 100644 --- a/INDEX.rst +++ b/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 diff --git a/modules/nats/README.rst b/modules/nats/README.rst new file mode 100644 index 000000000..a38a21466 --- /dev/null +++ b/modules/nats/README.rst @@ -0,0 +1 @@ +.. autoclass:: testcontainers.nats.NatsContainer diff --git a/modules/nats/testcontainers/nats/__init__.py b/modules/nats/testcontainers/nats/__init__.py new file mode 100644 index 000000000..5b6b26dd7 --- /dev/null +++ b/modules/nats/testcontainers/nats/__init__.py @@ -0,0 +1,75 @@ +# +# 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 nats import connect as nats_connect +from nats.aio.client import Client as NATSClient +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 get_conn_string(self): + return f"nats://{self.get_container_host_ip()}:{self.get_exposed_port(self.client_port)}" + + async def get_client(self, **kwargs) -> NATSClient: + """ + Get a nats client. + + Args: + **kwargs: Keyword arguments passed to `redis.Redis`. + + Returns: + client: Nats client to connect to the container. + """ + conn_string = self.get_conn_string() + client = await nats_connect(conn_string) + return client + + def start(self) -> "NatsContainer": + super().start() + self._healthcheck() + return self diff --git a/modules/nats/tests/test_nats.py b/modules/nats/tests/test_nats.py new file mode 100644 index 000000000..9d464cb12 --- /dev/null +++ b/modules/nats/tests/test_nats.py @@ -0,0 +1,52 @@ +from uuid import uuid4 + +import pytest +from nats.aio.client import Client as NATSClient + +from testcontainers.nats import NatsContainer + + +@pytest.mark.asyncio +async def test_basic_publishing(): + with NatsContainer() as container: + nc: NATSClient = await container.get_client() + + 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.asyncio +async def test_more_complex_example(): + with NatsContainer() as container: + nc: NATSClient = await container.get_client() + + 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() diff --git a/poetry.lock b/poetry.lock index f97690266..2ef48d2a1 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. [[package]] name = "alabaster" @@ -1491,6 +1491,21 @@ files = [ {file = "mypy_extensions-1.0.0.tar.gz", hash = "sha256:75dbf8955dc00442a438fc4d0666508a9a97b6bd41aa2f0ffe9d2f2725af0782"}, ] +[[package]] +name = "nats-py" +version = "2.7.2" +description = "NATS client for Python" +optional = true +python-versions = ">=3.7" +files = [ + {file = "nats-py-2.7.2.tar.gz", hash = "sha256:0c97b4a57bed0ef1ff9ae6c19bc115ec7ca8ede5ab3e001fd00a377056a547cf"}, +] + +[package.extras] +aiohttp = ["aiohttp"] +fast-parse = ["fast-mail-parser"] +nkeys = ["nkeys"] + [[package]] name = "neo4j" version = "5.16.0" @@ -2165,6 +2180,24 @@ tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""} [package.extras] testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"] +[[package]] +name = "pytest-asyncio" +version = "0.23.5" +description = "Pytest support for asyncio" +optional = false +python-versions = ">=3.8" +files = [ + {file = "pytest-asyncio-0.23.5.tar.gz", hash = "sha256:3a048872a9c4ba14c3e90cc1aa20cbc2def7d01c7c8db3777ec281ba9c057675"}, + {file = "pytest_asyncio-0.23.5-py3-none-any.whl", hash = "sha256:4e7093259ba018d58ede7d5315131d21923a60f8a6e9ee266ce1589685c89eac"}, +] + +[package.dependencies] +pytest = ">=7.0.0,<9" + +[package.extras] +docs = ["sphinx (>=5.3)", "sphinx-rtd-theme (>=1.0)"] +testing = ["coverage (>=6.2)", "hypothesis (>=5.7.1)"] + [[package]] name = "pytest-cov" version = "4.1.0" @@ -3147,6 +3180,7 @@ minio = ["minio"] mongodb = ["pymongo"] mssql = ["pymssql", "sqlalchemy"] mysql = ["pymysql", "sqlalchemy"] +nats = ["nats-py"] neo4j = ["neo4j"] nginx = [] opensearch = ["opensearch-py"] @@ -3159,4 +3193,4 @@ selenium = ["selenium"] [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0" -content-hash = "f4cb027301e265217ccb581b0ddd06fe6d91319fbcfbc3d20504a1fdbc45d7b1" +content-hash = "9a546ef8b1e1509a5d98576b207b05ec58fd29f5cc4c845cb4262116ee76a30e" diff --git a/pyproject.toml b/pyproject.toml index 7afb4cd96..326170809 100644 --- a/pyproject.toml +++ b/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 = ["nats-py"] 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",