test:Add search Iterator v2 test case for milvus client (milvus-io#39120

)

issue: milvus-io#37548

Signed-off-by: qixuan <[email protected]>

tests/python_client/base/client_v2_base.py

@@ -153,6 +153,25 @@ def search(self, client, collection_name, data, limit=10, filter=None, output_fi
                                        **kwargs).run()
         return res, check_result
 
+    @trace()
+    def search_interator(self, client, collection_name, data, batch_size=20, limit=100, filter=None, output_fields=None,
+                         search_params=None, timeout=None, check_task=None, check_items=None, **kwargs):
+        timeout = TIMEOUT if timeout is None else timeout
+        kwargs.update({"timeout": timeout})
+        func_name = sys._getframe().f_code.co_name
+        res, check = api_request([client.search_iterator, collection_name, data, batch_size, filter, limit,
+                                  output_fields, search_params], **kwargs)
+        if any(k in kwargs for k in ['use_rbac_mul_db', 'use_mul_db']):
+            self.using_database(client, kwargs.get('another_db'))
+        if kwargs.get('use_alias', False) is True:
+            alias = collection_name
+            self.alter_alias(client, kwargs.get('another_collection'), alias)
+        check_result = ResponseChecker(res, func_name, check_task, check_items, check,
+                                       collection_name=collection_name, data=data, batch_size=batch_size, filter=filter,
+                                       limit=limit, output_fields=output_fields, search_params=search_params,
+                                       **kwargs).run()
+        return res, check_result
+
     @trace()
     def query(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
         timeout = TIMEOUT if timeout is None else timeout
@@ -165,6 +184,18 @@ def query(self, client, collection_name, timeout=None, check_task=None, check_it
                                        **kwargs).run()
         return res, check_result
 
+    @trace()
+    def query_iterator(self, client, collection_name, timeout=None, check_task=None, check_items=None, **kwargs):
+        timeout = TIMEOUT if timeout is None else timeout
+        kwargs.update({"timeout": timeout})
+
+        func_name = sys._getframe().f_code.co_name
+        res, check = api_request([client.query_iterator, collection_name], **kwargs)
+        check_result = ResponseChecker(res, func_name, check_task, check_items, check,
+                                       collection_name=collection_name,
+                                       **kwargs).run()
+        return res, check_result
+
     @trace()
     def get(self, client, collection_name, ids, output_fields=None,
             timeout=None, check_task=None, check_items=None, **kwargs):

tests/python_client/milvus_client/test_milvus_client_rbac.py

@@ -1,4 +1,5 @@
 import time
+import numpy as np
 import pytest
 from common import common_func as cf
 from common import common_type as ct
@@ -621,3 +622,54 @@ def teardown_method(self, method):
         assert len(roles) == 2
 
         super().teardown_method(method)
+
+    @pytest.mark.skip("common.security.authorizationEnabled need to be set to true")
+    def test_milvus_client_search_iterator_rbac_mul_db(self):
+        """
+        target: test search iterator(high level api) normal case about mul db by rbac
+        method: create connection, collection, insert and search iterator
+        expected: search iterator permission deny after switch to  no permission db
+        """
+        uri = f"http://{cf.param_info.param_host}:{cf.param_info.param_port}"
+        client, _ = self.init_milvus_client(uri=uri, token="root:Milvus")
+        my_db = cf.gen_unique_str(prefix)
+        self.create_database(client, my_db)
+        collection_name = cf.gen_unique_str(prefix)
+        self.using_database(client, my_db)
+        # 1. create collection
+        self.create_collection(client, collection_name, default_dim, consistency_level="Bounded")
+        # 2. insert
+        rng = np.random.default_rng(seed=19530)
+        rows = [{default_primary_key_field_name: i, default_vector_field_name: list(rng.random((1, default_dim))[0]),
+                 default_float_field_name: i * 1.0, default_string_field_name: str(i)} for i in range(default_nb)]
+        self.insert(client, collection_name, rows)
+        self.flush(client, collection_name)
+        self.using_database(client, "default")
+        # 3. create collection
+        self.create_collection(client, collection_name, default_dim, consistency_level="Bounded")
+        # 4. insert
+        self.insert(client, collection_name, rows)
+        self.flush(client, collection_name)
+        user_name = cf.gen_unique_str(user_pre)
+        role_name = cf.gen_unique_str(role_pre)
+        password = cf.gen_str_by_length()
+        self.create_user(client, user_name=user_name, password=password)
+        self.create_role(client, role_name=role_name)
+        self.grant_role(client, user_name=user_name, role_name=role_name)
+        self.grant_privilege(client, role_name, "Collection", "Search", collection_name, 'default')
+        self.grant_privilege(client, role_name, "Collection", "Insert", collection_name, my_db)
+        client, _ = self.init_milvus_client(uri=uri, user=user_name, password=password)
+
+        # 5. search_iterator
+        vectors_to_search = rng.random((1, default_dim))
+        self.search_interator(client, collection_name, vectors_to_search, use_rbac_mul_db=True, another_db=my_db,
+                              check_task=CheckTasks.check_permission_deny)
+        client, _ = self.init_milvus_client(uri=uri, token="root:Milvus")
+        self.revoke_privilege(client, role_name, "Collection", "Search", collection_name, 'default')
+        self.revoke_privilege(client, role_name, "Collection", "Insert", collection_name, my_db)
+        self.revoke_role(client, user_name=user_name, role_name=role_name)
+        self.release_collection(client, collection_name)
+        self.drop_collection(client, collection_name)
+        self.using_database(client, 'default')
+        self.release_collection(client, collection_name)
+        self.drop_collection(client, collection_name)