Add resource group tests (#22101)

Signed-off-by: yanliang567 <yanliang.qiao@zilliz.com>

tests/python_client/base/client_base.py

@@ -25,6 +25,7 @@ class Base:
     collection_schema_wrap = None
     field_schema_wrap = None
     collection_object_list = []
+    resource_group_list = []
 
     def setup_class(self):
         log.info("[setup_class] Start setup class...")
@@ -64,6 +65,16 @@ class Base:
                 if collection_object.collection is not None and collection_object.name in collection_list:
                     collection_object.drop(check_task=ct.CheckTasks.check_nothing)
 
+            """ Clean up the rgs before disconnect """
+            for rg_name in self.resource_group_list:
+                rg = self.utility_wrap.describe_resource_group(name=rg_name)[0]
+                if rg is not None:
+                    if rg.num_available_node > 0:
+                        self.utility_wrap.transfer_node(source=rg_name,
+                                                        target=ct.default_resource_group_name,
+                                                        num_node=rg.num_available_node)
+                    self.utility_wrap.drop_resource_group(rg_name, check_task=ct.CheckTasks.check_nothing)
+
         except Exception as e:
             log.debug(str(e))
 
@@ -262,3 +273,14 @@ class TestcaseBase(Base):
             collection_w.insert(df)
             assert collection_w.num_entities == nb_of_segment * (i + 1)
         return collection_w
+
+    def init_resource_group(self, name, using="default", timeout=None, check_task=None, check_items=None, **kwargs):
+        if not self.connection_wrap.has_connection(alias=DefaultConfig.DEFAULT_USING)[0]:
+            self._connect()
+        utility_w = ApiUtilityWrapper()
+        res, check_result = utility_w.create_resource_group(name=name, using=using, timeout=timeout,
+                                                            check_task=check_task,
+                                                            check_items=check_items, **kwargs)
+        if res is None and check_result:
+            self.resource_group_list.append(name)
+        return res, check_result

tests/python_client/base/utility_wrapper.py

@@ -449,4 +449,41 @@ class ApiUtilityWrapper:
         func_name = sys._getframe().f_code.co_name
         res, check = api_request([self.role.list_grants], **kwargs)
         check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
-        return res, check_result
+        return res, check_result
+
+    def create_resource_group(self, name, using="default", timeout=None, check_task=None, check_items=None, **kwargs):
+        func_name = sys._getframe().f_code.co_name
+        res, check = api_request([self.ut.create_resource_group, name, using, timeout], **kwargs)
+        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
+        return res, check_result
+
+    def drop_resource_group(self, name, using="default", timeout=None, check_task=None, check_items=None, **kwargs):
+        func_name = sys._getframe().f_code.co_name
+        res, check = api_request([self.ut.drop_resource_group, name, using, timeout], **kwargs)
+        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
+        return res, check_result
+
+    def list_resource_groups(self, using="default", timeout=None, check_task=None, check_items=None, **kwargs):
+        func_name = sys._getframe().f_code.co_name
+        res, check = api_request([self.ut.list_resource_groups, using, timeout], **kwargs)
+        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
+        return res, check_result
+
+    def describe_resource_group(self, name, using="default", timeout=None, check_task=None, check_items=None, **kwargs):
+        func_name = sys._getframe().f_code.co_name
+        res, check = api_request([self.ut.describe_resource_group, name, using, timeout], **kwargs)
+        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
+        return res, check_result
+
+    def transfer_node(self, source, target, num_node, using="default", timeout=None, check_task=None, check_items=None, **kwargs):
+        func_name = sys._getframe().f_code.co_name
+        res, check = api_request([self.ut.transfer_node, source, target, num_node, using, timeout], **kwargs)
+        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
+        return res, check_result
+
+    def transfer_replica(self, source, target, collection_name, num_replica, using="default", timeout=None, check_task=None, check_items=None, **kwargs):
+        func_name = sys._getframe().f_code.co_name
+        res, check = api_request([self.ut.transfer_replica, source, target, collection_name,num_replica, using, timeout], **kwargs)
+        check_result = ResponseChecker(res, func_name, check_task, check_items, check, **kwargs).run()
+        return res, check_result
+

tests/python_client/check/func_check.py

@@ -6,7 +6,7 @@ from common import common_type as ct
 from common import common_func as cf
 from common.common_type import CheckTasks, Connect_Object_Name
 # from common.code_mapping import ErrorCode, ErrorMessage
-from pymilvus import Collection, Partition
+from pymilvus import Collection, Partition, ResourceGroupInfo
 from utils.api_request import Error
 import check.param_check as pc
 
@@ -83,6 +83,9 @@ class ResponseChecker:
         elif self.check_task == CheckTasks.check_permission_deny:
             # Collection interface response check
             result = self.check_permission_deny(self.response, self.succ)
+        elif self.check_task == CheckTasks.check_rg_property:
+            # describe resource group interface response check
+            result = self.check_rg_property(self.response, self.func_name, self.check_items)
 
         # Add check_items here if something new need verify
 
@@ -196,6 +199,29 @@ class ResponseChecker:
             assert partition.num_entities == check_items["num_entities"]
         return True
 
+    @staticmethod
+    def check_rg_property(rg, func_name, check_items):
+        exp_func_name = "describe_resource_group"
+        if func_name != exp_func_name:
+            log.warning("The function name is {} rather than {}".format(func_name, exp_func_name))
+        if not isinstance(rg, ResourceGroupInfo):
+            raise Exception("The result to check isn't ResourceGroupInfo type object")
+        if len(check_items) == 0:
+            raise Exception("No expect values found in the check task")
+        if check_items.get("name", None):
+            assert rg.name == check_items["name"]
+        if check_items.get("capacity", None):
+            assert rg.capacity == check_items["capacity"]
+        if check_items.get("num_available_node", None):
+            assert rg.num_available_node == check_items["num_available_node"]
+        if check_items.get("num_loaded_replica", None):
+            assert rg.num_loaded_replica == check_items["num_loaded_replica"]
+        if check_items.get("num_outgoing_node", None):
+            assert rg.num_outgoing_node == check_items["num_outgoing_node"]
+        if check_items.get("num_incoming_node", None):
+            assert rg.num_incoming_node == check_items["num_incoming_node"]
+        return True
+
     @staticmethod
     def check_search_results(search_res, func_name, check_items):
         """

tests/python_client/common/common_func.py

@@ -323,11 +323,11 @@ def gen_default_binary_dataframe_data(nb=ct.default_nb, dim=ct.default_dim, star
     return df, binary_raw_values
 
 
-def gen_default_list_data(nb=ct.default_nb, dim=ct.default_dim):
-    int_values = [i for i in range(nb)]
-    float_values = [np.float32(i) for i in range(nb)]
-    string_values = [str(i) for i in range(nb)]
-    float_vec_values = gen_vectors(nb, dim)
+def gen_default_list_data(nb=ct.default_nb, dim=ct.default_dim, start=0):
+    int_values = [i for i in range(start, start + nb)]
+    float_values = [np.float32(i) for i in range(start, start + nb)]
+    string_values = [str(i) for i in range(start, start + nb)]
+    float_vec_values = gen_vectors(start + nb, dim)
     data = [int_values, float_values, string_values, float_vec_values]
     return data
 

tests/python_client/common/common_type.py

@@ -37,6 +37,7 @@ default_float_vec_field_name = "float_vector"
 another_float_vec_field_name = "float_vector1"
 default_binary_vec_field_name = "binary_vector"
 default_partition_name = "_default"
+default_resource_group_name = '__default_resource_group'
 default_tag = "1970_01_01"
 row_count = "row_count"
 default_length = 65535
@@ -200,6 +201,7 @@ class CheckTasks:
     check_role_property = "check_role_property"
     check_permission_deny = "check_permission_deny"
     check_value_equal = "check_value_equal"
+    check_rg_property = "check_resource_group_property"
 
 
 class BulkLoadStates: