sdks/python: handle transient testcontainer startup/teardown errors

Author: mohamedawnallah

sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py

@@ -68,6 +68,9 @@ class TestContainerStartupError(Exception):
   """Raised when any test container fails to start."""
   pass
 
+class TestContainerTeardownError(Exception):
+  """Raised when any test container fails to teardown."""
+  pass
 
 def validate_enrichment_with_bigtable():
   expected = '''[START enrichment_with_bigtable]
@@ -186,7 +189,7 @@ def test_enrichment_with_external_pg(self, mock_stdout):
         output = mock_stdout.getvalue().splitlines()
         expected = validate_enrichment_with_external_pg()
         self.assertEqual(output, expected)
-    except TestContainerStartupError as e:
+    except (TestContainerStartupError, TestContainerTeardownError) as e:
       raise unittest.SkipTest(str(e))
     except Exception as e:
       self.fail(f"Test failed with unexpected error: {e}")
@@ -199,7 +202,7 @@ def test_enrichment_with_external_mysql(self, mock_stdout):
         output = mock_stdout.getvalue().splitlines()
         expected = validate_enrichment_with_external_mysql()
         self.assertEqual(output, expected)
-    except TestContainerStartupError as e:
+    except (TestContainerStartupError, TestContainerTeardownError) as e:
       raise unittest.SkipTest(str(e))
     except Exception as e:
       self.fail(f"Test failed with unexpected error: {e}")
@@ -212,7 +215,7 @@ def test_enrichment_with_external_sqlserver(self, mock_stdout):
         output = mock_stdout.getvalue().splitlines()
         expected = validate_enrichment_with_external_sqlserver()
         self.assertEqual(output, expected)
-    except TestContainerStartupError as e:
+    except (TestContainerStartupError, TestContainerTeardownError) as e:
       raise unittest.SkipTest(str(e))
     except Exception as e:
       self.fail(f"Test failed with unexpected error: {e}")
@@ -227,7 +230,7 @@ def test_enrichment_with_milvus(self, mock_stdout):
         output = parse_chunk_strings(output)
         expected = parse_chunk_strings(expected)
         assert_chunks_equivalent(output, expected)
-    except TestContainerStartupError as e:
+    except (TestContainerStartupError, TestContainerTeardownError) as e:
       raise unittest.SkipTest(str(e))
     except Exception as e:
       self.fail(f"Test failed with unexpected error: {e}")
@@ -373,19 +376,17 @@ def post_sql_enrichment_test(res: CloudSQLEnrichmentTestDataConstruct):
   def pre_milvus_enrichment() -> MilvusDBContainerInfo:
     try:
       db = MilvusEnrichmentTestHelper.start_db_container()
-    except Exception as e:
-      raise TestContainerStartupError(
-          f"Milvus container failed to start: {str(e)}")
-
-    connection_params = MilvusConnectionParameters(
+      connection_params = MilvusConnectionParameters(
         uri=db.uri,
         user=db.user,
         password=db.password,
         db_id=db.id,
         token=db.token)
-
-    collection_name = MilvusEnrichmentTestHelper.initialize_db_with_data(
+      collection_name = MilvusEnrichmentTestHelper.initialize_db_with_data(
         connection_params)
+    except Exception as e:
+      raise TestContainerStartupError(
+          f"Milvus container failed to start: {str(e)}")
 
     # Setup environment variables for db and collection configuration. This will
     # be used downstream by the milvus enrichment handler.
@@ -400,7 +401,12 @@ def pre_milvus_enrichment() -> MilvusDBContainerInfo:
 
   @staticmethod
   def post_milvus_enrichment(db: MilvusDBContainerInfo):
-    MilvusEnrichmentTestHelper.stop_db_container(db)
+    try:
+      MilvusEnrichmentTestHelper.stop_db_container(db)
+    except Exception:
+      raise TestContainerTeardownError(
+        f"Milvus container failed to tear down: {str(e)}")
+
     os.environ.pop('MILVUS_VECTOR_DB_URI', None)
     os.environ.pop('MILVUS_VECTOR_DB_USER', None)
     os.environ.pop('MILVUS_VECTOR_DB_PASSWORD', None)

sdks/python/apache_beam/ml/rag/enrichment/milvus_search_it_test.py

@@ -312,7 +312,7 @@ def start_db_container(
     with user_yaml_creator(service_container_port, max_vec_fields) as cfg:
       info = None
       original_tc_max_tries = testcontainers_config.max_tries
-      if not testcontainers_config.max_tries:
+      if testcontainers_config.max_tries is not None:
         testcontainers_config.max_tries = tc_max_retries
       for i in range(vector_client_max_retries):
         try:
@@ -326,7 +326,6 @@ def start_db_container(
           host = vector_db_container.get_container_host_ip()
           port = vector_db_container.get_exposed_port(service_container_port)
           info = MilvusDBContainerInfo(vector_db_container, host, port)
-          testcontainers_config.max_tries = original_tc_max_tries
           _LOGGER.info(
               "milvus db container started successfully on %s.", info.uri)
           break
@@ -351,20 +350,18 @@ def start_db_container(
                 stdout_logs,
                 stderr_logs)
             raise e
+        finally:
+          testcontainers_config.max_tries = original_tc_max_tries
       return info
 
   @staticmethod
   def stop_db_container(db_info: MilvusDBContainerInfo):
     if db_info is None:
       _LOGGER.warning("Milvus db info is None. Skipping stop operation.")
       return
-    try:
-      _LOGGER.debug("Stopping milvus db container.")
-      db_info.container.stop()
-      _LOGGER.info("milvus db container stopped successfully.")
-    except Exception as e:
-      _LOGGER.warning(
-          "Error encountered while stopping milvus db container: %s", e)
+    _LOGGER.debug("Stopping milvus db container.")
+    db_info.container.stop()
+    _LOGGER.info("milvus db container stopped successfully.")
 
   @staticmethod
   def initialize_db_with_data(connc_params: MilvusConnectionParameters):