test_utility.py 78.6 KB
Newer Older
T
ThreadDao 已提交
1
import threading
2
import time
T
ThreadDao 已提交
3

紫晴 已提交
4
import pytest
5
from pymilvus.exceptions import MilvusException
6
from base.client_base import TestcaseBase
7
from base.utility_wrapper import ApiUtilityWrapper
紫晴 已提交
8
from utils.util_log import test_log as log
D
del-zhenwu 已提交
9 10
from common import common_func as cf
from common import common_type as ct
D
del-zhenwu 已提交
11
from common.common_type import CaseLabel, CheckTasks
Z
zhuwenxing 已提交
12
from common.milvus_sys import MilvusSys
紫晴 已提交
13

D
del-zhenwu 已提交
14 15
prefix = "utility"
default_schema = cf.gen_default_collection_schema()
16
default_int64_field_name = ct.default_int64_field_name
D
del-zhenwu 已提交
17 18
default_field_name = ct.default_float_vec_field_name
default_index_params = {"index_type": "IVF_SQ8", "metric_type": "L2", "params": {"nlist": 64}}
19 20
default_dim = ct.default_dim
default_nb = ct.default_nb
21 22
num_loaded_entities = "num_loaded_entities"
num_total_entities = "num_total_entities"
23
loading_progress = "loading_progress"
24 25
num_loaded_partitions = "num_loaded_partitions"
not_loaded_partitions = "not_loaded_partitions"
紫晴 已提交
26

27

28
class TestUtilityParams(TestcaseBase):
D
del-zhenwu 已提交
29 30
    """ Test case of index interface """

B
binbin 已提交
31 32 33 34
    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_metric_type(self, request):
        if request.param == [] or request.param == "":
            pytest.skip("metric empty is valid for distance calculation")
35 36 37 38 39 40 41 42 43 44 45 46 47 48
        if isinstance(request.param, str):
            pytest.skip("string is valid type for metric")
        yield request.param

    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_metric_value(self, request):
        if request.param == [] or request.param == "":
            pytest.skip("metric empty is valid for distance calculation")
        if not isinstance(request.param, str):
            pytest.skip("Skip invalid type for metric")
        yield request.param

    @pytest.fixture(scope="function", params=["JACCARD", "Superstructure", "Substructure"])
    def get_not_support_metric(self, request):
B
binbin 已提交
49 50
        yield request.param

51 52 53 54
    @pytest.fixture(scope="function", params=["metric_type", "metric"])
    def get_support_metric_field(self, request):
        yield request.param

55 56 57 58 59 60 61 62 63
    @pytest.fixture(scope="function", params=ct.get_invalid_strs)
    def get_invalid_partition_names(self, request):
        if isinstance(request.param, list):
            if len(request.param) == 0:
                pytest.skip("empty is valid for partition")
        if request.param is None:
            pytest.skip("None is valid for partition")
        yield request.param

B
binbin 已提交
64 65 66 67 68
    """
    ******************************************************************
    #  The followings are invalid cases
    ******************************************************************
    """
69

70
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
71 72 73 74 75 76
    def test_has_collection_name_invalid(self, get_invalid_collection_name):
        """
        target: test has_collection with error collection name
        method: input invalid name
        expected: raise exception
        """
D
del-zhenwu 已提交
77
        self._connect()
D
del-zhenwu 已提交
78
        c_name = get_invalid_collection_name
D
del-zhenwu 已提交
79
        if isinstance(c_name, str) and c_name:
Y
yanliang567 已提交
80 81 82 83
            self.utility_wrap.has_collection(
                c_name,
                check_task=CheckTasks.err_res,
                check_items={ct.err_code: 1, ct.err_msg: "Invalid collection name"})
84 85
        # elif not isinstance(c_name, str): self.utility_wrap.has_collection(c_name, check_task=CheckTasks.err_res,
        # check_items={ct.err_code: 1, ct.err_msg: "illegal"})
D
del-zhenwu 已提交
86

87
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
88 89 90 91 92 93
    def test_has_partition_collection_name_invalid(self, get_invalid_collection_name):
        """
        target: test has_partition with error collection name
        method: input invalid name
        expected: raise exception
        """
D
del-zhenwu 已提交
94
        self._connect()
D
del-zhenwu 已提交
95 96
        c_name = get_invalid_collection_name
        p_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
97
        if isinstance(c_name, str) and c_name:
Y
yanliang567 已提交
98 99 100 101
            self.utility_wrap.has_partition(
                c_name, p_name,
                check_task=CheckTasks.err_res,
                check_items={ct.err_code: 1, ct.err_msg: "Invalid"})
D
del-zhenwu 已提交
102

103
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
104 105 106 107 108 109 110
    def test_has_partition_name_invalid(self, get_invalid_partition_name):
        """
        target: test has_partition with error partition name
        method: input invalid name
        expected: raise exception
        """
        self._connect()
111
        ut = ApiUtilityWrapper()
D
del-zhenwu 已提交
112 113
        c_name = cf.gen_unique_str(prefix)
        p_name = get_invalid_partition_name
D
del-zhenwu 已提交
114
        if isinstance(p_name, str) and p_name:
Y
yanliang567 已提交
115 116 117 118
            ex, _ = ut.has_partition(
                c_name, p_name,
                check_task=CheckTasks.err_res,
                check_items={ct.err_code: 1, ct.err_msg: "Invalid"})
D
del-zhenwu 已提交
119

120
    @pytest.mark.tags(CaseLabel.L2)
121 122 123 124 125 126
    def test_drop_collection_name_invalid(self, get_invalid_collection_name):
        self._connect()
        error = f'`collection_name` value {get_invalid_collection_name} is illegal'
        self.utility_wrap.drop_collection(get_invalid_collection_name, check_task=CheckTasks.err_res,
                                          check_items={ct.err_code: 1, ct.err_msg: error})

D
del-zhenwu 已提交
127
    # TODO: enable
128
    @pytest.mark.tags(CaseLabel.L2)
129
    def test_list_collections_using_invalid(self):
D
del-zhenwu 已提交
130 131 132 133 134 135 136
        """
        target: test list_collections with invalid using
        method: input invalid name
        expected: raise exception
        """
        self._connect()
        using = "empty"
D
del-zhenwu 已提交
137
        ut = ApiUtilityWrapper()
138
        ex, _ = ut.list_collections(using=using, check_task=CheckTasks.err_res,
Y
yanliang567 已提交
139
                                    check_items={ct.err_code: 0, ct.err_msg: "should create connect"})
D
del-zhenwu 已提交
140 141 142 143 144 145 146 147

    @pytest.mark.tags(CaseLabel.L1)
    def test_index_process_invalid_name(self, get_invalid_collection_name):
        """
        target: test building_process
        method: input invalid name
        expected: raise exception
        """
Y
yanliang567 已提交
148
        pass
149 150 151
        # self._connect() c_name = get_invalid_collection_name ut = ApiUtilityWrapper() if isinstance(c_name,
        # str) and c_name: ex, _ = ut.index_building_progress(c_name, check_items={ct.err_code: 1, ct.err_msg:
        # "Invalid collection name"})
D
del-zhenwu 已提交
152 153 154 155 156 157 158 159 160 161 162 163

    # TODO: not support index name
    @pytest.mark.tags(CaseLabel.L1)
    def _test_index_process_invalid_index_name(self, get_invalid_index_name):
        """
        target: test building_process
        method: input invalid index name
        expected: raise exception
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
        index_name = get_invalid_index_name
164
        ut = ApiUtilityWrapper()
D
del-zhenwu 已提交
165 166 167 168
        ex, _ = ut.index_building_progress(c_name, index_name)
        log.error(str(ex))
        assert "invalid" or "illegal" in str(ex)

169
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
170 171 172 173 174 175
    def test_wait_index_invalid_name(self, get_invalid_collection_name):
        """
        target: test wait_index
        method: input invalid name
        expected: raise exception
        """
Y
yanliang567 已提交
176 177 178 179
        pass
        # self._connect()
        # c_name = get_invalid_collection_name
        # ut = ApiUtilityWrapper()
D
del-zhenwu 已提交
180
        # if isinstance(c_name, str) and c_name:
181 182 183
        #     ex, _ = ut.wait_for_index_building_complete(c_name,
        #                                                 check_items={ct.err_code: 1,
        #                                                              ct.err_msg: "Invalid collection name"})
D
del-zhenwu 已提交
184 185

    @pytest.mark.tags(CaseLabel.L1)
186
    def _test_wait_index_invalid_index_name(self, get_invalid_index_name):
D
del-zhenwu 已提交
187 188 189 190 191 192 193 194
        """
        target: test wait_index
        method: input invalid index name
        expected: raise exception
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
        index_name = get_invalid_index_name
195
        ut = ApiUtilityWrapper()
D
del-zhenwu 已提交
196 197 198 199
        ex, _ = ut.wait_for_index_building_complete(c_name, index_name)
        log.error(str(ex))
        assert "invalid" or "illegal" in str(ex)

200
    @pytest.mark.tags(CaseLabel.L2)
201 202 203 204 205 206 207 208 209
    @pytest.mark.parametrize("invalid_c_name", ["12-s", "12 s", "(mn)", "中文", "%$#"])
    def test_loading_progress_invalid_collection_name(self, invalid_c_name):
        """
        target: test loading progress with invalid collection name
        method: input invalid collection name
        expected: raise exception
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
210
        df = cf.gen_default_dataframe_data()
211 212 213 214 215
        self.collection_wrap.construct_from_dataframe(c_name, df, primary_field=ct.default_int64_field_name)
        self.collection_wrap.load()
        error = {ct.err_code: 1, ct.err_msg: "Invalid collection name: {}".format(invalid_c_name)}
        self.utility_wrap.loading_progress(invalid_c_name, check_task=CheckTasks.err_res, check_items=error)

216
    @pytest.mark.tags(CaseLabel.L2)
217 218 219 220 221 222 223 224
    def test_loading_progress_not_existed_collection_name(self):
        """
        target: test loading progress with invalid collection name
        method: input invalid collection name
        expected: raise exception
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
225
        df = cf.gen_default_dataframe_data()
226 227 228 229 230
        self.collection_wrap.construct_from_dataframe(c_name, df, primary_field=ct.default_int64_field_name)
        self.collection_wrap.load()
        error = {ct.err_code: 1, ct.err_msg: "describe collection failed: can't find collection"}
        self.utility_wrap.loading_progress("not_existed_name", check_task=CheckTasks.err_res, check_items=error)

231
    @pytest.mark.tags(CaseLabel.L2)
232
    @pytest.mark.xfail(reason="pymilvus issue #677")
233 234 235 236 237 238 239 240 241 242 243 244 245
    def test_loading_progress_invalid_partition_names(self, get_invalid_partition_names):
        """
        target: test loading progress with invalid partition names
        method: input invalid partition names
        expected: raise an exception
        """
        collection_w = self.init_collection_general(prefix)[0]
        partition_names = get_invalid_partition_names
        err_msg = {ct.err_code: 0, ct.err_msg: "`partition_name_array` value {} is illegal".format(partition_names)}
        collection_w.load()
        self.utility_wrap.loading_progress(collection_w.name, partition_names,
                                           check_task=CheckTasks.err_res, check_items=err_msg)

B
binbin 已提交
246
    @pytest.mark.tags(CaseLabel.L1)
247 248 249 250 251 252 253 254 255 256
    @pytest.mark.parametrize("partition_names", [[ct.default_tag], [ct.default_partition_name, ct.default_tag]])
    def test_loading_progress_not_existed_partitions(self, partition_names):
        """
        target: test loading progress with not existed partitions
        method: input all or part not existed partition names
        expected: raise exception
        """
        collection_w = self.init_collection_general(prefix)[0]
        log.debug(collection_w.num_entities)
        collection_w.load()
B
binbin 已提交
257
        err_msg = {ct.err_code: -1, ct.err_msg: f"Partitions not exist: [{ct.default_tag}]"}
258 259 260
        self.utility_wrap.loading_progress(collection_w.name, partition_names,
                                           check_task=CheckTasks.err_res, check_items=err_msg)

261
    @pytest.mark.tags(CaseLabel.L2)
262 263 264 265 266 267 268 269 270 271 272 273 274
    def test_wait_for_loading_collection_not_existed(self):
        """
        target: test wait for loading
        method: input collection not created before
        expected: raise exception
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
        self.utility_wrap.wait_for_loading_complete(
            c_name,
            check_task=CheckTasks.err_res,
            check_items={ct.err_code: 1, ct.err_msg: "can't find collection"})

275
    @pytest.mark.tags(CaseLabel.L2)
276 277 278 279 280 281 282 283 284 285 286 287 288
    def test_wait_for_loading_partition_not_existed(self):
        """
        target: test wait for loading
        method: input partition not created before
        expected: raise exception
        """
        self._connect()
        collection_w = self.init_collection_wrap()
        self.utility_wrap.wait_for_loading_complete(
            collection_w.name, partition_names=[ct.default_tag],
            check_task=CheckTasks.err_res,
            check_items={ct.err_code: 1, ct.err_msg: f'partitionID of partitionName:{ct.default_tag} can not be find'})

289
    @pytest.mark.tags(CaseLabel.L2)
T
ThreadDao 已提交
290
    def test_drop_collection_not_existed(self):
291 292 293 294 295 296 297
        """
        target: test drop an not existed collection
        method: drop a not created collection
        expected: raise exception
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
298
        error = {ct.err_code: 1, ct.err_msg: f"DescribeCollection failed: can't find collection: {c_name}"}
299 300
        self.utility_wrap.drop_collection(c_name, check_task=CheckTasks.err_res, check_items=error)

301
    @pytest.mark.tags(CaseLabel.L2)
302 303 304 305 306 307 308 309 310 311 312 313
    def test_calc_distance_left_vector_invalid_type(self, get_invalid_vector_dict):
        """
        target: test calculated distance with invalid vectors
        method: input invalid vectors type
        expected: raise exception
        """
        self._connect()
        invalid_vector = get_invalid_vector_dict
        if not isinstance(invalid_vector, dict):
            self.utility_wrap.calc_distance(invalid_vector, invalid_vector,
                                            check_task=CheckTasks.err_res,
                                            check_items={"err_code": 1,
314 315
                                                         "err_msg": "vectors_left value {} "
                                                                    "is illegal".format(invalid_vector)})
316

317
    @pytest.mark.tags(CaseLabel.L2)
318 319 320 321 322 323 324 325 326 327 328 329
    def test_calc_distance_left_vector_invalid_value(self, get_invalid_vector_dict):
        """
        target: test calculated distance with invalid vectors
        method: input invalid vectors value
        expected: raise exception
        """
        self._connect()
        invalid_vector = get_invalid_vector_dict
        if isinstance(invalid_vector, dict):
            self.utility_wrap.calc_distance(invalid_vector, invalid_vector,
                                            check_task=CheckTasks.err_res,
                                            check_items={"err_code": 1,
330 331
                                                         "err_msg": "vectors_left value {} "
                                                                    "is illegal".format(invalid_vector)})
332

333
    @pytest.mark.tags(CaseLabel.L2)
334 335 336 337 338 339 340 341 342 343 344 345 346 347
    def test_calc_distance_right_vector_invalid_type(self, get_invalid_vector_dict):
        """
        target: test calculated distance with invalid vectors
        method: input invalid vectors type
        expected: raise exception
        """
        self._connect()
        invalid_vector = get_invalid_vector_dict
        vector = cf.gen_vectors(default_nb, default_dim)
        op_l = {"float_vectors": vector}
        if not isinstance(invalid_vector, dict):
            self.utility_wrap.calc_distance(op_l, invalid_vector,
                                            check_task=CheckTasks.err_res,
                                            check_items={"err_code": 1,
348 349
                                                         "err_msg": "vectors_right value {} "
                                                                    "is illegal".format(invalid_vector)})
350

351
    @pytest.mark.tags(CaseLabel.L2)
352 353 354 355 356 357 358 359 360 361 362 363 364 365
    def test_calc_distance_right_vector_invalid_value(self, get_invalid_vector_dict):
        """
        target: test calculated distance with invalid vectors
        method: input invalid vectors value
        expected: raise exception
        """
        self._connect()
        invalid_vector = get_invalid_vector_dict
        vector = cf.gen_vectors(default_nb, default_dim)
        op_l = {"float_vectors": vector}
        if isinstance(invalid_vector, dict):
            self.utility_wrap.calc_distance(op_l, invalid_vector,
                                            check_task=CheckTasks.err_res,
                                            check_items={"err_code": 1,
366 367
                                                         "err_msg": "vectors_right value {} "
                                                                    "is illegal".format(invalid_vector)})
368

B
binbin 已提交
369
    @pytest.mark.tags(CaseLabel.L2)
370
    def test_calc_distance_invalid_metric_type(self, get_support_metric_field, get_invalid_metric_type):
B
binbin 已提交
371 372 373 374 375 376 377 378 379 380
        """
        target: test calculated distance with invalid metric
        method: input invalid metric
        expected: raise exception
        """
        self._connect()
        vectors_l = cf.gen_vectors(default_nb, default_dim)
        vectors_r = cf.gen_vectors(default_nb, default_dim)
        op_l = {"float_vectors": vectors_l}
        op_r = {"float_vectors": vectors_r}
381
        metric_field = get_support_metric_field
B
binbin 已提交
382
        metric = get_invalid_metric_type
383
        params = {metric_field: metric}
384 385 386
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.err_res,
                                        check_items={"err_code": 1,
387 388
                                                     "err_msg": "params value {{'metric': {}}} "
                                                                "is illegal".format(metric)})
389 390

    @pytest.mark.tags(CaseLabel.L2)
391
    def test_calc_distance_invalid_metric_value(self, get_support_metric_field, get_invalid_metric_value):
392 393 394 395 396 397 398 399 400 401
        """
        target: test calculated distance with invalid metric
        method: input invalid metric
        expected: raise exception
        """
        self._connect()
        vectors_l = cf.gen_vectors(default_nb, default_dim)
        vectors_r = cf.gen_vectors(default_nb, default_dim)
        op_l = {"float_vectors": vectors_l}
        op_r = {"float_vectors": vectors_r}
402
        metric_field = get_support_metric_field
403
        metric = get_invalid_metric_value
404
        params = {metric_field: metric}
405 406 407
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.err_res,
                                        check_items={"err_code": 1,
408 409
                                                     "err_msg": "{} metric type is invalid for "
                                                                "float vector".format(metric)})
410 411

    @pytest.mark.tags(CaseLabel.L2)
412
    def test_calc_distance_not_support_metric(self, get_support_metric_field, get_not_support_metric):
413 414 415 416 417 418 419 420 421 422
        """
        target: test calculated distance with invalid metric
        method: input invalid metric
        expected: raise exception
        """
        self._connect()
        vectors_l = cf.gen_vectors(default_nb, default_dim)
        vectors_r = cf.gen_vectors(default_nb, default_dim)
        op_l = {"float_vectors": vectors_l}
        op_r = {"float_vectors": vectors_r}
423
        metric_field = get_support_metric_field
424
        metric = get_not_support_metric
425
        params = {metric_field: metric}
B
binbin 已提交
426 427 428
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.err_res,
                                        check_items={"err_code": 1,
429 430
                                                     "err_msg": "{} metric type is invalid for "
                                                                "float vector".format(metric)})
B
binbin 已提交
431

432
    @pytest.mark.tags(CaseLabel.L2)
433
    def test_calc_distance_invalid_using(self, get_support_metric_field):
434 435 436 437 438 439
        """
        target: test calculated distance with invalid using
        method: input invalid using
        expected: raise exception
        """
        self._connect()
B
binbin 已提交
440 441
        vectors_l = cf.gen_vectors(default_nb, default_dim)
        vectors_r = cf.gen_vectors(default_nb, default_dim)
442 443
        op_l = {"float_vectors": vectors_l}
        op_r = {"float_vectors": vectors_r}
444 445
        metric_field = get_support_metric_field
        params = {metric_field: "L2", "sqrt": True}
446 447 448 449 450 451
        using = "empty"
        self.utility_wrap.calc_distance(op_l, op_r, params, using=using,
                                        check_task=CheckTasks.err_res,
                                        check_items={"err_code": 1,
                                                     "err_msg": "should create connect"})

452
    @pytest.mark.tags(CaseLabel.L2)
453 454 455 456 457 458 459 460 461 462 463 464
    def test_calc_distance_not_match_dim(self):
        """
        target: test calculated distance with invalid vectors
        method: input invalid vectors type and value
        expected: raise exception
        """
        self._connect()
        dim = 129
        vector_l = cf.gen_vectors(default_nb, default_dim)
        vector_r = cf.gen_vectors(default_nb, dim)
        op_l = {"float_vectors": vector_l}
        op_r = {"float_vectors": vector_r}
B
binbin 已提交
465 466 467 468 469 470 471
        self.utility_wrap.calc_distance(op_l, op_r,
                                        check_task=CheckTasks.err_res,
                                        check_items={"err_code": 1,
                                                     "err_msg": "Cannot calculate distance between "
                                                                "vectors with different dimension"})

    @pytest.mark.tags(CaseLabel.L2)
472
    def test_calc_distance_collection_before_load(self, get_support_metric_field):
B
binbin 已提交
473 474 475 476 477 478 479
        """
        target: test calculated distance when entities is not ready
        method: calculate distance before load
        expected: raise exception
        """
        self._connect()
        nb = 10
480 481
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb,
                                                                               is_index=True)
B
binbin 已提交
482 483 484 485 486
        middle = len(insert_ids) // 2
        op_l = {"ids": insert_ids[:middle], "collection": collection_w.name,
                "field": default_field_name}
        op_r = {"ids": insert_ids[middle:], "collection": collection_w.name,
                "field": default_field_name}
487 488
        metric_field = get_support_metric_field
        params = {metric_field: "L2", "sqrt": True}
B
binbin 已提交
489 490 491
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.err_res,
                                        check_items={"err_code": 1,
492 493
                                                     "err_msg": "collection {} was not "
                                                                "loaded into memory)".format(collection_w.name)})
D
del-zhenwu 已提交
494

495

496
class TestUtilityBase(TestcaseBase):
D
del-zhenwu 已提交
497 498
    """ Test case of index interface """

499 500 501 502
    @pytest.fixture(scope="function", params=["metric_type", "metric"])
    def metric_field(self, request):
        yield request.param

503 504 505 506 507 508 509 510
    @pytest.fixture(scope="function", params=[True, False])
    def sqrt(self, request):
        yield request.param

    @pytest.fixture(scope="function", params=["L2", "IP"])
    def metric(self, request):
        yield request.param

511
    @pytest.fixture(scope="function", params=["HAMMING", "TANIMOTO"])
512 513 514
    def metric_binary(self, request):
        yield request.param

D
del-zhenwu 已提交
515 516 517 518 519 520 521
    @pytest.mark.tags(CaseLabel.L1)
    def test_has_collection(self):
        """
        target: test has_collection with collection name
        method: input collection name created before
        expected: True
        """
D
del-zhenwu 已提交
522 523
        cw = self.init_collection_wrap()
        res, _ = self.utility_wrap.has_collection(cw.name)
D
del-zhenwu 已提交
524 525
        assert res is True

Y
yanliang567 已提交
526
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
527 528 529 530 531 532 533
    def test_has_collection_not_created(self):
        """
        target: test has_collection with collection name which is not created
        method: input random collection name
        expected: False
        """
        c_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
534 535
        _ = self.init_collection_wrap()
        res, _ = self.utility_wrap.has_collection(c_name)
D
del-zhenwu 已提交
536 537 538 539 540 541 542 543 544 545
        assert res is False

    @pytest.mark.tags(CaseLabel.L1)
    def test_has_collection_after_drop(self):
        """
        target: test has_collection with collection name droped before
        method: input random collection name
        expected: False
        """
        c_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
546 547
        cw = self.init_collection_wrap(name=c_name)
        res, _ = self.utility_wrap.has_collection(c_name)
D
del-zhenwu 已提交
548
        assert res is True
D
del-zhenwu 已提交
549 550
        cw.drop()
        res, _ = self.utility_wrap.has_collection(c_name)
D
del-zhenwu 已提交
551 552
        assert res is False

553
    @pytest.mark.tags(CaseLabel.L1)
D
del-zhenwu 已提交
554 555 556 557 558 559 560
    def test_has_partition(self):
        """
        target: test has_partition with partition name
        method: input collection name and partition name created before
        expected: True
        """
        c_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
561 562 563 564
        p_name = cf.gen_unique_str(prefix)
        cw = self.init_collection_wrap(name=c_name)
        self.init_partition_wrap(cw, p_name)
        res, _ = self.utility_wrap.has_partition(c_name, p_name)
D
del-zhenwu 已提交
565 566
        assert res is True

567
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
568 569 570 571 572 573 574 575
    def test_has_partition_not_created(self):
        """
        target: test has_partition with partition name
        method: input collection name, and partition name not created before
        expected: True
        """
        c_name = cf.gen_unique_str(prefix)
        p_name = cf.gen_unique_str()
D
del-zhenwu 已提交
576 577
        self.init_collection_wrap(name=c_name)
        res, _ = self.utility_wrap.has_partition(c_name, p_name)
D
del-zhenwu 已提交
578 579 580 581 582 583 584 585 586 587 588
        assert res is False

    @pytest.mark.tags(CaseLabel.L1)
    def test_has_partition_after_drop(self):
        """
        target: test has_partition with partition name
        method: input collection name, and partition name dropped
        expected: True
        """
        c_name = cf.gen_unique_str(prefix)
        p_name = cf.gen_unique_str()
D
del-zhenwu 已提交
589 590 591
        cw = self.init_collection_wrap(name=c_name)
        pw = self.init_partition_wrap(cw, p_name)
        res, _ = self.utility_wrap.has_partition(c_name, p_name)
D
del-zhenwu 已提交
592
        assert res is True
D
del-zhenwu 已提交
593 594
        pw.drop()
        res, _ = self.utility_wrap.has_partition(c_name, p_name)
D
del-zhenwu 已提交
595 596
        assert res is False

597 598 599 600 601 602 603 604 605 606 607 608
    @pytest.mark.tags(CaseLabel.L2)
    def test_has_default_partition(self):
        """
        target: test has_partition with '_default' partition
        method: input collection name and partition name created before
        expected: True
        """
        c_name = cf.gen_unique_str(prefix)
        self.init_collection_wrap(name=c_name)
        res, _ = self.utility_wrap.has_partition(c_name, ct.default_partition_name)
        assert res is True

D
del-zhenwu 已提交
609 610 611 612 613 614 615 616
    @pytest.mark.tags(CaseLabel.L1)
    def test_list_collections(self):
        """
        target: test list_collections
        method: create collection, list_collections
        expected: in the result
        """
        c_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
617 618
        self.init_collection_wrap(name=c_name)
        res, _ = self.utility_wrap.list_collections()
D
del-zhenwu 已提交
619 620 621 622 623 624 625 626 627 628 629
        assert c_name in res

    # TODO: make sure all collections deleted
    @pytest.mark.tags(CaseLabel.L1)
    def _test_list_collections_no_collection(self):
        """
        target: test list_collections
        method: no collection created, list_collections
        expected: length of the result equals to 0
        """
        self._connect()
D
del-zhenwu 已提交
630
        res, _ = self.utility_wrap.list_collections()
D
del-zhenwu 已提交
631 632
        assert len(res) == 0

633
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
634 635 636 637 638 639 640 641
    def test_index_process_collection_not_existed(self):
        """
        target: test building_process
        method: input collection not created before
        expected: raise exception
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
Y
yanliang567 已提交
642
        self.utility_wrap.index_building_progress(
643 644 645
            c_name,
            check_task=CheckTasks.err_res,
            check_items={ct.err_code: 1, ct.err_msg: "can't find collection"})
D
del-zhenwu 已提交
646 647 648 649 650 651 652 653 654

    @pytest.mark.tags(CaseLabel.L1)
    def test_index_process_collection_empty(self):
        """
        target: test building_process
        method: input empty collection
        expected: no exception raised
        """
        c_name = cf.gen_unique_str(prefix)
655 656 657 658 659
        cw = self.init_collection_wrap(name=c_name)
        self.index_wrap.init_index(cw.collection, default_field_name, default_index_params)
        res, _ = self.utility_wrap.index_building_progress(c_name)
        exp_res = {'total_rows': 0, 'indexed_rows': 0}
        assert res == exp_res
D
del-zhenwu 已提交
660

661
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
662 663 664 665 666 667 668 669
    def test_index_process_collection_insert_no_index(self):
        """
        target: test building_process
        method: insert 1 entity, no index created
        expected: no exception raised
        """
        nb = 1
        c_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
670
        cw = self.init_collection_wrap(name=c_name)
D
del-zhenwu 已提交
671
        data = cf.gen_default_list_data(nb)
D
del-zhenwu 已提交
672
        cw.insert(data=data)
673 674
        error = {ct.err_code: 1, ct.err_msg: "no index is created"}
        self.utility_wrap.index_building_progress(c_name, check_task=CheckTasks.err_res, check_items=error)
D
del-zhenwu 已提交
675

D
del-zhenwu 已提交
676 677 678 679
    @pytest.mark.tags(CaseLabel.L1)
    def test_index_process_collection_index(self):
        """
        target: test building_process
