From b1453e87e1f5443f0f8d04c9b30a278aa835ca9b Mon Sep 17 00:00:00 2001
From: David Ankin <daveankin@gmail.com>
Date: Sat, 3 Aug 2024 00:19:00 -0400
Subject: [PATCH] feat(core): add ability to do OR & AND for waitforlogs (#661)

---
 core/testcontainers/core/waiting_utils.py     | 16 ++++++++--
 .../testcontainers/postgres/__init__.py       | 32 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/core/testcontainers/core/waiting_utils.py b/core/testcontainers/core/waiting_utils.py
index cc3351d1..82ea1f15 100644
--- a/core/testcontainers/core/waiting_utils.py
+++ b/core/testcontainers/core/waiting_utils.py
@@ -78,7 +78,12 @@ def wait_for(condition: Callable[..., bool]) -> bool:
 
 
 def wait_for_logs(
-    container: "DockerContainer", predicate: Union[Callable, str], timeout: float = config.timeout, interval: float = 1
+    container: "DockerContainer",
+    predicate: Union[Callable, str],
+    timeout: float = config.timeout,
+    interval: float = 1,
+    predicate_streams_and: bool = False,
+    #
 ) -> float:
     """
     Wait for the container to emit logs satisfying the predicate.
@@ -90,6 +95,7 @@ def wait_for_logs(
         timeout: Number of seconds to wait for the predicate to be satisfied. Defaults to wait
             indefinitely.
         interval: Interval at which to poll the logs.
+        predicate_streams_and: should the predicate be applied to both
 
     Returns:
         duration: Number of seconds until the predicate was satisfied.
@@ -101,7 +107,13 @@ def wait_for_logs(
         duration = time.time() - start
         stdout = container.get_logs()[0].decode()
         stderr = container.get_logs()[1].decode()
-        if predicate(stdout) or predicate(stderr):
+        predicate_result = (
+            predicate(stdout) or predicate(stderr)
+            if predicate_streams_and is False
+            else predicate(stdout) and predicate(stderr)
+            #
+        )
+        if predicate_result:
             return duration
         if duration > timeout:
             raise TimeoutError(f"Container did not emit logs satisfying predicate in {timeout:.3f} " "seconds")
diff --git a/modules/postgres/testcontainers/postgres/__init__.py b/modules/postgres/testcontainers/postgres/__init__.py
index 80baef75..c9ba2a22 100644
--- a/modules/postgres/testcontainers/postgres/__init__.py
+++ b/modules/postgres/testcontainers/postgres/__init__.py
@@ -91,7 +91,37 @@ def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] =
 
     @wait_container_is_ready()
     def _connect(self) -> None:
-        wait_for_logs(self, ".*database system is ready to accept connections.*", c.max_tries, c.sleep_time)
+        # postgres itself logs these messages to the standard error stream:
+        #
+        # $ /opt/homebrew/opt/postgresql@14/bin/postgres -D /opt/homebrew/var/postgresql@14 \
+        # > | grep -o -a -m 1 -h 'database system is ready to accept connections'
+        # 2024-08-03 00:13:02.799 EDT [70226] LOG:  starting PostgreSQL 14.11 (Homebrew) ....
+        # 2024-08-03 00:13:02.804 EDT [70226] LOG:  listening on IPv4 address "127.0.0.1", port 5432
+        # ...
+        # ^C2024-08-03 00:13:04.226 EDT [70226] LOG:  received fast shutdown request
+        # ...
+        #
+        # $ /opt/homebrew/opt/postgresql@14/bin/postgres -D /opt/homebrew/var/postgresql@14 2>&1 \
+        # > | grep -o -a -m 1 -h 'database system is ready to accept connections'
+        # database system is ready to accept connections
+        #
+        # and the setup script inside docker library postgres
+        # uses pg_ctl:
+        # https://github.com/docker-library/postgres/blob/66da3846b40396249936938ee17e9684e6968a57/16/alpine3.20/docker-entrypoint.sh#L261-L282
+        # which prints logs to stdout:
+        # https://www.postgresql.org/docs/current/app-pg-ctl.html#:~:text=the%20server%27s%20standard%20output%20and%20standard%20error%20are%20sent%20to%20pg_ctl%27s%20standard%20output
+        #
+        # so we must wait for both the setup and real startup:
+        predicate_streams_and = True
+
+        wait_for_logs(
+            self,
+            ".*database system is ready to accept connections.*",
+            c.max_tries,
+            c.sleep_time,
+            predicate_streams_and=predicate_streams_and,
+            #
+        )
 
         count = 0
         while count < c.max_tries: