Skip to content

Commit

Permalink
test:[cp2.5] add search iterator v2 case for milvus client (#39447)
Browse files Browse the repository at this point in the history 
issue: #37548 
pr: #39400

Signed-off-by: qixuan <[email protected]>
  • Loading branch information
@qixuan0212
qixuan0212 authored Jan 22, 2025
1 parent 8475d8b commit a90bd9e
Show file tree
Hide file tree
Showing 4 changed files with 581 additions and 23 deletions.
19 changes: 19 additions & 0 deletions tests/python_client/base/client_v2_base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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_iterator(self, client, collection_name, data, batch_size, limit=-1, 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
Expand Down
44 changes: 24 additions & 20 deletions tests/python_client/check/func_check.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -429,26 +429,30 @@ def check_search_iterator(search_res, func_name, check_items):
search_iterator = search_res
pk_list = []
while True:
res = search_iterator.next()
if len(res) == 0:
log.info("search iteration finished, close")
search_iterator.close()
break
if check_items.get("batch_size", None):
assert len(res) <= check_items["batch_size"]
if check_items.get("radius", None):
for distance in res.distances():
if check_items["metric_type"] == "L2":
assert distance < check_items["radius"]
else:
assert distance > check_items["radius"]
if check_items.get("range_filter", None):
for distance in res.distances():
if check_items["metric_type"] == "L2":
assert distance >= check_items["range_filter"]
else:
assert distance <= check_items["range_filter"]
pk_list.extend(res.ids())
try:
res = search_iterator.next()
if len(res) == 0:
log.info("search iteration finished, close")
search_iterator.close()
break
if check_items.get("batch_size", None):
assert len(res) <= check_items["batch_size"]
if check_items.get("radius", None):
for distance in res.distances():
if check_items["metric_type"] == "L2":
assert distance < check_items["radius"]
else:
assert distance > check_items["radius"]
if check_items.get("range_filter", None):
for distance in res.distances():
if check_items["metric_type"] == "L2":
assert distance >= check_items["range_filter"]
else:
assert distance <= check_items["range_filter"]
pk_list.extend(res.ids())
except Exception as e:
assert check_items["err_msg"] in str(e)
return False
assert len(pk_list) == len(set(pk_list))
log.info("check: total %d results" % len(pk_list))

Expand Down
54 changes: 51 additions & 3 deletions tests/python_client/milvus_client/test_milvus_client_rbac.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -611,13 +611,61 @@ def teardown_method(self, method):
roles, _ = self.list_roles(client)
for role in roles:
if role not in ['admin', 'public']:
privileges, _ = self.describe_role(client, role)
if privileges:
for privilege in privileges:
role_info, _ = self.describe_role(client, role)
if role_info:
for privilege in role_info.get("privileges", []):
self.revoke_privilege(client, role, privilege["object_type"],
privilege["privilege"], privilege["object_name"])
self.drop_role(client, role)
roles, _ = self.list_roles(client)
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
"""
batch_size = 20
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_iterator(client, collection_name, vectors_to_search, batch_size, 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", "Insert", collection_name, my_db)
self.drop_collection(client, collection_name)
self.using_database(client, 'default')
self.drop_collection(client, collection_name)
Loading

0 comments on commit a90bd9e

Please sign in to comment.