680 681 682
        method: 1.insert 1024 (because minSegmentSizeToEnableIndex=1024)
                2.build(server does create index) and call building_process
        expected: indexed_rows=0
D
del-zhenwu 已提交
683
        """
684
        nb = 1024
D
del-zhenwu 已提交
685
        c_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
686
        cw = self.init_collection_wrap(name=c_name)
D
del-zhenwu 已提交
687
        data = cf.gen_default_list_data(nb)
D
del-zhenwu 已提交
688
        cw.insert(data=data)
689 690 691 692
        cw.create_index(default_field_name, default_index_params)
        res, _ = self.utility_wrap.index_building_progress(c_name)
        assert res['indexed_rows'] == 0
        assert res['total_rows'] == nb
D
del-zhenwu 已提交
693 694 695 696 697

    @pytest.mark.tags(CaseLabel.L1)
    def test_index_process_collection_indexing(self):
        """
        target: test building_process
698 699 700
        method: 1.insert 2048 entities to ensure that server will build
                2.call building_process during building
        expected: 2048 or less entities indexed
D
del-zhenwu 已提交
701
        """
702
        nb = 2048
D
del-zhenwu 已提交
703
        c_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
704
        cw = self.init_collection_wrap(name=c_name)
D
del-zhenwu 已提交
705
        data = cf.gen_default_list_data(nb)
D
del-zhenwu 已提交
706 707
        cw.insert(data=data)
        cw.create_index(default_field_name, default_index_params)
708 709 710 711 712 713 714
        start = time.time()
        while True:
            time.sleep(1)
            res, _ = self.utility_wrap.index_building_progress(c_name)
            if 0 < res['indexed_rows'] <= nb:
                break
            if time.time() - start > 5:
715
                raise MilvusException(1, f"Index build completed in more than 5s")
D
del-zhenwu 已提交
716

717
    @pytest.mark.tags(CaseLabel.L2)
D
del-zhenwu 已提交
718 719 720 721 722 723 724 725
    def test_wait_index_collection_not_existed(self):
        """
        target: test wait_index
        method: input collection not created before
        expected: raise exception
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
Y
yanliang567 已提交
726
        self.utility_wrap.wait_for_index_building_complete(
727 728 729
            c_name,
            check_task=CheckTasks.err_res,
            check_items={ct.err_code: 1, ct.err_msg: "can't find collection"})
D
del-zhenwu 已提交
730 731 732 733 734 735 736 737 738 739

    @pytest.mark.tags(CaseLabel.L1)
    def test_wait_index_collection_empty(self):
        """
        target: test wait_index
        method: input empty collection
        expected: no exception raised
        """
        self._connect()
        c_name = cf.gen_unique_str(prefix)
740 741 742 743 744 745
        cw = self.init_collection_wrap(name=c_name)
        cw.create_index(default_field_name, default_index_params)
        assert self.utility_wrap.wait_for_index_building_complete(c_name)[0]
        res, _ = self.utility_wrap.index_building_progress(c_name)
        exp_res = {'total_rows': 0, 'indexed_rows': 0}
        assert res == exp_res
D
del-zhenwu 已提交
746 747 748 749 750

    @pytest.mark.tags(CaseLabel.L1)
    def test_wait_index_collection_index(self):
        """
        target: test wait_index
Y
yanliang567 已提交
751 752
        method: insert 5000 entities, build and call wait_index
        expected: 5000 entity indexed
D
del-zhenwu 已提交
753
        """
Y
yanliang567 已提交
754
        nb = 5000
D
del-zhenwu 已提交
755
        c_name = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
756
        cw = self.init_collection_wrap(name=c_name)
D
del-zhenwu 已提交
757
        data = cf.gen_default_list_data(nb)
D
del-zhenwu 已提交
758 759 760
        cw.insert(data=data)
        cw.create_index(default_field_name, default_index_params)
        res, _ = self.utility_wrap.wait_for_index_building_complete(c_name)
Y
yanliang567 已提交
761
        assert res is True
D
del-zhenwu 已提交
762 763
        res, _ = self.utility_wrap.index_building_progress(c_name)
        assert res["indexed_rows"] == nb
D
del-zhenwu 已提交
764

765
    @pytest.mark.tags(CaseLabel.L2)
766 767 768 769
    def test_loading_progress_without_loading(self):
        """
        target: test loading progress without loading
        method: insert and flush data, call loading_progress without loading
770
        expected: raise exception
771 772 773 774 775
        """
        collection_w = self.init_collection_wrap()
        df = cf.gen_default_dataframe_data()
        collection_w.insert(df)
        assert collection_w.num_entities == ct.default_nb
776 777 778
        error = {ct.err_code: 1, ct.err_msg: {"has not been loaded into QueryNode"}}
        self.utility_wrap.loading_progress(collection_w.name,
                                           check_task=CheckTasks.err_res, check_items=error)
779

780
    @pytest.mark.tags(CaseLabel.L1)
781 782 783 784 785 786 787 788 789 790
    @pytest.mark.parametrize("nb", [ct.default_nb, 5000])
    def test_loading_progress_collection(self, nb):
        """
        target: test loading progress
        method: 1.insert flush and load 2.call loading_progress
        expected: all entities is loafed, because load is synchronous
        """
        # create, insert default_nb, flush and load
        collection_w = self.init_collection_general(prefix, insert_data=True, nb=nb)[0]
        res, _ = self.utility_wrap.loading_progress(collection_w.name)
791
        assert res[loading_progress] == '100%'
792

793
    @pytest.mark.tags(CaseLabel.L2)
794 795 796 797 798 799 800 801 802 803 804 805
    def test_loading_progress_with_async_load(self):
        """
        target: test loading progress with async collection load
        method: 1.load collection with async=True 2.loading_progress
        expected: loading part entities
        """
        collection_w = self.init_collection_wrap()
        df = cf.gen_default_dataframe_data()
        collection_w.insert(df)
        assert collection_w.num_entities == ct.default_nb
        collection_w.load(_async=True)
        res, _ = self.utility_wrap.loading_progress(collection_w.name)
806 807 808 809 810
        loading_int = cf.percent_to_int(res[loading_progress])
        if -1 != loading_int:
            assert (0 <= loading_int <= 100)
        else:
            log.info("The output of loading progress is not a string or a percentage")
811

812
    @pytest.mark.tags(CaseLabel.L2)
813 814
    def test_loading_progress_empty_collection(self):
        """
X
Xieql 已提交
815
        target: test loading_progress on an empty collection
816 817 818 819 820 821
        method: 1.create collection and no insert 2.loading_progress
        expected: 0 entities is loaded
        """
        collection_w = self.init_collection_wrap()
        collection_w.load()
        res, _ = self.utility_wrap.loading_progress(collection_w.name)
822 823
        exp_res = {loading_progress: '100%', num_loaded_partitions: 1, not_loaded_partitions: []}

824 825
        assert exp_res == res

826
    @pytest.mark.tags(CaseLabel.L1)
827 828
    def test_loading_progress_after_release(self):
        """
829 830
        target: test loading progress after release
        method: insert and flush data, call loading_progress after release
B
binbin 已提交
831
        expected: return successfully with 0%
832 833 834
        """
        collection_w = self.init_collection_general(prefix, insert_data=True)[0]
        collection_w.release()
B
binbin 已提交
835 836 837 838
        res = self.utility_wrap.loading_progress(collection_w.name)[0]
        exp_res = {loading_progress: '0%', num_loaded_partitions: 0, not_loaded_partitions: ['_default']}

        assert exp_res == res
839

840
    @pytest.mark.tags(CaseLabel.L2)
841 842 843 844
    def test_loading_progress_with_release_partition(self):
        """
        target: test loading progress after release part partitions
        method: 1.insert data into two partitions and flush
845
                2.load one partition and release one partition
846 847 848 849
        expected: loaded one partition entities
        """
        half = ct.default_nb
        # insert entities into two partitions, collection flush and load
J
jingkl 已提交
850 851
        collection_w, partition_w, _, _ = self.insert_entities_into_two_partitions_in_half(half)
        partition_w.release()
852
        res = self.utility_wrap.loading_progress(collection_w.name)[0]
853
        assert res[loading_progress] == '50%'
854

855
    @pytest.mark.tags(CaseLabel.L2)
856 857 858 859 860 861 862 863 864 865 866 867
    def test_loading_progress_with_load_partition(self):
        """
        target: test loading progress after load partition
        method: 1.insert data into two partitions and flush
                2.load one partition and loading progress
        expected: loaded one partition entities
        """
        half = ct.default_nb
        collection_w, partition_w, _, _ = self.insert_entities_into_two_partitions_in_half(half)
        collection_w.release()
        partition_w.load()
        res = self.utility_wrap.loading_progress(collection_w.name)[0]
868
        assert res[loading_progress] == '50%'
869

870
    @pytest.mark.tags(CaseLabel.L1)
