Add more retries for starting Redis test container

Author: Amar3tto

sdks/python/apache_beam/io/requestresponse_it_test.py

@@ -206,7 +206,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 @pytest.mark.uses_testcontainer
 class TestRedisCache(unittest.TestCase):
   def setUp(self) -> None:
-    self.retries = 3
+    self.retries = 10
     self._start_container()
 
   def test_rrio_cache_all_miss(self):
@@ -288,7 +288,8 @@ def test_rrio_no_coder_exception(self):
       res.wait_until_finish()
 
   def tearDown(self) -> None:
-    self.container.stop()
+    if hasattr(self, 'container') and self.container:
+      self.container.stop()
 
   def _start_container(self):
     for i in range(self.retries):
@@ -298,6 +299,15 @@ def _start_container(self):
         self.host = self.container.get_container_host_ip()
         self.port = self.container.get_exposed_port(6379)
         self.client = self.container.get_client()
+
+        # Ping Redis to ensure it's ready
+        for _ in range(10):
+          try:
+            if self.client.ping():
+              break
+          except Exception:
+            time.sleep(0.5)
+
         break
       except Exception as e:
         if i == self.retries - 1:

sdks/python/apache_beam/transforms/enrichment_handlers/bigquery_it_test.py

@@ -140,7 +140,7 @@ def setUpClass(cls):
 
   def setUp(self) -> None:
     self.condition_template = "id = {}"
-    self.retries = 3
+    self.retries = 10
     self._start_container()
 
   def _start_container(self):
@@ -151,6 +151,15 @@ def _start_container(self):
         self.host = self.container.get_container_host_ip()
         self.port = self.container.get_exposed_port(6379)
         self.client = self.container.get_client()
+
+        # Ping Redis to ensure it's ready
+        for _ in range(10):
+          try:
+            if self.client.ping():
+              break
+          except Exception:
+            time.sleep(0.5)
+
         break
       except Exception as e:
         if i == self.retries - 1:
@@ -160,7 +169,8 @@ def _start_container(self):
           raise e
 
   def tearDown(self) -> None:
-    self.container.stop()
+    if hasattr(self, 'container') and self.container:
+      self.container.stop()
     self.client = None
 
   def test_bigquery_enrichment(self):

sdks/python/apache_beam/transforms/enrichment_handlers/bigtable_it_test.py

@@ -167,7 +167,7 @@ def setUp(self):
     instance = client.instance(self.instance_id)
     self.table = instance.table(self.table_id)
     create_rows(self.table)
-    self.retries = 3
+    self.retries = 10
     self._start_container()
 
   def _start_container(self):
@@ -178,14 +178,24 @@ def _start_container(self):
         self.host = self.container.get_container_host_ip()
         self.port = self.container.get_exposed_port(6379)
         self.client = self.container.get_client()
+
+        # Ping Redis to ensure it's ready
+        for _ in range(10):
+          try:
+            if self.client.ping():
+              break
+          except Exception:
+            time.sleep(0.5)
+
         break
       except Exception as e:
         if i == self.retries - 1:
           _LOGGER.error('Unable to start redis container for RRIO tests.')
           raise e
 
   def tearDown(self) -> None:
-    self.container.stop()
+    if hasattr(self, 'container') and self.container:
+      self.container.stop()
     self.table = None
 
   def test_enrichment_with_bigtable(self):

sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_it_test.py

@@ -71,25 +71,35 @@ def setUp(self) -> None:
     self.entity_type_name = "entity_id"
     self.api_endpoint = "us-central1-aiplatform.googleapis.com"
     self.feature_ids = ['title', 'genres']
-    self.retries = 3
+    self.retries = 10
     self._start_container()
 
   def _start_container(self):
-    for i in range(3):
+    for i in range(self.retries):
       try:
         self.container = RedisContainer(image='redis:7.2.4')
         self.container.start()
         self.host = self.container.get_container_host_ip()
         self.port = self.container.get_exposed_port(6379)
         self.client = self.container.get_client()
+
+        # Ping Redis to ensure it's ready
+        for _ in range(10):
+          try:
+            if self.client.ping():
+              break
+          except Exception:
+            time.sleep(0.5)
+
         break
       except Exception as e:
         if i == self.retries - 1:
           _LOGGER.error('Unable to start redis container for RRIO tests.')
           raise e
 
   def tearDown(self) -> None:
-    self.container.stop()
+    if hasattr(self, 'container') and self.container:
+      self.container.stop()
     self.client = None
 
   def test_vertex_ai_feature_store_bigtable_serving_enrichment(self):