Use new _wait_until_ready hook in google module

modules/google/testcontainers/google/__init__.py

@@ -1,2 +1,4 @@
-from .datastore import DatastoreContainer  # noqa: F401
-from .pubsub import PubSubContainer  # noqa: F401
+from .datastore import DatastoreContainer
+from .pubsub import PubSubContainer
+
+__all__ = ["DatastoreContainer", "PubSubContainer"]

modules/google/testcontainers/google/datastore.py

@@ -13,6 +13,8 @@
 import os
 from unittest.mock import patch
 
+from typing_extensions import override
+
 from google.cloud import datastore
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
@@ -32,8 +34,7 @@ class DatastoreContainer(DockerContainer):
 
             >>> from testcontainers.google import DatastoreContainer
 
-            >>> config = DatastoreContainer()
-            >>> with config as datastore:
+            >>> with DatastoreContainer("google/cloud-sdk:471.0.0-emulators") as datastore:
             ...    datastore_client = datastore.get_datastore_client()
     """
 
@@ -56,7 +57,6 @@ def get_datastore_emulator_host(self) -> str:
         return f"{self.get_container_host_ip()}:{self.get_exposed_port(self.port)}"
 
     def get_datastore_client(self, **kwargs) -> datastore.Client:
-        wait_for_logs(self, "Dev App Server is now running.", timeout=30.0)
         env_vars = {
             "DATASTORE_DATASET": self.project,
             "DATASTORE_EMULATOR_HOST": self.get_datastore_emulator_host(),
@@ -66,3 +66,7 @@ def get_datastore_client(self, **kwargs) -> datastore.Client:
         }
         with patch.dict(os.environ, env_vars):
             return datastore.Client(**kwargs)
+
+    @override
+    def _wait_until_ready(self) -> None:
+        wait_for_logs(self, "Dev App Server is now running.", timeout=30.0)

modules/google/testcontainers/google/pubsub.py

@@ -32,8 +32,7 @@ class PubSubContainer(DockerContainer):
 
             >>> from testcontainers.google import PubSubContainer
 
-            >>> config = PubSubContainer()
-            >>> with config as pubsub:
+            >>> with PubSubContainer("google/cloud-sdk:471.0.0-emulators") as pubsub:
             ...    publisher = pubsub.get_publisher_client()
             ...    topic_path = publisher.topic_path(pubsub.project, "my-topic")
             ...    topic = publisher.create_topic(name=topic_path)

modules/google/tests/test_google.py

@@ -6,8 +6,7 @@
 
 
 def test_pubsub_container():
-    pubsub: PubSubContainer
-    with PubSubContainer() as pubsub:
+    with PubSubContainer("google/cloud-sdk:471.0.0-emulators") as pubsub:
         wait_for_logs(pubsub, r"Server started, listening on \d+", timeout=60)
         # Create a new topic
         publisher = pubsub.get_publisher_client()
@@ -32,7 +31,7 @@ def test_pubsub_container():
 
 def test_datastore_container_creation():
     # Initialize the Datastore emulator container
-    with DatastoreContainer() as datastore:
+    with DatastoreContainer("google/cloud-sdk:471.0.0-emulators") as datastore:
         # Obtain a datastore client configured to connect to the emulator
         client = datastore.get_datastore_client()
 
@@ -54,7 +53,7 @@ def test_datastore_container_creation():
 
 def test_datastore_container_isolation():
     # Initialize the Datastore emulator container
-    with DatastoreContainer() as datastore:
+    with DatastoreContainer("google/cloud-sdk:471.0.0-emulators") as datastore:
         # Obtain a datastore client configured to connect to the emulator
         client = datastore.get_datastore_client()
 
@@ -67,7 +66,7 @@ def test_datastore_container_isolation():
         client.put(entity)
 
         # Create a second container and try to fetch the entity to makesure its a different container
-        with DatastoreContainer() as datastore2:
+        with DatastoreContainer("google/cloud-sdk:471.0.0-emulators") as datastore2:
             assert (
                 datastore.get_datastore_emulator_host() != datastore2.get_datastore_emulator_host()
             ), "Datastore containers use the same port."