871 872 873 874 875 876 877 878 879 880
    def test_loading_progress_with_partition(self):
        """
        target: test loading progress with partition
        method: 1.insert data into two partitions and flush, and load
                2.loading progress with one partition
        expected: loaded one partition entities
        """
        half = ct.default_nb
        collection_w, partition_w, _, _ = self.insert_entities_into_two_partitions_in_half(half)
        res = self.utility_wrap.loading_progress(collection_w.name, partition_names=[partition_w.name])[0]
881
        assert res[loading_progress] == '100%'
882

883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917
    @pytest.mark.tags(CaseLabel.ClusterOnly)
    def test_loading_progress_multi_replicas(self):
        """
        target: test loading progress with multi replicas
        method: 1.Create collection and insert data
                2.Load replicas and get loading progress
                3.Create partitions and insert data
                4.Get loading progress
                5.Release and re-load replicas, get loading progress
        expected: Verify loading progress result
        """
        collection_w = self.init_collection_wrap()
        collection_w.insert(cf.gen_default_dataframe_data())
        assert collection_w.num_entities == ct.default_nb
        collection_w.load(partition_names=[ct.default_partition_name], replica_number=2)
        res_collection, _ = self.utility_wrap.loading_progress(collection_w.name)
        assert res_collection == {loading_progress: '100%', num_loaded_partitions: 1, not_loaded_partitions: []}

        # create partition and insert
        partition_w = self.init_partition_wrap(collection_wrap=collection_w)
        partition_w.insert(cf.gen_default_dataframe_data(start=ct.default_nb))
        assert partition_w.num_entities == ct.default_nb
        res_part_partition, _ = self.utility_wrap.loading_progress(collection_w.name)
        assert res_part_partition == {'loading_progress': '50%', 'num_loaded_partitions': 1,
                                      'not_loaded_partitions': [partition_w.name]}

        res_part_partition, _ = self.utility_wrap.loading_progress(collection_w.name, partition_names=[partition_w.name])
        assert res_part_partition == {'loading_progress': '0%', 'num_loaded_partitions': 0,
                                      'not_loaded_partitions': [partition_w.name]}

        collection_w.release()
        collection_w.load(replica_number=2)
        res_all_partitions, _ = self.utility_wrap.loading_progress(collection_w.name)
        assert res_all_partitions == {'loading_progress': '100%', 'num_loaded_partitions': 2, 'not_loaded_partitions': []}

918
    @pytest.mark.tags(CaseLabel.L1)
919 920 921 922 923 924 925 926 927 928 929
    def test_wait_loading_collection_empty(self):
        """
        target: test wait_for_loading
        method: input empty collection
        expected: no exception raised
        """
        self._connect()
        cw = self.init_collection_wrap(name=cf.gen_unique_str(prefix))
        cw.load()
        self.utility_wrap.wait_for_loading_complete(cw.name)
        res, _ = self.utility_wrap.loading_progress(cw.name)
930
        exp_res = {loading_progress: "100%", not_loaded_partitions: [], num_loaded_partitions: 1}
931 932
        assert res == exp_res

933
    @pytest.mark.tags(CaseLabel.L1)
934 935 936 937 938 939 940 941 942
    def test_wait_for_loading_complete(self):
        """
        target: test wait for loading collection
        method: insert 10000 entities and wait for loading complete
        expected: after loading complete, loaded entities is 10000
        """
        nb = 6000
        collection_w = self.init_collection_wrap()
        df = cf.gen_default_dataframe_data(nb)
943
        collection_w.insert(df, timeout=60)
944 945 946 947
        assert collection_w.num_entities == nb
        collection_w.load(_async=True)
        self.utility_wrap.wait_for_loading_complete(collection_w.name)
        res, _ = self.utility_wrap.loading_progress(collection_w.name)
948
        assert res[loading_progress] == '100%'
949

950
    @pytest.mark.tags(CaseLabel.L0)
T
ThreadDao 已提交
951 952 953 954 955 956 957 958 959 960 961 962
    def test_drop_collection(self):
        """
        target: test utility drop collection by name
        method: input collection name and drop collection
        expected: collection is dropped
        """
        c_name = cf.gen_unique_str(prefix)
        self.init_collection_wrap(c_name)
        assert self.utility_wrap.has_collection(c_name)[0]
        self.utility_wrap.drop_collection(c_name)
        assert not self.utility_wrap.has_collection(c_name)[0]

963
    @pytest.mark.tags(CaseLabel.L0)
T
ThreadDao 已提交
964 965 966 967 968 969 970 971 972 973 974 975 976 977
    def test_drop_collection_repeatedly(self):
        """
        target: test drop collection repeatedly
        method: 1.collection.drop 2.utility.drop_collection
        expected: raise exception
        """
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(c_name)
        assert self.utility_wrap.has_collection(c_name)[0]
        collection_w.drop()
        assert not self.utility_wrap.has_collection(c_name)[0]
        error = {ct.err_code: 1, ct.err_msg: {"describe collection failed: can't find collection:"}}
        self.utility_wrap.drop_collection(c_name, check_task=CheckTasks.err_res, check_items=error)

978
    @pytest.mark.tags(CaseLabel.L2)
T
ThreadDao 已提交
979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994
    def test_drop_collection_create_repeatedly(self):
        """
        target: test repeatedly create and drop same name collection
        method: repeatedly create and drop collection
        expected: no exception
        """
        from time import sleep
        loops = 3
        c_name = cf.gen_unique_str(prefix)
        for _ in range(loops):
            self.init_collection_wrap(c_name)
            assert self.utility_wrap.has_collection(c_name)[0]
            self.utility_wrap.drop_collection(c_name)
            assert not self.utility_wrap.has_collection(c_name)[0]
            sleep(1)

995 996 997 998 999 1000 1001
    @pytest.mark.tags(CaseLabel.L1)
    def test_calc_distance_default(self):
        """
        target: test calculated distance with default params
        method: calculated distance between two random vectors
        expected: distance calculated successfully
        """
1002
        log.info("Creating connection")
1003
        self._connect()
1004
        log.info("Creating vectors for distance calculation")
1005 1006 1007 1008
        vectors_l = cf.gen_vectors(default_nb, default_dim)
        vectors_r = cf.gen_vectors(default_nb, default_dim)
        op_l = {"float_vectors": vectors_l}
        op_r = {"float_vectors": vectors_r}
1009
        log.info("Calculating distance for generated vectors")
1010 1011 1012 1013 1014 1015
        self.utility_wrap.calc_distance(op_l, op_r,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r})

    @pytest.mark.tags(CaseLabel.L2)
1016
    def test_calc_distance_default_sqrt(self, metric_field, metric):
1017 1018 1019 1020 1021
        """
        target: test calculated distance with default param
        method: calculated distance with default sqrt
        expected: distance calculated successfully
        """
B
binbin 已提交
1022
        log.info("Creating connection")
1023
        self._connect()
B
binbin 已提交
1024
        log.info("Creating vectors for distance calculation")
1025 1026 1027 1028
        vectors_l = cf.gen_vectors(default_nb, default_dim)
        vectors_r = cf.gen_vectors(default_nb, default_dim)
        op_l = {"float_vectors": vectors_l}
        op_r = {"float_vectors": vectors_r}
B
binbin 已提交
1029
        log.info("Calculating distance for generated vectors within default sqrt")
1030
        params = {metric_field: metric}
1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r,
                                                     "metric": metric})

    @pytest.mark.tags(CaseLabel.L2)
    def test_calc_distance_default_metric(self, sqrt):
        """
        target: test calculated distance with default param
        method: calculated distance with default metric
        expected: distance calculated successfully
        """
B
binbin 已提交
1044
        log.info("Creating connection")
1045
        self._connect()
B
binbin 已提交
1046
        log.info("Creating vectors for distance calculation")
1047 1048 1049 1050
        vectors_l = cf.gen_vectors(default_nb, default_dim)
        vectors_r = cf.gen_vectors(default_nb, default_dim)
        op_l = {"float_vectors": vectors_l}
        op_r = {"float_vectors": vectors_r}
B
binbin 已提交
1051
        log.info("Calculating distance for generated vectors within default metric")
1052 1053 1054 1055 1056 1057 1058 1059
        params = {"sqrt": sqrt}
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r,
                                                     "sqrt": sqrt})

    @pytest.mark.tags(CaseLabel.L2)
1060
    def test_calc_distance_binary_metric(self, metric_field, metric_binary):
1061 1062 1063 1064 1065
        """
        target: test calculate distance with binary vectors
        method: calculate distance between binary vectors
        expected: distance calculated successfully
        """
1066
        log.info("Creating connection")
1067
        self._connect()
1068
        log.info("Creating vectors for distance calculation")
1069 1070 1071 1072 1073
        nb = 10
        raw_vectors_l, vectors_l = cf.gen_binary_vectors(nb, default_dim)
        raw_vectors_r, vectors_r = cf.gen_binary_vectors(nb, default_dim)
        op_l = {"bin_vectors": vectors_l}
        op_r = {"bin_vectors": vectors_r}
1074
        log.info("Calculating distance for binary vectors")
1075
        params = {metric_field: metric_binary}
1076 1077
        vectors_l = raw_vectors_l
        vectors_r = raw_vectors_r
1078 1079 1080 1081 1082 1083 1084
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r,
                                                     "metric": metric_binary})

    @pytest.mark.tags(CaseLabel.L1)
1085
    def test_calc_distance_from_collection_ids(self, metric_field, metric, sqrt):
1086 1087 1088 1089 1090
        """
        target: test calculated distance from collection entities
        method: both left and right vectors are from collection
        expected: distance calculated successfully
        """
