From 5dad0e7e20e99fa610668bdb9e7fbcb994447406 Mon Sep 17 00:00:00 2001
From: Balint Bartha <39852431+totallyzen@users.noreply.github.com>
Date: Mon, 18 Nov 2024 11:14:37 +0100
Subject: [PATCH] feat(compose): support for setting profiles

---
 core/testcontainers/compose/compose.py        |  4 ++++
 .../profile_support/compose.yaml              | 16 +++++++++++++
 core/tests/test_compose.py                    | 24 +++++++++++++++++++
 3 files changed, 44 insertions(+)
 create mode 100644 core/tests/compose_fixtures/profile_support/compose.yaml

diff --git a/core/testcontainers/compose/compose.py b/core/testcontainers/compose/compose.py
index 564eda8f..3f543818 100644
--- a/core/testcontainers/compose/compose.py
+++ b/core/testcontainers/compose/compose.py
@@ -171,6 +171,7 @@ class DockerCompose:
     env_file: Optional[str] = None
     services: Optional[list[str]] = None
     docker_command_path: Optional[str] = None
+    profiles: Optional[list[str]] = None
 
     def __post_init__(self):
         if isinstance(self.compose_file_name, str):
@@ -225,6 +226,9 @@ def start(self) -> None:
             # we run in detached mode instead of blocking
             up_cmd.append("--detach")
 
+        if self.profiles:
+            up_cmd.extend([item for profile in self.profiles for item in ["--profile", profile]])
+
         if self.services:
             up_cmd.extend(self.services)
 
diff --git a/core/tests/compose_fixtures/profile_support/compose.yaml b/core/tests/compose_fixtures/profile_support/compose.yaml
new file mode 100644
index 00000000..c7bec7cc
--- /dev/null
+++ b/core/tests/compose_fixtures/profile_support/compose.yaml
@@ -0,0 +1,16 @@
+services:
+  runs-always: &simple-service
+    image: alpine:latest
+    init: true
+    command:
+      - sh
+      - -c
+      - 'while true; do sleep 0.1 ; date -Ins; done'
+  runs-profile-a:
+    <<: *simple-service
+    profiles:
+      - profile-a
+  runs-profile-b:
+    <<: *simple-service
+    profiles:
+      - profile-b
diff --git a/core/tests/test_compose.py b/core/tests/test_compose.py
index b43da28c..088934ac 100644
--- a/core/tests/test_compose.py
+++ b/core/tests/test_compose.py
@@ -352,3 +352,27 @@ def fetch(req: Union[Request, str]):
         if 200 < res.getcode() >= 400:
             raise Exception(f"HTTP Error: {res.getcode()} - {res.reason}: {body}")
         return res.getcode(), body
+
+
+@pytest.mark.parametrize(
+    argnames=["profiles", "running", "not_running"],
+    argvalues=[
+        pytest.param(None, ["runs-always"], ["runs-profile-a", "runs-profile-b"], id="default"),
+        pytest.param(
+            ["profile-a"], ["runs-always", "runs-profile-a"], ["runs-profile-b"], id="one-additional-profile-via-str"
+        ),
+        pytest.param(
+            ["profile-a", "profile-b"],
+            ["runs-always", "runs-profile-a", "runs-profile-b"],
+            [],
+            id="all-profiles-explicitly",
+        ),
+    ],
+)
+def test_compose_profile_support(profiles: list[str] | None, running: list[str], not_running: list[str]):
+    with DockerCompose(context=FIXTURES / "profile_support", profiles=profiles) as compose:
+        for service in running:
+            assert compose.get_container(service) is not None
+        for service in not_running:
+            with pytest.raises(ContainerIsNotRunning):
+                compose.get_container(service)