using testcontainers module

pyinfra/locations.py

@@ -12,3 +12,5 @@ TEST_DIR = MODULE_DIR / "test"
 CONFIG_FILE = PACKAGE_ROOT_DIR / "config.yaml"
 
 TEST_CONFIG_FILE = TEST_DIR / "config.yaml"
+
+COMPOSE_PATH = PACKAGE_ROOT_DIR

pyinfra/test/unit_tests/conftest.py

@@ -1,11 +1,14 @@
 import json
 import logging
+import time
 from unittest.mock import Mock
 
+import pika
 import pytest
+import testcontainers.compose
 
 from pyinfra.exceptions import UnknownClient
-from pyinfra.locations import TEST_DIR
+from pyinfra.locations import TEST_DIR, COMPOSE_PATH
 from pyinfra.queue.queue_manager.pika_queue_manager import PikaQueueManager, get_connection_params
 from pyinfra.queue.queue_manager.queue_manager import QueueManager
 from pyinfra.storage.adapters.azure import AzureStorageAdapter
@@ -67,6 +70,17 @@ def storage(client_name, bucket_name, request):
     storage.clear_bucket(bucket_name)
 
 
+@pytest.fixture(scope="session", autouse=True)
+def docker_compose(sleep_seconds=30):
+    logger.info(f"Starting docker containers with {COMPOSE_PATH}/docker-compose.yml...")
+    compose = testcontainers.compose.DockerCompose(COMPOSE_PATH)
+    compose.start()
+    logger.info(f"Sleeping for {sleep_seconds} seconds to wait for containers to finish startup... ")
+    time.sleep(sleep_seconds)
+    yield compose
+    compose.stop()
+
+
 def get_endpoint_from_running_container(container_name, config):
     port = config.endpoint.split(":")[-1]
     endpoint = f"{docker.get_ip_address(container_name)}:{port}"
@@ -112,7 +126,10 @@ def queue_manager(queue_manager_name):
     queue_manager = get_queue_manager(queue_manager_name)
     yield queue_manager
     if queue_manager_name == "pika":
-        queue_manager.connection.close()
+        try:
+            queue_manager.connection.close()
+        except (pika.exceptions.StreamLostError, ConnectionResetError):
+            pass
 
 
 @pytest.fixture(scope="session")