changed things and linted

Signed-off-by: Francisco Javier Arceo <[email protected]>

sdk/python/tests/integration/feature_repos/universal/online_store/milvus.py

@@ -1,8 +1,7 @@
-import time
 from typing import Any, Dict
 
-from pymilvus import connections
 from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
 
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
@@ -12,7 +11,7 @@
 class MilvusOnlineStoreCreator(OnlineStoreCreator):
     def __init__(self, project_name: str, **kwargs):
         super().__init__(project_name)
-        self.container = DockerContainer("milvusdb/milvus:v2.2.9").with_exposed_ports(
+        self.container = DockerContainer("milvusdb/milvus:v2.4.4").with_exposed_ports(
             19530
         )
 
@@ -22,23 +21,14 @@ def create_online_store(self) -> Dict[str, Any]:
         host = "localhost"
         port = self.container.get_exposed_port(19530)
 
-        max_attempts = 12
-        for attempt in range(1, max_attempts + 1):
-            try:
-                print(
-                    f"Attempting to connect to Milvus at {host}:{port}, attempt {attempt}"
-                )
-                connections.connect(alias="default", host=host, port=port)
-                if connections.has_connection(alias="default"):
-                    print("Successfully connected to Milvus")
-                    break
-            except Exception as e:
-                print(f"Connection attempt failed: {e}")
-                time.sleep(5)
-        else:
-            raise RuntimeError(
-                "Cannot connect to Milvus server after multiple attempts"
-            )
+        log_string_to_wait_for = "database system is ready to accept connections"
+        wait_for_logs(
+            container=self.container, predicate=log_string_to_wait_for, timeout=10
+        )
+        init_log_string_to_wait_for = "Milvus DB init process complete"
+        wait_for_logs(
+            container=self.container, predicate=init_log_string_to_wait_for, timeout=10
+        )
 
         return {
             "type": "milvus",