fix: Allow passing in a custom wait strategy string in MySQL, Cassandra, Kafka and Trino (#953)

Hey,

Just a small change in adding the ability to pass in a custom wait
strategy message for MySQL. The current implementation had '.*: ready
for connections.*: ready for connections.*' as a default regex with no
option to change it. This PR moves it into a parameter.
No new functionality has been added so I presume the tests can stay the
same.

If you find this useful, I can also go over the other implementations
and make a similar change wherever this is a problem.

Author: SubhadityaMukherjee

modules/cassandra/testcontainers/cassandra/__init__.py

@@ -39,15 +39,17 @@ class CassandraContainer(DockerContainer):
     CQL_PORT = 9042
     DEFAULT_LOCAL_DATACENTER = "datacenter1"
 
-    def __init__(self, image: str = "cassandra:latest", **kwargs) -> None:
+    def __init__(
+        self, image: str = "cassandra:latest", wait_strategy_check_string: str = "Startup complete", **kwargs
+    ) -> None:
         super().__init__(image=image, **kwargs)
         self.with_exposed_ports(self.CQL_PORT)
         self.with_env("JVM_OPTS", "-Dcassandra.skip_wait_for_gossip_to_settle=0 -Dcassandra.initial_token=0")
         self.with_env("HEAP_NEWSIZE", "128M")
         self.with_env("MAX_HEAP_SIZE", "1024M")
         self.with_env("CASSANDRA_ENDPOINT_SNITCH", "GossipingPropertyFileSnitch")
         self.with_env("CASSANDRA_DC", self.DEFAULT_LOCAL_DATACENTER)
-        self.waiting_for(LogMessageWaitStrategy("Startup complete"))
+        self.waiting_for(LogMessageWaitStrategy(wait_strategy_check_string))
 
     def get_contact_points(self) -> list[tuple[str, int]]:
         return [(self.get_container_host_ip(), int(self.get_exposed_port(self.CQL_PORT)))]

modules/kafka/testcontainers/kafka/__init__.py

@@ -55,12 +55,18 @@ class KafkaContainer(DockerContainer):
     TC_START_SCRIPT = "/tc-start.sh"
     MIN_KRAFT_TAG = "7.0.0"
 
-    def __init__(self, image: str = "confluentinc/cp-kafka:7.6.0", port: int = 9093, **kwargs) -> None:
+    def __init__(
+        self,
+        image: str = "confluentinc/cp-kafka:7.6.0",
+        port: int = 9093,
+        wait_strategy_check_string: str = r".*\[KafkaServer id=\d+\] started.*",
+        **kwargs,
+    ) -> None:
         raise_for_deprecated_parameter(kwargs, "port_to_expose", "port")
         super().__init__(image, **kwargs)
         self.port = port
         self.kraft_enabled = False
-        self.wait_for: re.Pattern[str] = re.compile(r".*\[KafkaServer id=\d+\] started.*")
+        self.wait_for: re.Pattern[str] = re.compile(wait_strategy_check_string)
         self.boot_command = ""
         self.cluster_id = "MkU3OEVBNTcwNTJENDM2Qk"
         self.listeners = f"PLAINTEXT://0.0.0.0:{self.port},BROKER://0.0.0.0:9092"

modules/mysql/testcontainers/mysql/__init__.py

@@ -71,6 +71,7 @@ def __init__(
         dbname: Optional[str] = None,
         port: int = 3306,
         seed: Optional[str] = None,
+        wait_strategy_check_string: str = r".*: ready for connections.*: ready for connections.*",
         **kwargs,
     ) -> None:
         if dialect is not None and dialect.startswith("mysql+"):
@@ -96,6 +97,7 @@ def __init__(
         if self.username == "root":
             self.root_password = self.password
         self.seed = seed
+        self.wait_strategy_check_string = wait_strategy_check_string
 
     def _configure(self) -> None:
         self.with_env("MYSQL_ROOT_PASSWORD", self.root_password)
@@ -107,7 +109,7 @@ def _configure(self) -> None:
 
     def _connect(self) -> None:
         wait_strategy = LogMessageWaitStrategy(
-            re.compile(r".*: ready for connections.*: ready for connections.*", flags=re.DOTALL | re.MULTILINE),
+            re.compile(self.wait_strategy_check_string, flags=re.DOTALL | re.MULTILINE),
         )
         wait_strategy.wait_until_ready(self)
 

modules/trino/testcontainers/trino/__init__.py

@@ -24,14 +24,15 @@ def __init__(
         user: str = "test",
         port: int = 8080,
         container_start_timeout: int = 30,
+        wait_strategy_check_string: str = ".*======== SERVER STARTED ========.*",
         **kwargs,
     ):
         super().__init__(image=image, **kwargs)
         self.user = user
         self.port = port
         self.with_exposed_ports(self.port)
         self.waiting_for(
-            LogMessageWaitStrategy(re.compile(".*======== SERVER STARTED ========.*", re.MULTILINE))
+            LogMessageWaitStrategy(re.compile(wait_strategy_check_string, re.MULTILINE))
             .with_poll_interval(c.sleep_time)
             .with_startup_timeout(container_start_timeout)
         )