format code of test_utility.py (#12505)

Signed-off-by: zhuwenxing <wenxing.zhu@zilliz.com>
pull/12564/head
zhuwenxing 2021-12-01 20:13:32 +08:00 committed by GitHub
parent dde2795bcb
commit 7f399b077d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 11 deletions

View File

@ -19,6 +19,7 @@ default_nb = ct.default_nb
num_loaded_entities = "num_loaded_entities"
num_total_entities = "num_total_entities"
def get_segment_distribution(res):
"""
Get segment distribution
@ -1347,7 +1348,7 @@ class TestUtilityAdvanced(TestcaseBase):
c_names = []
num = 5
for i in range(thread_num*num):
for i in range(thread_num * num):
c_name = cf.gen_unique_str(prefix)
self.init_collection_wrap(c_name)
c_names.append(c_name)
@ -1359,7 +1360,7 @@ class TestUtilityAdvanced(TestcaseBase):
assert not self.utility_wrap.has_collection(name)[0]
for i in range(thread_num):
x = threading.Thread(target=create_and_drop_collection, args=(c_names[i*num:(i+1)*num],))
x = threading.Thread(target=create_and_drop_collection, args=(c_names[i * num:(i + 1) * num],))
threads.append(x)
x.start()
for t in threads:
@ -1455,16 +1456,16 @@ class TestUtilityAdvanced(TestcaseBase):
# prepare load balance params
res, _ = self.utility_wrap.get_query_segment_info(c_name)
segment_distribution = get_segment_distribution(res)
all_querynodes = [node["identifier"]for node in ms.query_nodes]
all_querynodes = [node["identifier"] for node in ms.query_nodes]
assert len(all_querynodes) > 1
all_querynodes = sorted(all_querynodes,
key=lambda x: len(segment_distribution[x]["sealed"]) \
key=lambda x: len(segment_distribution[x]["sealed"])
if x in segment_distribution else 0, reverse=True)
src_node_id = all_querynodes[0]
des_node_ids = all_querynodes[1:]
sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
# load balance
self.utility_wrap.load_balance(src_node_id,des_node_ids,sealed_segment_ids)
self.utility_wrap.load_balance(src_node_id, des_node_ids, sealed_segment_ids)
# get segments distribution after load balance
res, _ = self.utility_wrap.get_query_segment_info(c_name)
segment_distribution = get_segment_distribution(res)