feat(core): Add RunFunctionWaitStrategy

This is usefull for converting old wait_container_is_ready
which use container.exec() to check for a condition to become True.

Author: jankatins

core/testcontainers/core/wait_strategies.py

@@ -703,6 +703,44 @@ def _get_status_compose_container(container: DockerCompose) -> str:
         raise NotImplementedError
 
 
+class RunFunctionWaitStrategy(WaitStrategy):
+    """Runs a functions and waits until it succeeds.
+
+    The function must take a single argument, the WaitStrategyTarget (= DockerContainer)
+    (use a lambda to capture other arguments) and must return a Boolean or raise an Exception.
+
+    Args:
+        func: The function to run. It must return True when the wait is over.
+    """
+
+    def __init__(
+        self,
+        func: Callable[[WaitStrategyTarget], bool],
+    ):
+        super().__init__()
+        self.func = func
+
+    def wait_until_ready(self, container: WaitStrategyTarget) -> Any:
+        start_time = time.time()
+        last_exception = None
+        while True:
+            try:
+                result = self.func(container)
+                if result:
+                    return result
+            except tuple(self._transient_exceptions) as e:
+                logger.debug(f"Check attempt failed: {e!s}")
+                last_exception = str(e)
+            if time.time() - start_time > self._startup_timeout:
+                raise TimeoutError(
+                    f"Wait time ({self._startup_timeout}s) exceeded for {self.func.__name__}"
+                    f"Exception: {last_exception}. "
+                    f"Hint: Check if the container is ready, "
+                    f"and the expected conditions are met for the function to succeed."
+                )
+            time.sleep(self._poll_interval)
+
+
 class CompositeWaitStrategy(WaitStrategy):
     """
     Wait for multiple conditions to be satisfied in sequence.
@@ -787,6 +825,7 @@ def wait_until_ready(self, container: WaitStrategyTarget) -> None:
     "HttpWaitStrategy",
     "LogMessageWaitStrategy",
     "PortWaitStrategy",
+    "RunFunctionWaitStrategy",
     "WaitStrategy",
     "WaitStrategyTarget",
 ]

core/tests/test_wait_strategies.py

@@ -15,6 +15,7 @@
     LogMessageWaitStrategy,
     PortWaitStrategy,
     WaitStrategy,
+    RunFunctionWaitStrategy,
 )
 
 
@@ -550,6 +551,60 @@ def test_wait_until_ready(self, mock_sleep, mock_time, mock_is_file, file_exists
                 strategy.wait_until_ready(mock_container)
 
 
+class TestRunFunctionWaitStrategy:
+    """Test the RunFunctionWaitStrategy class."""
+
+    def test_run_function_wait_strategy_initialization(self):
+        func = lambda x: True
+        strategy = RunFunctionWaitStrategy(func)
+        assert strategy.func == func
+
+    def test_run_function_wait_strategy_wait_until_ready(self):
+        returns = [False, False, True]
+        mock_container = object()
+
+        def func(target) -> bool:
+            assert target is mock_container
+            return returns.pop(0)
+
+        strategy = RunFunctionWaitStrategy(func).with_poll_interval(0)
+        strategy.wait_until_ready(mock_container)  # type: ignore[arg-type]
+
+    def test_run_function_wait_strategy_wait_until_ready_with_unknown_exception(self):
+        mock_container = object()
+
+        def func(target) -> bool:
+            assert target is mock_container
+            raise RuntimeError("Unknown error, abort!")
+
+        strategy = RunFunctionWaitStrategy(func).with_poll_interval(0)
+        with pytest.raises(RuntimeError, match="Unknown error, abort!"):
+            strategy.wait_until_ready(mock_container)  # type: ignore[arg-type]
+
+    @pytest.mark.parametrize("transient_exception", [ConnectionError, NotImplementedError])
+    def test_run_function_wait_strategy_wait_until_ready_with_transient_exception(self, transient_exception):
+        mock_container = object()
+        returns = [False, False, True]
+
+        def func(target) -> bool:
+            assert target is mock_container
+            if returns.pop(0):
+                return True
+            raise transient_exception("Go on")
+
+        # ConnectionError should be in the default transient exceptions, but NotImplementedError ist not
+        strategy = (
+            RunFunctionWaitStrategy(func).with_poll_interval(0.001).with_transient_exceptions(NotImplementedError)
+        )
+        strategy.wait_until_ready(mock_container)  # type: ignore[arg-type]
+
+    def test_run_function_wait_strategy_wait_until_ready_with_timeout(self):
+        mock_container = object()
+        strategy = RunFunctionWaitStrategy(lambda x: False).with_poll_interval(0).with_startup_timeout(0)
+        with pytest.raises(TimeoutError, match=r"Wait time (.*) exceeded for"):
+            strategy.wait_until_ready(mock_container)  # type: ignore[arg-type]
+
+
 class TestCompositeWaitStrategy:
     """Test the CompositeWaitStrategy class."""