Update rate limit test file (#18942)

Signed-off-by: Binbin Lv <binbin.lv@zilliz.com>

Signed-off-by: Binbin Lv <binbin.lv@zilliz.com>

tests/python_client/common/common_func.py

@@ -10,6 +10,7 @@ from pymilvus import DataType
 from base.schema_wrapper import ApiCollectionSchemaWrapper, ApiFieldSchemaWrapper
 from common import common_type as ct
 from utils.util_log import test_log as log
+from customize.milvus_operator import MilvusOperator
 
 """" Methods of processing data """
 
@@ -778,3 +779,60 @@ def gen_grant_list(collection_name):
                   {"object": "User", "object_name": "*", "privilege": "UpdateUser"},
                   {"object": "User", "object_name": "*", "privilege": "SelectUser"}]
     return grant_list
+
+def install_milvus_operator_specific_config(namespace, milvus_mode, release_name, image,
+                                            rate_limit_enable, collection_rate_limit):
+    """
+    namespace : str
+    milvus_mode : str -> standalone or cluster
+    release_name : str
+    image: str -> image tag including repository
+    rate_limit_enable: str -> true or false, switch for rate limit
+    collection_rate_limit: int -> collection rate limit numbers
+    input_content ：the content that need to insert to the file
+    return: milvus host name
+    """
+
+    if not isinstance(namespace, str):
+        log.error("[namespace] is not a string.")
+
+    if not isinstance(milvus_mode, str):
+        log.error("[milvus_mode] is not a string.")
+
+    if not isinstance(release_name, str):
+        log.error("[release_name] is not a string.")
+
+    if not isinstance(image, str):
+        log.error("[image] is not a string.")
+
+    if not isinstance(rate_limit_enable, str):
+        log.error("[rate_limit_enable] is not a string.")
+
+    if not isinstance(collection_rate_limit, int):
+        log.error("[collection_rate_limit] is not an integer.")
+
+    if milvus_mode not in ["standalone", "cluster"]:
+        log.error("[milvus_mode] is not 'standalone' or 'cluster'")
+    
+    if rate_limit_enable not in ["true", "false"]:
+        log.error("[rate_limit_enable] is not 'true' or 'false'")
+
+    data_config = {
+        'metadata.namespace': namespace,
+        'spec.mode': milvus_mode,
+        'metadata.name': release_name,
+        'spec.components.image': image,
+        'spec.components.proxy.serviceType': 'LoadBalancer',
+        'spec.components.dataNode.replicas': 2,
+        'spec.config.common.retentionDuration': 60,
+        'spec.config.quotaAndLimits.enable': rate_limit_enable,
+        'spec.config.quotaAndLimits.ddl.collectionRate': collection_rate_limit,
+    }
+    mil = MilvusOperator()
+    mil.install(data_config)
+    if mil.wait_for_healthy(release_name, NAMESPACE, timeout=TIMEOUT):
+        host = mic.endpoint(release_name, NAMESPACE).split(':')[0]
+    else:
+        raise MilvusException(message=f'Milvus healthy timeout 1800s')
+    
+    return host

tests/python_client/common/common_type.py

@@ -57,6 +57,8 @@ compact_retention_duration = 40  # compaction travel time retention range 20s
 max_compaction_interval = 60  # the max time interval (s) from the last compaction
 max_field_num = 256  # Maximum number of fields in a collection
 default_replica_num = 1
+IMAGE_REPOSITORY_MILVUS = "harbor.milvus.io/dockerhub/milvusdb/milvus"
+NAMESPACE_CHAOS_TESTING = "chaos-testing"
 
 Not_Exist = "Not_Exist"
 Connect_Object_Name = True

tests/python_client/rate_limit/test_rate_limit.py

