Skip to content

Commit 302c73d

Browse files
bearritoalexanderankin
bearrito
and
alexanderankin
authored
fix(nats): Client-Free(ish) NATS container (testcontainers#462)
Co-authored-by: bstrausser <bstrausser@locusrobotics.com> Co-authored-by: David Ankin <daveankin@gmail.com>
1 parent e8876f4 commit 302c73d

6 files changed

Lines changed: 205 additions & 2 deletions

File tree

‎index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ testcontainers-python facilitates the use of Docker containers for functional an
3131
modules/mongodb/README
3232
modules/mssql/README
3333
modules/mysql/README
34+
modules/nats/README
3435
modules/neo4j/README
3536
modules/nginx/README
3637
modules/opensearch/README

‎modules/nats/README.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
.. autoclass:: testcontainers.nats.NatsContainer
2+
.. title:: testcontainers.nats.NatsContainer

‎modules/nats/testcontainers/nats/__init__.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
#
2+
# Licensed under the Apache License, Version 2.0 (the "License"); you may
3+
# not use this file except in compliance with the License. You may obtain
4+
# a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
10+
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
11+
# License for the specific language governing permissions and limitations
12+
# under the License.
13+
14+
15+
from testcontainers.core.container import DockerContainer
16+
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs
17+
18+
19+
class NatsContainer(DockerContainer):
20+
"""
21+
Nats container.
22+
23+
Example:
24+
25+
.. doctest::
26+
27+
>>> import asyncio
28+
>>> from nats import connect as nats_connect
29+
>>> from testcontainers.nats import NatsContainer
30+
31+
>>> async def test_doctest_usage():
32+
... with NatsContainer() as nats_container:
33+
... client = await nats_connect(nats_container.nats_uri())
34+
... sub_tc = await client.subscribe("tc")
35+
... await client.publish("tc", b"Test-Containers")
36+
... next_message = await sub_tc.next_msg(timeout=5.0)
37+
... await client.close()
38+
... return next_message.data
39+
>>> asyncio.run(test_doctest_usage())
40+
b'Test-Containers'
41+
"""
42+
43+
def __init__(
44+
self,
45+
image: str = "nats:latest",
46+
client_port: int = 4222,
47+
management_port: int = 8222,
48+
expected_ready_log: str = "Server is ready",
49+
ready_timeout_secs: int = 120,
50+
**kwargs,
51+
) -> None:
52+
super().__init__(image, **kwargs)
53+
self.client_port = client_port
54+
self.management_port = management_port
55+
self._expected_ready_log = expected_ready_log
56+
self._ready_timeout_secs = max(ready_timeout_secs, 0)
57+
self.with_exposed_ports(self.client_port, self.management_port)
58+
59+
@wait_container_is_ready()
60+
def _healthcheck(self) -> None:
61+
wait_for_logs(self, self._expected_ready_log, timeout=self._ready_timeout_secs)
62+
63+
def nats_uri(self) -> str:
64+
return f"nats://{self.get_container_host_ip()}:{self.get_exposed_port(self.client_port)}"
65+
66+
def nats_host_and_port(self) -> tuple[str, int]:
67+
return self.get_container_host_ip(), self.get_exposed_port(self.client_port)
68+
69+
def nats_management_uri(self) -> str:
70+
return f"nats://{self.get_container_host_ip()}:{self.get_exposed_port(self.management_port)}"
71+
72+
def start(self) -> "NatsContainer":
73+
super().start()
74+
self._healthcheck()
75+
return self

‎modules/nats/tests/test_nats.py‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
from testcontainers.nats import NatsContainer
2+
from uuid import uuid4
3+
import pytest
4+
5+
from nats import connect as nats_connect
6+
from nats.aio.client import Client as NATSClient
7+
8+
9+
async def get_client(container: NatsContainer) -> "NATSClient":
10+
"""
11+
Get a nats client.
12+
13+
Returns:
14+
client: Nats client to connect to the container.
15+
"""
16+
conn_string = container.nats_uri()
17+
client = await nats_connect(conn_string)
18+
return client
19+
20+
21+
def test_basic_container_ops():
22+
with NatsContainer() as container:
23+
# Not sure how to get type information without doing this
24+
container: NatsContainer = container
25+
h, p = container.nats_host_and_port()
26+
assert h == "localhost"
27+
uri = container.nats_uri()
28+
management_uri = container.nats_management_uri()
29+
30+
assert uri != management_uri
31+
32+
33+
@pytest.mark.asyncio
34+
async def test_pubsub(anyio_backend):
35+
with NatsContainer() as container:
36+
nc: NATSClient = await get_client(container)
37+
38+
topic = str(uuid4())
39+
40+
sub = await nc.subscribe(topic)
41+
sent_message = b"Test-Containers"
42+
await nc.publish(topic, b"Test-Containers")
43+
received_msg = await sub.next_msg()
44+
print("Received:", received_msg)
45+
assert sent_message == received_msg.data
46+
await nc.flush()
47+
await nc.close()
48+
49+
50+
@pytest.mark.asyncio
51+
async def test_more_complex_example(anyio_backend):
52+
with NatsContainer() as container:
53+
nc: NATSClient = await get_client(container)
54+
55+
sub = await nc.subscribe("greet.*")
56+
await nc.publish("greet.joe", b"hello")
57+
58+
try:
59+
await sub.next_msg(timeout=0.1)
60+
except TimeoutError:
61+
pass
62+
63+
await nc.publish("greet.joe", b"hello.joe")
64+
await nc.publish("greet.pam", b"hello.pam")
65+
66+
first = await sub.next_msg(timeout=0.1)
67+
assert b"hello.joe" == first.data
68+
69+
second = await sub.next_msg(timeout=0.1)
70+
assert b"hello.pam" == second.data
71+
72+
await nc.publish("greet.bob", b"hello")
73+
74+
await sub.unsubscribe()
75+
await nc.drain()
76+
77+
78+
@pytest.mark.asyncio
79+
async def test_doctest_usage():
80+
"""simpler to run test to mirror what is in the doctest"""
81+
with NatsContainer() as nats_container:
82+
client = await nats_connect(nats_container.nats_uri())
83+
sub_tc = await client.subscribe("tc")
84+
await client.publish("tc", b"Test-Containers")
85+
next_message = await sub_tc.next_msg(timeout=5.0)
86+
await client.close()
87+
assert next_message.data == b"Test-Containers"

‎poetry.lock‎

Lines changed: 36 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pyproject.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ packages = [
4444
{ include = "testcontainers", from = "modules/mongodb" },
4545
{ include = "testcontainers", from = "modules/mssql" },
4646
{ include = "testcontainers", from = "modules/mysql" },
47+
{ include = "testcontainers", from = "modules/nats" },
4748
{ include = "testcontainers", from = "modules/neo4j" },
4849
{ include = "testcontainers", from = "modules/nginx" },
4950
{ include = "testcontainers", from = "modules/opensearch" },
@@ -79,6 +80,7 @@ pyyaml = { version = "*", optional = true }
7980
python-keycloak = { version = "*", optional = true }
8081
boto3 = { version = "*", optional = true }
8182
minio = { version = "*", optional = true }
83+
nats-py = { version = "*", optional = true }
8284
pymongo = { version = "*", optional = true }
8385
sqlalchemy = { version = "*", optional = true }
8486
pymssql = { version = "*", optional = true }
@@ -109,6 +111,7 @@ minio = ["minio"]
109111
mongodb = ["pymongo"]
110112
mssql = ["sqlalchemy", "pymssql"]
111113
mysql = ["sqlalchemy", "pymysql"]
114+
nats = ["nats-py"]
112115
neo4j = ["neo4j"]
113116
nginx = []
114117
opensearch = ["opensearch-py"]
@@ -136,6 +139,7 @@ sqlalchemy = "*"
136139
psycopg = "*"
137140
kafka-python = "^2.0.2"
138141
cassandra-driver = "*"
142+
pytest-asyncio = "0.23.5"
139143

140144
[[tool.poetry.source]]
141145
name = "PyPI"

0 commit comments

Comments
 (0)