B
binbin 已提交
1091
        log.info("Creating connection")
1092 1093
        self._connect()
        nb = 10
1094
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb)
1095 1096 1097 1098 1099
        middle = len(insert_ids) // 2
        vectors = vectors[0].loc[:, default_field_name]
        vectors_l = vectors[:middle]
        vectors_r = []
        for i in range(middle):
1100
            vectors_r.append(vectors[middle + i])
B
binbin 已提交
1101
        log.info("Creating vectors from collections for distance calculation")
1102 1103 1104 1105
        op_l = {"ids": insert_ids[:middle], "collection": collection_w.name,
                "field": default_field_name}
        op_r = {"ids": insert_ids[middle:], "collection": collection_w.name,
                "field": default_field_name}
B
binbin 已提交
1106
        log.info("Creating vectors for entities")
1107
        params = {metric_field: metric, "sqrt": sqrt}
1108 1109 1110 1111 1112 1113 1114 1115
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r,
                                                     "metric": metric,
                                                     "sqrt": sqrt})

    @pytest.mark.tags(CaseLabel.L2)
1116
    def test_calc_distance_from_collections(self, metric_field, metric, sqrt):
1117 1118 1119 1120 1121
        """
        target: test calculated distance between entities from collections
        method: calculated distance between entities from two collections
        expected: distance calculated successfully
        """
B
binbin 已提交
1122
        log.info("Creating connection")
1123 1124 1125
        self._connect()
        nb = 10
        prefix_1 = "utility_distance"
B
binbin 已提交
1126
        log.info("Creating two collections")
1127 1128
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb)
        collection_w_1, vectors_1, _, insert_ids_1, _ = self.init_collection_general(prefix_1, True, nb)
1129 1130
        vectors_l = vectors[0].loc[:, default_field_name]
        vectors_r = vectors_1[0].loc[:, default_field_name]
B
binbin 已提交
1131
        log.info("Extracting entities from collections for distance calculating")
1132 1133 1134 1135
        op_l = {"ids": insert_ids, "collection": collection_w.name,
                "field": default_field_name}
        op_r = {"ids": insert_ids_1, "collection": collection_w_1.name,
                "field": default_field_name}
1136
        params = {metric_field: metric, "sqrt": sqrt}
B
binbin 已提交
1137
        log.info("Calculating distance for entities from two collections")
1138 1139 1140 1141 1142 1143 1144 1145
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r,
                                                     "metric": metric,
                                                     "sqrt": sqrt})

    @pytest.mark.tags(CaseLabel.L2)
1146
    def test_calc_distance_left_vector_and_collection_ids(self, metric_field, metric, sqrt):
1147 1148 1149 1150 1151
        """
        target: test calculated distance from collection entities
        method: set left vectors as random vectors, right vectors from collection
        expected: distance calculated successfully
        """
B
binbin 已提交
1152
        log.info("Creating connection")
1153 1154
        self._connect()
        nb = 10
1155
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb)
1156 1157 1158 1159 1160 1161 1162
        middle = len(insert_ids) // 2
        vectors = vectors[0].loc[:, default_field_name]
        vectors_l = cf.gen_vectors(nb, default_dim)
        vectors_r = []
        for i in range(middle):
            vectors_r.append(vectors[middle + i])
        op_l = {"float_vectors": vectors_l}
B
binbin 已提交
1163
        log.info("Extracting entities from collections for distance calculating")
1164 1165
        op_r = {"ids": insert_ids[middle:], "collection": collection_w.name,
                "field": default_field_name}
1166
        params = {metric_field: metric, "sqrt": sqrt}
B
binbin 已提交
1167
        log.info("Calculating distance between vectors and entities")
1168 1169 1170 1171 1172 1173 1174 1175
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r,
                                                     "metric": metric,
                                                     "sqrt": sqrt})

    @pytest.mark.tags(CaseLabel.L2)
1176
    def test_calc_distance_right_vector_and_collection_ids(self, metric_field, metric, sqrt):
1177 1178 1179 1180 1181
        """
        target: test calculated distance from collection entities
        method: set right vectors as random vectors, left vectors from collection
        expected: distance calculated successfully
        """
B
binbin 已提交
1182
        log.info("Creating connection")
1183 1184
        self._connect()
        nb = 10
1185
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb)
1186 1187 1188 1189
        middle = len(insert_ids) // 2
        vectors = vectors[0].loc[:, default_field_name]
        vectors_l = vectors[:middle]
        vectors_r = cf.gen_vectors(nb, default_dim)
B
binbin 已提交
1190
        log.info("Extracting entities from collections for distance calculating")
1191 1192 1193
        op_l = {"ids": insert_ids[:middle], "collection": collection_w.name,
                "field": default_field_name}
        op_r = {"float_vectors": vectors_r}
1194
        params = {metric_field: metric, "sqrt": sqrt}
B
binbin 已提交
1195
        log.info("Calculating distance between right vector and entities")
1196 1197 1198 1199 1200 1201 1202 1203
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r,
                                                     "metric": metric,
                                                     "sqrt": sqrt})

    @pytest.mark.tags(CaseLabel.L2)
1204
    def test_calc_distance_from_partition_ids(self, metric_field, metric, sqrt):
1205 1206 1207 1208 1209
        """
        target: test calculated distance from one partition entities
        method: both left and right vectors are from partition
        expected: distance calculated successfully
        """
B
binbin 已提交
1210
        log.info("Creating connection")
1211 1212
        self._connect()
        nb = 10
1213
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb, partition_num=1)
1214 1215
        partitions = collection_w.partitions
        middle = len(insert_ids) // 2
1216
        params = {metric_field: metric, "sqrt": sqrt}
1217 1218
        start = 0
        end = middle
1219
        for i in range(len(partitions)):
B
binbin 已提交
1220
            log.info("Extracting entities from partitions for distance calculating")
1221 1222
            vectors_l = vectors[i].loc[:, default_field_name]
            vectors_r = vectors[i].loc[:, default_field_name]
1223
            op_l = {"ids": insert_ids[start:end], "collection": collection_w.name,
1224
                    "partition": partitions[i].name, "field": default_field_name}
1225
            op_r = {"ids": insert_ids[start:end], "collection": collection_w.name,
1226
                    "partition": partitions[i].name, "field": default_field_name}
1227 1228
            start += middle
            end += middle
B
binbin 已提交
1229
            log.info("Calculating distance between entities from one partition")
1230 1231 1232 1233 1234 1235 1236 1237
            self.utility_wrap.calc_distance(op_l, op_r, params,
                                            check_task=CheckTasks.check_distance,
                                            check_items={"vectors_l": vectors_l,
                                                         "vectors_r": vectors_r,
                                                         "metric": metric,
                                                         "sqrt": sqrt})

    @pytest.mark.tags(CaseLabel.L2)
1238
    def test_calc_distance_from_partitions(self, metric_field, metric, sqrt):
1239 1240 1241 1242 1243
        """
        target: test calculated distance between entities from partitions
        method: calculate distance between entities from two partitions
        expected: distance calculated successfully
        """
B
binbin 已提交
1244
        log.info("Create connection")
1245 1246
        self._connect()
        nb = 10
1247
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb, partition_num=1)
1248 1249
        partitions = collection_w.partitions
        middle = len(insert_ids) // 2
1250
        params = {metric_field: metric, "sqrt": sqrt}
1251 1252
        vectors_l = vectors[0].loc[:, default_field_name]
        vectors_r = vectors[1].loc[:, default_field_name]
B
binbin 已提交
1253
        log.info("Extract entities from two partitions for distance calculating")
1254 1255 1256 1257
        op_l = {"ids": insert_ids[:middle], "collection": collection_w.name,
                "partition": partitions[0].name, "field": default_field_name}
        op_r = {"ids": insert_ids[middle:], "collection": collection_w.name,
                "partition": partitions[1].name, "field": default_field_name}
B
binbin 已提交
1258
        log.info("Calculate distance between entities from two partitions")
1259 1260 1261 1262 1263 1264 1265 1266
        self.utility_wrap.calc_distance(op_l, op_r, params,
                                        check_task=CheckTasks.check_distance,
                                        check_items={"vectors_l": vectors_l,
                                                     "vectors_r": vectors_r,
                                                     "metric": metric,
                                                     "sqrt": sqrt})

    @pytest.mark.tags(CaseLabel.L2)
1267
    def test_calc_distance_left_vectors_and_partition_ids(self, metric_field, metric, sqrt):
1268 1269 1270 1271 1272
        """
        target: test calculated distance between vectors and partition entities
        method: set left vectors as random vectors, right vectors are entities
        expected: distance calculated successfully
        """
B
binbin 已提交
1273
        log.info("Creating connection")
1274 1275
        self._connect()
        nb = 10