@@ -0,0 +1,107 @@
+import multiprocessing
+import numbers
+import random
+
+import pytest
+import pandas as pd
+from time import sleep
+
+from base.client_base import TestcaseBase
+from utils.util_log import test_log as log
+from common import common_func as cf
+from common import common_type as ct
+from common.common_type import CaseLabel, CheckTasks
+from utils.util_pymilvus import *
+from utils.util_k8s import wait_pods_ready, read_pod_log
+from utils.util_pymilvus import get_latest_tag
+from common.constants import *
+from customize.milvus_operator import MilvusOperator
+
+
+prefix = "rate_limit_collection"
+exp_name = "name"
+exp_schema = "schema"
+TIMEOUT = 1800
+milvus_port = "19530"
+default_schema = cf.gen_default_collection_schema()
+IMAGE_REPOSITORY = ct.IMAGE_REPOSITORY_MILVUS
+NAMESPACE = ct.NAMESPACE_CHAOS_TESTING
+rate_limit_period = 61
+
+
+class TestRateLimit(TestcaseBase):
+    """ Test case of search interface """
+
+    """
+    ******************************************************************
+    #  The followings are valid cases
+    ******************************************************************
+    """
+
+    @pytest.mark.tags(CaseLabel.L3)
+    @pytest.mark.parametrize("MILVUS_MODE", ["standalone", "cluster"])
+    @pytest.mark.parametrize("rate_limit_enable", ["true", "false"])
+    @pytest.mark.parametrize("collection_rate_limit", ["10", "500"])
+    def test_rate_limit_create_drop_collection(self, MILVUS_MODE, rate_limit_enable, collection_rate_limit):
+        """
+        target: test rate limit for ddl (create collection)
+        method: 1. install milvus with different rate limit parameters
+                2. create/drop collectionRateLimit+1 collections
+        expected: 1. raise exception with rate limit enabled in one rate limit period
+                  2. create/drop collections successfully with rate limit disabled
+                  3. create/drop collections successfully in next rate limit period
+        """
+        # 1. install milvus with operator
+        release_name = "rate_limit" + MILVUS_MODE + rate_limit_enable + collection_rate_limit
+        image_tag = get_latest_tag()
+        image = f'{IMAGE_REPOSITORY}:{image_tag}'
+        host = cf.install_milvus_operator_specific_config(NAMESPACE, MILVUS_MODE, release_name, image,
+                                                          rate_limit_enable, collection_rate_limit)
+
+        # 2. connect milvus
+        self.connection_wrap.add_connection(default={"host": host, "port": milvus_port})
+        self.connection_wrap.connect(alias='default')
+
+        # 3. create maximum numbers of collections in one two limit periods
+        for period in range(2):
+            log.info("test_rate_limit_create_collection: starting to check rate limit period %d" % (period+1))
+            for collection_num in range(collectionRateLimit):
+                log.info("test_rate_limit_create_collection: creating collection %d" % (collection_num+1))
+                c_name = cf.gen_unique_str(prefix)
+                collection_w = self.init_collection_wrap(c_name,
+                                                         check_task=CheckTasks.check_collection_property,
+                                                         check_items={exp_name: c_name, exp_schema: default_schema})
+        # 4. create one more collection
+            log.info("test_rate_limit_create_collection: creating one more collection")
+            c_name = cf.gen_unique_str(prefix)
+            error = {ct.err_code: 0, ct.err_msg: "Fail to create collection"}
+            if rate_limit_enable:
+                collection_w = self.init_collection_wrap(c_name, check_task=CheckTasks.err_res, check_items=error)
+                # 5. sleep 61s to verify create collections in next rate limit period
+                sleep(rate_limit_period)
+        # 6. drop maximum+1 numbers of collections in two rate limit period
+        collection_list = self.utility_wrap.list_collections()[0]
+        drop_num = 0
+        error = {ct.err_code: 0, ct.err_msg: "Fail to drop collection"}
+        for collection_object in self.collection_object_list:
+            if collection_object.collection is not None and collection_object.name in collection_list:
+                log.info("test_rate_limit_create_collection: dropping collection %d" % (drop_num + 1))
+                if drop_num <= collection_rate_limit:
+                    collection_list.remove(collection_object.name)
+                    collection_object.drop()
+                    drop_num += 1
+                else:
+                    if rate_limit_enable:
+                        collection_object.drop(check_task=CheckTasks.err_res, check_items=error)
+                        # 7. sleep 61s to verify drop collections in next rate limit period
+                        sleep(rate_limit_period)
+                    drop_num = 0
+        # 8. export milvus logs
+        label = f"app.kubernetes.io/instance={release_name}"
+        log.info('Start to export milvus pod logs')
+        read_pod_log(namespace=NAMESPACE, label_selector=label, release_name=release_name)
+
+        # 9. uninstall milvus
+        mil.uninstall(release_name, namespace=NAMESPACE)
+
+

tests/python_client/testcases/test_search.py

@@ -27,7 +27,7 @@ default_nq = ct.default_nq
 default_dim = ct.default_dim
 default_limit = ct.default_limit
 default_search_exp = "int64 >= 0"
-default_search_string_exp =  "varchar >= \"0\""
+default_search_string_exp = "varchar >= \"0\""
 default_search_mix_exp = "int64 >= 0 && varchar >= \"0\""
 default_invaild_string_exp = "varchar >= 0"
 perfix_expr = 'varchar like "0%"'