From f667bc72ddd9bc3dcd0be85ad4c8b70d119db9bd Mon Sep 17 00:00:00 2001
From: Matthias Bisping <matthias.bisping@iqser.com>
Date: Tue, 22 Mar 2022 14:37:47 +0100
Subject: [PATCH] using testcontainers module

---
 docker-compose.yaml => docker-compose.yml |  0
 pyinfra/locations.py                      |  2 ++
 pyinfra/test/unit_tests/conftest.py       | 21 +++++++++++++++++++--
 3 files changed, 21 insertions(+), 2 deletions(-)
 rename docker-compose.yaml => docker-compose.yml (100%)

diff --git a/docker-compose.yaml b/docker-compose.yml
similarity index 100%
rename from docker-compose.yaml
rename to docker-compose.yml
diff --git a/pyinfra/locations.py b/pyinfra/locations.py
index 5183a8b..a1878da 100644
--- a/pyinfra/locations.py
+++ b/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
diff --git a/pyinfra/test/unit_tests/conftest.py b/pyinfra/test/unit_tests/conftest.py
index 9a7b67d..b6510a9 100644
--- a/pyinfra/test/unit_tests/conftest.py
+++ b/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")