replace more stuff

Author: alexanderankin

core/testcontainers/core/generic.py

@@ -16,7 +16,6 @@
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.exceptions import ContainerStartException
 from testcontainers.core.utils import raise_for_deprecated_parameter
-from testcontainers.core.waiting_utils import wait_container_is_ready
 
 ADDITIONAL_TRANSIENT_ERRORS = []
 try:
@@ -34,8 +33,11 @@ class DbContainer(DockerContainer):
     Generic database container.
     """
 
-    @wait_container_is_ready(*ADDITIONAL_TRANSIENT_ERRORS)
     def _connect(self) -> None:
+        from testcontainers.core.wait_strategies import ContainerStatusWaitStrategy as C
+
+        C().with_transient_exceptions(*ADDITIONAL_TRANSIENT_ERRORS).wait_until_ready(self)
+
         import sqlalchemy
 
         engine = sqlalchemy.create_engine(self.get_connection_url())

core/testcontainers/core/wait_strategies.py

@@ -35,6 +35,8 @@
 from urllib.error import HTTPError, URLError
 from urllib.request import Request, urlopen
 
+from typing_extensions import Self
+
 from testcontainers.compose import DockerCompose
 from testcontainers.core.utils import setup_logger
 # Import base classes from waiting_utils to make them available for tests
@@ -725,6 +727,24 @@ def __init__(self, *strategies: WaitStrategy) -> None:
         super().__init__()
         self._strategies = list(strategies)
 
+    def with_poll_interval(self, interval: Union[float, timedelta]) -> Self:
+        super().with_poll_interval(interval)
+        for _strategy in self._strategies:
+            _strategy.with_poll_interval(interval)
+        return self
+
+    def with_startup_timeout(self, timeout: Union[int, timedelta]) -> Self:
+        super().with_startup_timeout(timeout)
+        for _strategy in self._strategies:
+            _strategy.with_startup_timeout(timeout)
+        return self
+
+    def with_transient_exceptions(self, *transient_exceptions: type[Exception]) -> Self:
+        super().with_transient_exceptions(*transient_exceptions)
+        for _strategy in self._strategies:
+            _strategy.with_transient_exceptions(*transient_exceptions)
+        return self
+
     def wait_until_ready(self, container: WaitStrategyTarget) -> None:
         """
         Wait until all contained strategies are ready.

core/testcontainers/core/waiting_utils.py

@@ -75,6 +75,7 @@ class WaitStrategy(ABC):
     def __init__(self) -> None:
         self._startup_timeout: float = testcontainers_config.timeout
         self._poll_interval: float = testcontainers_config.sleep_time
+        self._transient_exceptions: list[type[Exception]] = [*TRANSIENT_EXCEPTIONS]
 
     def with_startup_timeout(self, timeout: Union[int, timedelta]) -> "WaitStrategy":
         """Set the maximum time to wait for the container to be ready."""
@@ -92,12 +93,21 @@ def with_poll_interval(self, interval: Union[float, timedelta]) -> "WaitStrategy
             self._poll_interval = interval
         return self
 
+    def with_transient_exceptions(self, *transient_exceptions: type[Exception]) -> "WaitStrategy":
+        self._transient_exceptions.extend(transient_exceptions)
+        return self
+
     @abstractmethod
     def wait_until_ready(self, container: WaitStrategyTarget) -> None:
         """Wait until the container is ready."""
         pass
 
-    def _poll(self, check: Callable[[], bool]) -> bool:
+    def _poll(self, check: Callable[[], bool], transient_exceptions: list[type[Exception]] = None) -> bool:
+        if not transient_exceptions:
+            all_te_types = self._transient_exceptions
+        else:
+            all_te_types = [*self._transient_exceptions, *(transient_exceptions or [])]
+
         start = time.time()
         while True:
             start_attempt = time.time()
@@ -112,8 +122,13 @@ def _poll(self, check: Callable[[], bool]) -> bool:
                     return result
             except StopIteration:
                 return False
-            except:  # noqa: E722, RUF100
-                pass
+            except Exception as e:  # noqa: E722, RUF100
+                is_transient = False
+                for et in all_te_types:
+                    if isinstance(e, et):
+                        is_transient = True
+                if not is_transient:
+                    raise RuntimeError(f"exception while checking for strategy {self}") from e
 
             seconds_left_until_next = self._poll_interval - (time.time() - start_attempt)
             time.sleep(max(0.0, seconds_left_until_next))

core/tests/test_new_docker_api.py

@@ -4,12 +4,11 @@
 
 
 def test_docker_custom_image():
-    container = DockerContainer("mysql:5.7.17")
-    container.with_exposed_ports(3306)
-    container.with_env("MYSQL_ROOT_PASSWORD", "root")
+    container = DockerContainer("nginx:alpine-slim")
+    container.with_exposed_ports(80)
 
     with container:
-        port = container.get_exposed_port(3306)
+        port = container.get_exposed_port(80)
         assert int(port) > 0
 
 

core/tests/test_waiting_utils.py

@@ -1,6 +1,7 @@
 import pytest
 
 from testcontainers.core.container import DockerContainer
+from testcontainers.core.wait_strategies import ContainerStatusWaitStrategy
 from testcontainers.core.waiting_utils import wait_for_logs, wait_for, wait_container_is_ready
 
 
@@ -28,8 +29,15 @@ def simple_check() -> bool:
 def test_wait_container_is_ready_decorator_with_container() -> None:
     """Test wait_container_is_ready decorator with a real container."""
 
-    @wait_container_is_ready()
     def check_container_logs(container: DockerContainer) -> bool:
+        # wait until it becomes running.
+        # if it is too late, it is actually fine in this case.
+        # we are happy with an exited (even crashed) container that has logs.
+        try:
+            ContainerStatusWaitStrategy().wait_until_ready(container)
+        except TimeoutError:
+            pass
+
         stdout, stderr = container.get_logs()
         return b"Hello from Docker!" in stdout or b"Hello from Docker!" in stderr