From fa2081a7b325cdd316de28c99b029150022db203 Mon Sep 17 00:00:00 2001 From: Morgan Trench Date: Wed, 31 Jul 2024 21:14:15 +0800 Subject: [PATCH] fix(rabbitmq): add `vhost` as parameter to RabbitMqContainer (#656) Adds a `vhost` parameter to the RabbitMQContainer constructor that allows the `RABBITMQ_DEFAULT_VHOST` [environment variable](https://www.rabbitmq.com/docs/configure#supported-environment-variables) to be modified. Subsequently `vhost` is then also used inside the `get_connection_params` method for the `pika` connection parameters, which is used to test if the container is ready. --- .../rabbitmq/testcontainers/rabbitmq/__init__.py | 4 ++++ modules/rabbitmq/tests/test_rabbitmq.py | 15 ++++++++++----- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/modules/rabbitmq/testcontainers/rabbitmq/__init__.py b/modules/rabbitmq/testcontainers/rabbitmq/__init__.py index 3e5ad0b3..cd7b079a 100644 --- a/modules/rabbitmq/testcontainers/rabbitmq/__init__.py +++ b/modules/rabbitmq/testcontainers/rabbitmq/__init__.py @@ -31,6 +31,7 @@ def __init__( port: Optional[int] = None, username: Optional[str] = None, password: Optional[str] = None, + vhost: Optional[str] = None, **kwargs, ) -> None: """Initialize the RabbitMQ test container. @@ -45,11 +46,13 @@ def __init__( self.port = port or int(os.environ.get("RABBITMQ_NODE_PORT", 5672)) self.username = username or os.environ.get("RABBITMQ_DEFAULT_USER", "guest") self.password = password or os.environ.get("RABBITMQ_DEFAULT_PASS", "guest") + self.vhost = vhost or os.environ.get("RABBITMQ_DEFAULT_VHOST", "/") self.with_exposed_ports(self.port) self.with_env("RABBITMQ_NODE_PORT", self.port) self.with_env("RABBITMQ_DEFAULT_USER", self.username) self.with_env("RABBITMQ_DEFAULT_PASS", self.password) + self.with_env("RABBITMQ_DEFAULT_VHOST", self.vhost) @wait_container_is_ready(pika.exceptions.IncompatibleProtocolError, pika.exceptions.AMQPConnectionError) def readiness_probe(self) -> bool: @@ -71,6 +74,7 @@ def get_connection_params(self) -> pika.ConnectionParameters: return pika.ConnectionParameters( host=self.get_container_host_ip(), port=self.get_exposed_port(self.port), + virtual_host=self.vhost, credentials=credentials, ) diff --git a/modules/rabbitmq/tests/test_rabbitmq.py b/modules/rabbitmq/tests/test_rabbitmq.py index 98fb7e6d..c0c1894d 100644 --- a/modules/rabbitmq/tests/test_rabbitmq.py +++ b/modules/rabbitmq/tests/test_rabbitmq.py @@ -13,14 +13,17 @@ @pytest.mark.parametrize( - argnames=["port", "username", "password"], + argnames=["port", "username", "password", "vhost"], argvalues=[ - [None, None, None], # use the defaults - [5673, None, None], # test with custom port - [None, "my_test_user", "my_secret_password"], # test with custom credentials + [None, None, None, None], # use the defaults + [5673, None, None, None], # test with custom port + [None, "my_test_user", "my_secret_password", None], # test with custom credentials + [None, None, None, "vhost"], # test with custom vhost ], ) -def test_docker_run_rabbitmq(port: Optional[int], username: Optional[str], password: Optional[str]): +def test_docker_run_rabbitmq( + port: Optional[int], username: Optional[str], password: Optional[str], vhost: Optional[str] +): """Run rabbitmq test container and use it to deliver a simple message.""" kwargs = {} if port is not None: @@ -29,6 +32,8 @@ def test_docker_run_rabbitmq(port: Optional[int], username: Optional[str], passw kwargs["username"] = username if password is not None: kwargs["password"] = password + if vhost is not None: + kwargs["vhost"] = vhost rabbitmq_container = RabbitMqContainer("rabbitmq:latest", **kwargs) with rabbitmq_container as rabbitmq: