sdks/python: mark milvus itests with require_docker_in_docker

Author: mohamedawnallah

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

@@ -34,18 +34,6 @@
 
 import pytest
 import yaml
-from pymilvus import CollectionSchema
-from pymilvus import DataType
-from pymilvus import FieldSchema
-from pymilvus import Function
-from pymilvus import FunctionType
-from pymilvus import MilvusClient
-from pymilvus import RRFRanker
-from pymilvus.milvus_client import IndexParams
-from testcontainers.core.config import MAX_TRIES as TC_MAX_TRIES
-from testcontainers.core.config import testcontainers_config
-from testcontainers.core.generic import DbContainer
-from testcontainers.milvus import MilvusContainer
 
 import apache_beam as beam
 from apache_beam.ml.rag.types import Chunk
@@ -54,7 +42,21 @@
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 
+# pylint: disable=ungrouped-imports
 try:
+  from pymilvus import (
+      CollectionSchema,
+      DataType,
+      FieldSchema,
+      Function,
+      FunctionType,
+      MilvusClient,
+      RRFRanker)
+  from pymilvus.milvus_client import IndexParams
+  from testcontainers.core.config import MAX_TRIES as TC_MAX_TRIES
+  from testcontainers.core.config import testcontainers_config
+  from testcontainers.core.generic import DbContainer
+  from testcontainers.milvus import MilvusContainer
   from apache_beam.transforms.enrichment import Enrichment
   from apache_beam.ml.rag.enrichment.milvus_search import (
       MilvusSearchEnrichmentHandler,
@@ -467,7 +469,7 @@ def create_user_yaml(service_port: int, max_vector_field_num=5):
         os.remove(path)
 
 
-@pytest.mark.uses_testcontainer
+@pytest.mark.require_docker_in_docker
 @unittest.skipUnless(
     platform.system() == "Linux",
     "Test runs only on Linux due to lack of support, as yet, for nested "
@@ -483,22 +485,16 @@ class TestMilvusSearchEnrichment(unittest.TestCase):
 
   @classmethod
   def setUpClass(cls):
-    try:
-      cls._db = MilvusEnrichmentTestHelper.start_db_container(
-          cls._version, vector_client_max_retries=1, tc_max_retries=1)
-      cls._connection_params = MilvusConnectionParameters(
-          uri=cls._db.uri,
-          user=cls._db.user,
-          password=cls._db.password,
-          db_id=cls._db.id,
-          token=cls._db.token)
-      cls._collection_load_params = MilvusCollectionLoadParameters()
-      cls._collection_name = MilvusEnrichmentTestHelper.initialize_db_with_data(
-          cls._connection_params)
-    except Exception as e:
-      pytest.skip(
-          f"Skipping all tests in {cls.__name__} due to DB startup failure: {e}"
-      )
+    cls._db = MilvusEnrichmentTestHelper.start_db_container(cls._version)
+    cls._connection_params = MilvusConnectionParameters(
+        uri=cls._db.uri,
+        user=cls._db.user,
+        password=cls._db.password,
+        db_id=cls._db.id,
+        token=cls._db.token)
+    cls._collection_load_params = MilvusCollectionLoadParameters()
+    cls._collection_name = MilvusEnrichmentTestHelper.initialize_db_with_data(
+        cls._connection_params)
 
   @classmethod
   def tearDownClass(cls):
@@ -1368,4 +1364,4 @@ def assert_chunks_equivalent(
 
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()