feat(core, postgres): Add RunFunctionWaitStrategy and use it in pg

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."""
 

modules/postgres/testcontainers/postgres/__init__.py

@@ -15,7 +15,8 @@
 
 from testcontainers.core.generic import DbContainer
 from testcontainers.core.utils import raise_for_deprecated_parameter
-from testcontainers.core.waiting_utils import wait_container_is_ready
+from testcontainers.core.wait_strategies import RunFunctionWaitStrategy
+from testcontainers.core.waiting_utils import WaitStrategyTarget
 
 _UNSET = object()
 
@@ -64,6 +65,7 @@ def __init__(
  self.driver = f"+{driver}" if driver else ""
 
  self.with_exposed_ports(self.port)
+ self.waiting_for(RunFunctionWaitStrategy(_check_postgres_ready))
 
  def _configure(self) -> None:
  self.with_env("POSTGRES_USER", self.username)
@@ -87,7 +89,6 @@ def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] =
  port=self.port,
  )
 
- @wait_container_is_ready()
  def _connect(self) -> None:
  escaped_single_password = self.password.replace("'", "'\"'\"'")
  result = self.exec(
@@ -99,3 +100,11 @@ def _connect(self) -> None:
  )
  if result.exit_code:
  raise ConnectionError("pg_isready is not ready yet")
+
+
+def _check_postgres_ready(container: WaitStrategyTarget) -> bool:
+ if not isinstance(container, PostgresContainer):
+ raise AssertionError("This check can only wait for postgres containers to start up")
+ # This raises a ConnectionError if not ready yet
+ container._connect()
+ return True