From 628773291f9edbdb787dc42d19d9794313acc261 Mon Sep 17 00:00:00 2001
From: rhoban <rhoban@bloomberg.net>
Date: Sun, 1 Jun 2025 11:45:04 -0400
Subject: [PATCH] Introduce wait_for_healthcheck

---
 core/testcontainers/core/container.py | 16 +++++++
 core/tests/test_container.py          | 69 ++++++++++++++++++++++++++-
 2 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/core/testcontainers/core/container.py b/core/testcontainers/core/container.py
index b7979a61..b92fa985 100644
--- a/core/testcontainers/core/container.py
+++ b/core/testcontainers/core/container.py
@@ -1,4 +1,5 @@
 import contextlib
+import time
 from os import PathLike
 from socket import socket
 from typing import TYPE_CHECKING, Optional, Union
@@ -251,6 +252,17 @@ def _configure(self) -> None:
         # placeholder if subclasses want to define this and use the default start method
         pass
 
+    def wait_for_healthcheck(self, timeout: int = 10):
+        start_time = time.time()
+        underlying = self.get_wrapped_container()
+        while time.time() - start_time < timeout:
+            underlying.reload()
+            if underlying.health == "healthy":
+                break
+            time.sleep(0.1)
+        else:
+            raise NotHealthy()
+
 
 class Reaper:
     _instance: "Optional[Reaper]" = None
@@ -327,3 +339,7 @@ def _create_instance(cls) -> "Reaper":
         Reaper._instance = Reaper()
 
         return Reaper._instance
+
+
+class NotHealthy(Exception):
+    pass
diff --git a/core/tests/test_container.py b/core/tests/test_container.py
index bb7dd059..a6cc71a0 100644
--- a/core/tests/test_container.py
+++ b/core/tests/test_container.py
@@ -1,8 +1,12 @@
+import pathlib
+import textwrap
+
 import pytest
 
-from testcontainers.core.container import DockerContainer
+from testcontainers.core.container import DockerContainer, NotHealthy
 from testcontainers.core.docker_client import DockerClient
 from testcontainers.core.config import ConnectionMode
+from testcontainers.core.image import DockerImage
 
 FAKE_ID = "ABC123"
 
@@ -96,3 +100,66 @@ def test_attribute(init_attr, init_value, class_attr, stored_value):
     """Test that the attributes set through the __init__ function are properly stored."""
     with DockerContainer("ubuntu", **{init_attr: init_value}) as container:
         assert getattr(container, class_attr) == stored_value
+
+
+@pytest.fixture(name="with_never_healthy_image")
+def with_never_health_image_fixture(tmp_path):
+    DOCKERFILE = """FROM alpine:latest
+    HEALTHCHECK --interval=1s CMD test -e /testfile
+    CMD sleep infinity
+    """
+    dockerfile = tmp_path / "Dockerfile"
+    dockerfile.write_text(textwrap.dedent(DOCKERFILE))
+    with DockerImage(tmp_path) as image:
+        yield image
+
+
+def test_wait_for_healthcheck_never_healthy(with_never_healthy_image: DockerImage):
+    # Given
+    with DockerContainer(image=str(with_never_healthy_image)) as container:
+        # Expect
+        with pytest.raises(NotHealthy):
+            # When
+            container.wait_for_healthcheck()
+
+
+@pytest.fixture(name="with_immediately_healthy_image")
+def with_immediately_healthy_image_fixture(tmp_path):
+    DOCKERFILE = """FROM alpine:latest
+    RUN touch /testfile
+
+    HEALTHCHECK --interval=1s CMD test -e /testfile
+    CMD sleep infinity
+    """
+    dockerfile = tmp_path / "Dockerfile"
+    dockerfile.write_text(textwrap.dedent(DOCKERFILE))
+    with DockerImage(tmp_path) as image:
+        yield image
+
+
+def test_wait_for_healthcheck_immediate_healthy(with_immediately_healthy_image: DockerImage):
+    # Given
+    with DockerContainer(image=str(with_immediately_healthy_image)) as container:
+        # When
+        container.wait_for_healthcheck()
+
+
+@pytest.fixture(name="with_eventually_healthy_image")
+def with_eventually_healthy_image_fixture(tmp_path):
+    DOCKERFILE = """FROM alpine:latest
+    RUN touch /testfile
+
+    HEALTHCHECK --interval=1s CMD test -e /testfile
+    CMD sleep 4 && touch /testfile
+    """
+    dockerfile = tmp_path / "Dockerfile"
+    dockerfile.write_text(textwrap.dedent(DOCKERFILE))
+    with DockerImage(tmp_path) as image:
+        yield image
+
+
+def test_wait_for_healthcheck_eventually_healthy(with_eventually_healthy_image: DockerImage):
+    # Given
+    with DockerContainer(str(with_eventually_healthy_image)) as container:
+        # When
+        container.wait_for_healthcheck()