1276
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb, partition_num=1)
1277 1278 1279
        middle = len(insert_ids) // 2
        partitions = collection_w.partitions
        vectors_l = cf.gen_vectors(nb // 2, default_dim)
B
binbin 已提交
1280
        log.info("Extract entities from collection as right vectors")
1281
        op_l = {"float_vectors": vectors_l}
1282
        params = {metric_field: metric, "sqrt": sqrt}
1283 1284
        start = 0
        end = middle
B
binbin 已提交
1285
        log.info("Calculate distance between vector and entities")
1286 1287
        for i in range(len(partitions)):
            vectors_r = vectors[i].loc[:, default_field_name]
1288
            op_r = {"ids": insert_ids[start:end], "collection": collection_w.name,
1289
                    "partition": partitions[i].name, "field": default_field_name}
1290 1291
            start += middle
            end += middle
1292 1293 1294 1295 1296 1297 1298 1299
            self.utility_wrap.calc_distance(op_l, op_r, params,
                                            check_task=CheckTasks.check_distance,
                                            check_items={"vectors_l": vectors_l,
                                                         "vectors_r": vectors_r,
                                                         "metric": metric,
                                                         "sqrt": sqrt})

    @pytest.mark.tags(CaseLabel.L2)
1300
    def test_calc_distance_right_vectors_and_partition_ids(self, metric_field, metric, sqrt):
1301 1302 1303 1304 1305
        """
        target: test calculated distance between vectors and partition entities
        method: set right vectors as random vectors, left vectors are entities
        expected: distance calculated successfully
        """
B
binbin 已提交
1306
        log.info("Create connection")
1307 1308
        self._connect()
        nb = 10
1309
        collection_w, vectors, _, insert_ids, _ = self.init_collection_general(prefix, True, nb, partition_num=1)
1310 1311 1312 1313
        middle = len(insert_ids) // 2
        partitions = collection_w.partitions
        vectors_r = cf.gen_vectors(nb // 2, default_dim)
        op_r = {"float_vectors": vectors_r}
1314
        params = {metric_field: metric, "sqrt": sqrt}
1315 1316
        start = 0
        end = middle
1317 1318
        for i in range(len(partitions)):
            vectors_l = vectors[i].loc[:, default_field_name]
B
binbin 已提交
1319
            log.info("Extract entities from partition %d as left vector" % i)
1320
            op_l = {"ids": insert_ids[start:end], "collection": collection_w.name,
1321
                    "partition": partitions[i].name, "field": default_field_name}
1322 1323
            start += middle
            end += middle
B
binbin 已提交
1324
            log.info("Calculate distance between vector and entities from partition %d" % i)
1325 1326 1327 1328 1329 1330
            self.utility_wrap.calc_distance(op_l, op_r, params,
                                            check_task=CheckTasks.check_distance,
                                            check_items={"vectors_l": vectors_l,
                                                         "vectors_r": vectors_r,
                                                         "metric": metric,
                                                         "sqrt": sqrt})
D
del-zhenwu 已提交
1331

1332

1333
class TestUtilityAdvanced(TestcaseBase):
D
del-zhenwu 已提交
1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344
    """ Test case of index interface """

    @pytest.mark.tags(CaseLabel.L2)
    def test_has_collection_multi_collections(self):
        """
        target: test has_collection with collection name
        method: input collection name created before
        expected: True
        """
        c_name = cf.gen_unique_str(prefix)
        c_name_2 = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
1345 1346
        self.init_collection_wrap(name=c_name)
        self.init_collection_wrap(name=c_name_2)
D
del-zhenwu 已提交
1347
        for name in [c_name, c_name_2]:
D
del-zhenwu 已提交
1348
            res, _ = self.utility_wrap.has_collection(name)
D
del-zhenwu 已提交
1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359
            assert res is True

    @pytest.mark.tags(CaseLabel.L2)
    def test_list_collections_multi_collection(self):
        """
        target: test list_collections
        method: create collection, list_collections
        expected: in the result
        """
        c_name = cf.gen_unique_str(prefix)
        c_name_2 = cf.gen_unique_str(prefix)
D
del-zhenwu 已提交
1360 1361 1362
        self.init_collection_wrap(name=c_name)
        self.init_collection_wrap(name=c_name_2)
        res, _ = self.utility_wrap.list_collections()
D
del-zhenwu 已提交
1363 1364
        for name in [c_name, c_name_2]:
            assert name in res
T
ThreadDao 已提交
1365

1366
    @pytest.mark.tags(CaseLabel.L2)
T
ThreadDao 已提交
1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377
    def test_drop_multi_collection_concurrent(self):
        """
        target: test concurrent drop collection
        method: multi thread drop one collection
        expected: drop successfully
        """
        thread_num = 3
        threads = []
        c_names = []
        num = 5

1378
        for i in range(thread_num * num):
T
ThreadDao 已提交
1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389
            c_name = cf.gen_unique_str(prefix)
            self.init_collection_wrap(c_name)
            c_names.append(c_name)

        def create_and_drop_collection(names):
            for name in names:
                assert self.utility_wrap.has_collection(name)[0]
                self.utility_wrap.drop_collection(name)
                assert not self.utility_wrap.has_collection(name)[0]

        for i in range(thread_num):
1390
            x = threading.Thread(target=create_and_drop_collection, args=(c_names[i * num:(i + 1) * num],))
T
ThreadDao 已提交
1391 1392 1393 1394 1395
            threads.append(x)
            x.start()
        for t in threads:
            t.join()
        log.debug(self.utility_wrap.list_collections()[0])
1396

1397
    @pytest.mark.tags(CaseLabel.L2)
1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408
    def test_get_query_segment_info_empty_collection(self):
        """
        target: test getting query segment info of empty collection
        method: init a collection and get query segment info
        expected: length of segment is 0
        """
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        collection_w.load()
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        assert len(res) == 0
1409

1410
    @pytest.mark.tags(CaseLabel.L1)
1411
    def test_get_growing_query_segment_info(self):
1412
        """
1413
        target: test getting growing query segment info of collection with data
1414 1415
        method: init a collection, insert data, load, search, and get query segment info
        expected:
1416 1417 1418
            1. length of segment is greater than 0
            2. the sum num_rows of each segment is equal to num of entities
        """
1419 1420
        import random
        dim = 128
1421 1422 1423
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        nb = 3000
1424
        nq = 2
1425 1426
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
1427 1428 1429
        collection_w.load()
        vectors = [[random.random() for _ in range(dim)] for _ in range(nq)]
        collection_w.search(vectors, default_field_name, ct.default_search_params, ct.default_limit)
1430 1431
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        assert len(res) > 0
1432
        segment_ids = []
1433 1434
        cnt = 0
        for r in res:
1435 1436 1437 1438
            log.info(f"segmentID {r.segmentID}: state: {r.state}; num_rows: {r.num_rows} ")
            if r.segmentID not in segment_ids:
                segment_ids.append(r.segmentID)
                cnt += r.num_rows
1439 1440
        assert cnt == nb

1441
    @pytest.mark.tags(CaseLabel.L2)
1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454
    def test_get_growing_segment_info_after_load(self):
        """
        target: test get growing segment info
        method: 1.create and load collection
                2.insert data and no flush
                3.get the growing segment
        expected: Verify growing segment num entities
        """
        from pymilvus.grpc_gen.common_pb2 import SegmentState
        collection_w = self.init_collection_wrap(cf.gen_unique_str(prefix))

        collection_w.load()
        collection_w.insert(cf.gen_default_dataframe_data())
1455
        collection_w.search(cf.gen_vectors(1, ct.default_dim), default_field_name, ct.default_search_params, ct.default_limit)
1456 1457 1458 1459 1460 1461 1462
        seg_info = self.utility_wrap.get_query_segment_info(collection_w.name)[0]
        num_entities = 0
        for seg in seg_info:
            assert seg.state == SegmentState.Growing
            num_entities += seg.num_rows
        assert num_entities == ct.default_nb

1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480
    @pytest.mark.tags(CaseLabel.L1)
    def test_get_sealed_query_segment_info(self):
        """
        target: test getting sealed query segment info of collection with data
        method: init a collection, insert data, flush, load, and get query segment info
        expected:
            1. length of segment is greater than 0
            2. the sum num_rows of each segment is equal to num of entities
        """
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        nb = 3000
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
        collection_w.num_entities
        collection_w.load()
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        assert len(res) > 0
1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509
        segment_ids = []
        cnt = 0
        for r in res:
            log.info(f"segmentID {r.segmentID}: state: {r.state}; num_rows: {r.num_rows} ")
            if r.segmentID not in segment_ids:
                segment_ids.append(r.segmentID)
                cnt += r.num_rows
        assert cnt == nb

    @pytest.mark.tags(CaseLabel.L1)
    def test_get_sealed_query_segment_info_after_create_index(self):
        """
        target: test getting sealed query segment info of collection with data
        method: init a collection, insert data, flush, create index, load, and get query segment info
        expected:
            1. length of segment is greater than 0
            2. the sum num_rows of each segment is equal to num of entities
        """
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        nb = 3000
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
        collection_w.num_entities
        collection_w.create_index(default_field_name, default_index_params)
        collection_w.load()
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        assert len(res) > 0
        segment_ids = []
1510 1511
        cnt = 0
        for r in res:
1512 1513 1514 1515
            log.info(f"segmentID {r.segmentID}: state: {r.state}; num_rows: {r.num_rows} ")
            if r.segmentID not in segment_ids:
                segment_ids.append(r.segmentID)
                cnt += r.num_rows
1516
        assert cnt == nb
Z
zhuwenxing 已提交
1517

1518
    @pytest.mark.tags(CaseLabel.L2)
Z
zhuwenxing 已提交
1519 1520 1521 1522 1523 1524 1525
    def test_load_balance_normal(self):
        """
        target: test load balance of collection
        method: init a collection and load balance
        expected: sealed_segment_ids is subset of des_sealed_segment_ids
        """
        # init a collection
1526 1527 1528 1529
        self._connect()
        querynode_num = len(MilvusSys().query_nodes)
        if querynode_num < 2:
            pytest.skip("skip load balance testcase when querynode number less than 2")
Z
zhuwenxing 已提交
1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        ms = MilvusSys()
        nb = 3000
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
        # get sealed segments
        collection_w.num_entities
        # get growing segments
        collection_w.insert(df)
        collection_w.load()
        # prepare load balance params
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
1543
        segment_distribution = cf.get_segment_distribution(res)
1544
        all_querynodes = [node["identifier"] for node in ms.query_nodes]
Z
zhuwenxing 已提交
1545 1546
        assert len(all_querynodes) > 1
        all_querynodes = sorted(all_querynodes,
1547 1548
                                key=lambda x: len(segment_distribution[x]["sealed"])
                                if x in segment_distribution else 0, reverse=True)
Z
zhuwenxing 已提交
1549 1550 1551 1552
        src_node_id = all_querynodes[0]
        des_node_ids = all_querynodes[1:]
        sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
        # load balance
1553
        self.utility_wrap.load_balance(collection_w.name, src_node_id, des_node_ids, sealed_segment_ids)
Z
zhuwenxing 已提交
1554 1555
        # get segments distribution after load balance
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
1556
        segment_distribution = cf.get_segment_distribution(res)
1557 1558 1559
        sealed_segment_ids_after_load_banalce = segment_distribution[src_node_id]["sealed"]
        # assert src node has no sealed segments
        assert sealed_segment_ids_after_load_banalce == []
Z
zhuwenxing 已提交
1560 1561 1562 1563
        des_sealed_segment_ids = []
        for des_node_id in des_node_ids:
            des_sealed_segment_ids += segment_distribution[des_node_id]["sealed"]
        # assert sealed_segment_ids is subset of des_sealed_segment_ids
1564
        assert set(sealed_segment_ids).issubset(des_sealed_segment_ids)
1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593 1594 1595 1596 1597

    @pytest.mark.tags(CaseLabel.L1)
    def test_load_balance_with_src_node_not_exist(self):
        """
        target: test load balance of collection
        method: init a collection and load balance with src_node not exist
        expected: raise exception
        """
        # init a collection
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        ms = MilvusSys()
        nb = 3000
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
        # get sealed segments
        collection_w.num_entities
        # get growing segments
        collection_w.insert(df)
        collection_w.load()
        # prepare load balance params
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        segment_distribution = cf.get_segment_distribution(res)
        all_querynodes = [node["identifier"] for node in ms.query_nodes]
        all_querynodes = sorted(all_querynodes,
                                key=lambda x: len(segment_distribution[x]["sealed"])
                                if x in segment_distribution else 0, reverse=True)
        # set src_node_id as the id of indexnode's id, which is not exist for querynode
        invalid_src_node_id = [node["identifier"] for node in ms.index_nodes][0]
        src_node_id = all_querynodes[0]
        dst_node_ids = all_querynodes[1:]
        sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
        # load balance
1598
        self.utility_wrap.load_balance(collection_w.name, invalid_src_node_id, dst_node_ids, sealed_segment_ids,
1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631 1632
                                       check_task=CheckTasks.err_res,
                                       check_items={ct.err_code: 1, ct.err_msg: "is not exist to balance"})

    @pytest.mark.tags(CaseLabel.L1)
    def test_load_balance_with_all_dst_node_not_exist(self):
        """
        target: test load balance of collection
        method: init a collection and load balance with all dst_node not exist
        expected: raise exception
        """
        # init a collection
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        ms = MilvusSys()
        nb = 3000
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
        # get sealed segments
        collection_w.num_entities
        # get growing segments
        collection_w.insert(df)
        collection_w.load()
        # prepare load balance params
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        segment_distribution = cf.get_segment_distribution(res)
        all_querynodes = [node["identifier"] for node in ms.query_nodes]
        all_querynodes = sorted(all_querynodes,
                                key=lambda x: len(segment_distribution[x]["sealed"])
                                if x in segment_distribution else 0, reverse=True)
        src_node_id = all_querynodes[0]
        # add indexnode's id, which is not exist for querynode, to dst_node_ids
        dst_node_ids = [node["identifier"] for node in ms.index_nodes]
        sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
        # load balance
1633
        self.utility_wrap.load_balance(collection_w.name, src_node_id, dst_node_ids, sealed_segment_ids,
1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666 1667 1668 1669 1670 1671 1672
                                       check_task=CheckTasks.err_res,
                                       check_items={ct.err_code: 1, ct.err_msg: "no available queryNode to allocate"})

    @pytest.mark.tags(CaseLabel.L1)
    def test_load_balance_with_one_sealed_segment_id_not_exist(self):
        """
        target: test load balance of collection
        method: init a collection and load balance with one of sealed segment ids not exist
        expected: raise exception
        """
        # init a collection
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        ms = MilvusSys()
        nb = 3000
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
        # get sealed segments
        collection_w.num_entities
        # get growing segments
        collection_w.insert(df)
        collection_w.load()
        # prepare load balance params
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        segment_distribution = cf.get_segment_distribution(res)
        all_querynodes = [node["identifier"] for node in ms.query_nodes]
        all_querynodes = sorted(all_querynodes,
                                key=lambda x: len(segment_distribution[x]["sealed"])
                                if x in segment_distribution else 0, reverse=True)
        src_node_id = all_querynodes[0]
        dst_node_ids = all_querynodes[1:]
        dst_node_ids.append([node["identifier"] for node in ms.index_nodes][0])
        sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
        # add a segment id which is not exist or a growing segment
        if len(segment_distribution[src_node_id]["growing"]) > 0:
            sealed_segment_ids.append(segment_distribution[src_node_id]["growing"][0])
        else:
            sealed_segment_ids.append(max(segment_distribution[src_node_id]["sealed"]) + 1)
        # load balance
1673
        self.utility_wrap.load_balance(collection_w.name, src_node_id, dst_node_ids, sealed_segment_ids,
1674 1675
                                       check_task=CheckTasks.err_res,
                                       check_items={ct.err_code: 1, ct.err_msg: "is not exist"})
1676

1677
    @pytest.mark.tags(CaseLabel.L2)
1678 1679 1680 1681 1682 1683
    def test_load_balance_in_one_group(self):
        """
        target: test load balance of collection in one group
        method: init a collection, load with multi replicas and load balance among the querynodes in one group
        expected: load balance successfully
        """
1684 1685 1686 1687
        self._connect()
        querynode_num = len(MilvusSys().query_nodes)
        if querynode_num < 3:
            pytest.skip("skip load balance for multi replicas testcase when querynode number less than 3")
1688 1689 1690 1691 1692 1693 1694 1695 1696 1697 1698 1699 1700 1701 1702 1703 1704 1705 1706 1707 1708 1709
        # init a collection
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        nb = 3000
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
        # get sealed segments
        collection_w.num_entities
        collection_w.load(replica_number=2)
        # get growing segments
        collection_w.insert(df)
        # get replicas information
        res, _ = collection_w.get_replicas()
        # prepare load balance params
        # find a group which has multi nodes
        group_nodes = []
        for g in res.groups:
            if len(g.group_nodes) >= 2:
                group_nodes = list(g.group_nodes)
                break
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        segment_distribution = cf.get_segment_distribution(res)
1710 1711 1712 1713 1714 1715
        group_nodes = sorted(group_nodes,
                             key=lambda x: len(
                                 segment_distribution[x]["sealed"])
                             if x in segment_distribution else 0, reverse=True)
        src_node_id = group_nodes[0]
        dst_node_ids = group_nodes[1:]
1716 1717 1718 1719 1720 1721 1722
        sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
        # load balance
        self.utility_wrap.load_balance(collection_w.name, src_node_id, dst_node_ids, sealed_segment_ids)
        # get segments distribution after load balance
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        segment_distribution = cf.get_segment_distribution(res)
        sealed_segment_ids_after_load_banalce = segment_distribution[src_node_id]["sealed"]
1723
        # assert src node has no sealed segments
1724 1725 1726 1727 1728 1729 1730 1731 1732 1733 1734 1735 1736 1737 1738 1739 1740 1741 1742 1743 1744 1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758 1759 1760 1761 1762 1763 1764 1765 1766 1767 1768
        assert sealed_segment_ids_after_load_banalce == []
        des_sealed_segment_ids = []
        for des_node_id in dst_node_ids:
            des_sealed_segment_ids += segment_distribution[des_node_id]["sealed"]
        # assert sealed_segment_ids is subset of des_sealed_segment_ids
        assert set(sealed_segment_ids).issubset(des_sealed_segment_ids)

    @pytest.mark.tags(CaseLabel.L3)
    def test_load_balance_not_in_one_group(self):
        """
        target: test load balance of collection in one group
        method: init a collection, load with multi replicas and load balance among the querynodes in different group
        expected: load balance failed
        """
        # init a collection
        c_name = cf.gen_unique_str(prefix)
        collection_w = self.init_collection_wrap(name=c_name)
        ms = MilvusSys()
        nb = 3000
        df = cf.gen_default_dataframe_data(nb)
        collection_w.insert(df)
        # get sealed segments
        collection_w.num_entities
        collection_w.load(replica_number=2)
        # get growing segments
        collection_w.insert(df)
        # get replicas information
        res, _ = collection_w.get_replicas()
        # prepare load balance params
        all_querynodes = [node["identifier"] for node in ms.query_nodes]
        # find a group which has multi nodes
        group_nodes = []
        for g in res.groups:
            if len(g.group_nodes) >= 2:
                group_nodes = list(g.group_nodes)
                break
        src_node_id = group_nodes[0]              
        dst_node_ids = list(set(all_querynodes) - set(group_nodes))
        res, _ = self.utility_wrap.get_query_segment_info(c_name)
        segment_distribution = cf.get_segment_distribution(res)
        sealed_segment_ids = segment_distribution[src_node_id]["sealed"]
        # load balance
        self.utility_wrap.load_balance(collection_w.name, src_node_id, dst_node_ids, sealed_segment_ids,
                                       check_task=CheckTasks.err_res,
                                       check_items={ct.err_code: 1, ct.err_msg: "must be in the same replica group"})