1515
1616from testcontainers .core .generic import DbContainer
1717from testcontainers .core .utils import raise_for_deprecated_parameter
18- from testcontainers .core .waiting_utils import wait_container_is_ready
18+ from testcontainers .core .wait_strategies import RunFunctionWaitStrategy
19+ from testcontainers .core .waiting_utils import WaitStrategyTarget
1920
2021_UNSET = object ()
2122
@@ -64,6 +65,7 @@ def __init__(
6465 self .driver = f"+{ driver } " if driver else ""
6566
6667 self .with_exposed_ports (self .port )
68+ self .waiting_for (RunFunctionWaitStrategy (_check_postgres_ready ))
6769
6870 def _configure (self ) -> None :
6971 self .with_env ("POSTGRES_USER" , self .username )
@@ -87,15 +89,18 @@ def get_connection_url(self, host: Optional[str] = None, driver: Optional[str] =
8789 port = self .port ,
8890 )
8991
90- @wait_container_is_ready ()
91- def _connect (self ) -> None :
92- escaped_single_password = self .password .replace ("'" , "'\" '\" '" )
93- result = self .exec (
94- [
95- "sh" ,
96- "-c" ,
97- f"PGPASSWORD='{ escaped_single_password } ' psql --username { self .username } --dbname { self .dbname } --host 127.0.0.1 -c 'select version();'" ,
98- ]
99- )
100- if result .exit_code :
101- raise ConnectionError ("pg_isready is not ready yet" )
92+
93+ def _check_postgres_ready (container : WaitStrategyTarget ) -> bool :
94+ if not isinstance (container , PostgresContainer ):
95+ raise AssertionError ("This check can only wait for postgres containers to start up" )
96+ escaped_single_password = container .password .replace ("'" , "'\" '\" '" )
97+ result = container .exec (
98+ [
99+ "sh" ,
100+ "-c" ,
101+ f"PGPASSWORD='{ escaped_single_password } ' psql --username { container .username } --dbname { container .dbname } --host 127.0.0.1 -c 'select version();'" ,
102+ ]
103+ )
104+ if result .exit_code :
105+ raise ConnectionError ("pg is not ready yet" )
106+ return True
0 commit comments