test_index.py 32.7 KB
Newer Older
J
JinHai-CN 已提交
1 2 3 4 5 6
import logging
import time
import pdb
import threading
from multiprocessing import Pool, Process
import numpy
7
import pytest
Z
zhenwu 已提交
8
import sklearn.preprocessing
J
JinHai-CN 已提交
9 10
from utils import *

G
groot 已提交
11
nb = 6000
J
JinHai-CN 已提交
12
dim = 128
G
groot 已提交
13
index_file_size = 10
Z
zhenwu 已提交
14
BUILD_TIMEOUT = 300
J
JinHai-CN 已提交
15
nprobe = 1
16
top_k = 5
Z
zhenwu 已提交
17
tag = "1970-01-01"
18 19
NLIST = 4046
INVALID_NLIST = 100000000
20 21 22 23 24 25 26 27
field_name = "float_vector"
binary_field_name = "binary_vector"
collection_id = "index"
default_index_type = "FLAT"
entity = gen_entities(1)
entities = gen_entities(nb)
raw_vector, binary_entity = gen_binary_entities(1)
raw_vectors, binary_entities = gen_binary_entities(nb)
28 29
query, query_vecs = gen_query_vectors(field_name, entities, top_k, 1)
default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 1024}, "metric_type": "L2"}
J
JinHai-CN 已提交
30 31 32 33 34


class TestIndexBase:
    @pytest.fixture(
        scope="function",
35
        params=gen_simple_index()
J
JinHai-CN 已提交
36
    )
37 38 39 40
    def get_simple_index(self, request, connect):
        logging.getLogger().info(request.param)
        if str(connect._cmd("mode")) == "CPU":
            if request.param["index_type"] in index_cpu_not_support():
41
                pytest.skip("sq8h not support in CPU mode")
42
        return request.param
J
JinHai-CN 已提交
43 44 45

    @pytest.fixture(
        scope="function",
46 47 48 49 50
        params=[
            1,
            10,
            1500
        ],
J
JinHai-CN 已提交
51
    )
52 53
    def get_nq(self, request):
        yield request.param
54

J
JinHai-CN 已提交
55 56 57 58 59 60 61
    """
    ******************************************************************
      The following cases are used to test `create_index` function
    ******************************************************************
    """

    @pytest.mark.timeout(BUILD_TIMEOUT)
X
Xiaohai Xu 已提交
62
    def test_create_index(self, connect, collection, get_simple_index):
J
JinHai-CN 已提交
63 64
        '''
        target: test create index interface
65 66
        method: create collection and add entities in it, create index
        expected: return search success
J
JinHai-CN 已提交
67
        '''
68
        ids = connect.insert(collection, entities)
69
        connect.create_index(collection, field_name, get_simple_index)
J
JinHai-CN 已提交
70

Z
zhenwu 已提交
71
    @pytest.mark.timeout(BUILD_TIMEOUT)
X
Xiaohai Xu 已提交
72
    def test_create_index_no_vectors(self, connect, collection, get_simple_index):
Z
zhenwu 已提交
73 74
        '''
        target: test create index interface
75 76
        method: create collection and add entities in it, create index
        expected: return search success
Z
zhenwu 已提交
77
        '''
78
        connect.create_index(collection, field_name, get_simple_index)
Z
zhenwu 已提交
79

Z
zhenwu 已提交
80
    @pytest.mark.timeout(BUILD_TIMEOUT)
X
Xiaohai Xu 已提交
81
    def test_create_index_partition(self, connect, collection, get_simple_index):
Z
zhenwu 已提交
82 83
        '''
        target: test create index interface
84 85
        method: create collection, create partition, and add entities in it, create index
        expected: return search success
Z
zhenwu 已提交
86
        '''
87 88 89
        connect.create_partition(collection, tag)
        ids = connect.insert(collection, entities, partition_tag=tag)
        connect.flush([collection])
90
        connect.create_index(collection, field_name, get_simple_index)
91 92

    @pytest.mark.timeout(BUILD_TIMEOUT)
X
Xiaohai Xu 已提交
93
    def test_create_index_partition_flush(self, connect, collection, get_simple_index):
94 95
        '''
        target: test create index interface
96 97
        method: create collection, create partition, and add entities in it, create index
        expected: return search success
98
        '''
99 100
        connect.create_partition(collection, tag)
        ids = connect.insert(collection, entities, partition_tag=tag)
101
        connect.flush()
102
        connect.create_index(collection, field_name, get_simple_index)
103 104 105 106 107 108 109 110

    def test_create_index_without_connect(self, dis_connect, collection):
        '''
        target: test create index without connection
        method: create collection and add entities in it, check if added successfully
        expected: raise exception
        '''
        with pytest.raises(Exception) as e:
111
            dis_connect.create_index(collection, field_name, get_simple_index)
J
JinHai-CN 已提交
112 113

    @pytest.mark.timeout(BUILD_TIMEOUT)
114
    def test_create_index_search_with_query_vectors(self, connect, collection, get_simple_index, get_nq):
J
JinHai-CN 已提交
115 116
        '''
        target: test create index interface, search with more query vectors
117 118
        method: create collection and add entities in it, create index
        expected: return search success
J
JinHai-CN 已提交
119
        '''
120
        ids = connect.insert(collection, entities)
121
        connect.create_index(collection, field_name, get_simple_index)
122 123
        logging.getLogger().info(connect.get_collection_stats(collection))
        nq = get_nq
124 125
        index_type = get_simple_index["index_type"]
        search_param = get_search_param(index_type)
126
        query, vecs = gen_query_vectors(field_name, entities, top_k, nq, search_params=search_param)
127 128
        res = connect.search(collection, query)
        assert len(res) == nq
J
JinHai-CN 已提交
129

130 131
    @pytest.mark.timeout(BUILD_TIMEOUT)
    @pytest.mark.level(2)
X
Xiaohai Xu 已提交
132
    def test_create_index_multithread(self, connect, collection, args):
133 134
        '''
        target: test create index interface with multiprocess
135 136
        method: create collection and add entities in it, create index
        expected: return search success
137
        '''
138
        ids = connect.insert(collection, entities)
139 140

        def build(connect):
141
            connect.create_index(collection, field_name, default_index)
142 143 144 145

        threads_num = 8
        threads = []
        for i in range(threads_num):
D
del-zhenwu 已提交
146
            m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"])
147 148 149 150 151 152 153
            t = threading.Thread(target=build, args=(m,))
            threads.append(t)
            t.start()
            time.sleep(0.2)
        for t in threads:
            t.join()

154 155 156 157 158
    def test_create_index_collection_not_existed(self, connect):
        '''
        target: test create index interface when collection name not existed
        method: create collection and add entities in it, create index
            , make sure the collection name not in index
159
        expected: create index failed
160 161 162
        '''
        collection_name = gen_unique_str(collection_id)
        with pytest.raises(Exception) as e:
163
            connect.create_index(collection_name, field_name, default_index)
164

165
    @pytest.mark.level(2)
166
    @pytest.mark.timeout(BUILD_TIMEOUT)
167
    def test_create_index_insert_flush(self, connect, collection, get_simple_index):
168
        '''
169 170 171
        target: test create index
        method: create collection and create index, add entities in it
        expected: create index ok, and count correct
172
        '''
173
        connect.create_index(collection, field_name, get_simple_index)
174 175 176 177
        ids = connect.insert(collection, entities)
        connect.flush([collection])
        count = connect.count_entities(collection)
        assert count == nb
178

D
del-zhenwu 已提交
179
    @pytest.mark.level(2)
180
    @pytest.mark.timeout(BUILD_TIMEOUT)
181
    def test_create_same_index_repeatedly(self, connect, collection, get_simple_index):
182
        '''
183 184 185
        target: check if index can be created repeatedly, with the same create_index params
        method: create index after index have been built
        expected: return code success, and search ok
186
        '''
187 188
        connect.create_index(collection, field_name, get_simple_index)
        connect.create_index(collection, field_name, get_simple_index)
189

190
    # TODO:
191 192 193 194 195 196 197 198 199
    @pytest.mark.level(2)
    @pytest.mark.timeout(BUILD_TIMEOUT)
    def test_create_different_index_repeatedly(self, connect, collection):
        '''
        target: check if index can be created repeatedly, with the different create_index params
        method: create another index with different index_params after index have been built
        expected: return code 0, and describe index result equals with the second index params
        '''
        ids = connect.insert(collection, entities)
200
        indexs = [default_index, {"metric_type":"L2", "index_type": "FLAT", "params":{"nlist": 1024}}]
201
        for index in indexs:
202
            connect.create_index(collection, field_name, index)
203
            stats = connect.get_collection_stats(collection)
204
            # assert stats["partitions"][0]["segments"][0]["index_name"] == index["index_type"]
205
            assert stats["row_count"] == nb
206 207

    @pytest.mark.timeout(BUILD_TIMEOUT)
208
    def test_create_index_ip(self, connect, collection, get_simple_index):
209 210 211 212 213
        '''
        target: test create index interface
        method: create collection and add entities in it, create index
        expected: return search success
        '''
214
        ids = connect.insert(collection, entities)
215
        get_simple_index["metric_type"] = "IP"
216
        connect.create_index(collection, field_name, get_simple_index)
217

218
    @pytest.mark.timeout(BUILD_TIMEOUT)
219
    def test_create_index_no_vectors_ip(self, connect, collection, get_simple_index):
220 221 222 223 224
        '''
        target: test create index interface
        method: create collection and add entities in it, create index
        expected: return search success
        '''
225
        get_simple_index["metric_type"] = "IP"
226
        connect.create_index(collection, field_name, get_simple_index)
227

228
    @pytest.mark.timeout(BUILD_TIMEOUT)
229
    def test_create_index_partition_ip(self, connect, collection, get_simple_index):
230 231 232 233 234
        '''
        target: test create index interface
        method: create collection, create partition, and add entities in it, create index
        expected: return search success
        '''
235 236 237
        connect.create_partition(collection, tag)
        ids = connect.insert(collection, entities, partition_tag=tag)
        connect.flush([collection])
238
        get_simple_index["metric_type"] = "IP"
239
        connect.create_index(collection, field_name, get_simple_index)
240

J
JinHai-CN 已提交
241
    @pytest.mark.timeout(BUILD_TIMEOUT)
242
    def test_create_index_partition_flush_ip(self, connect, collection, get_simple_index):
243 244 245 246 247
        '''
        target: test create index interface
        method: create collection, create partition, and add entities in it, create index
        expected: return search success
        '''
248 249
        connect.create_partition(collection, tag)
        ids = connect.insert(collection, entities, partition_tag=tag)
250
        connect.flush()
251
        get_simple_index["metric_type"] = "IP"
252
        connect.create_index(collection, field_name, get_simple_index)
J
JinHai-CN 已提交
253

254
    @pytest.mark.timeout(BUILD_TIMEOUT)
255
    def test_create_index_search_with_query_vectors_ip(self, connect, collection, get_simple_index, get_nq):
256 257 258 259 260
        '''
        target: test create index interface, search with more query vectors
        method: create collection and add entities in it, create index
        expected: return search success
        '''
261 262 263 264 265
        metric_type = "IP"
        ids = connect.insert(collection, entities)
        get_simple_index["metric_type"] = metric_type
        connect.create_index(collection, field_name, get_simple_index)
        logging.getLogger().info(connect.get_collection_stats(collection))
266 267 268
        nq = get_nq
        index_type = get_simple_index["index_type"]
        search_param = get_search_param(index_type)
269 270
        query, vecs = gen_query_vectors(field_name, entities, top_k, nq, metric_type=metric_type, search_params=search_param)
        res = connect.search(collection, query)
271
        assert len(res) == nq
J
JinHai-CN 已提交
272 273

    @pytest.mark.timeout(BUILD_TIMEOUT)
274
    @pytest.mark.level(2)
275
    def test_create_index_multithread_ip(self, connect, collection, args):
J
JinHai-CN 已提交
276 277
        '''
        target: test create index interface with multiprocess
278 279
        method: create collection and add entities in it, create index
        expected: return search success
J
JinHai-CN 已提交
280
        '''
281
        ids = connect.insert(collection, entities)
J
JinHai-CN 已提交
282

283
        def build(connect):
284
            default_index["metric_type"] = "IP"
285
            connect.create_index(collection, field_name, default_index)
J
JinHai-CN 已提交
286

287 288 289
        threads_num = 8
        threads = []
        for i in range(threads_num):
D
del-zhenwu 已提交
290
            m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"])
291 292 293
            t = threading.Thread(target=build, args=(m,))
            threads.append(t)
            t.start()
J
JinHai-CN 已提交
294
            time.sleep(0.2)
295 296
        for t in threads:
            t.join()
J
JinHai-CN 已提交
297

298
    def test_create_index_collection_not_existed_ip(self, connect, collection):
J
JinHai-CN 已提交
299
        '''
X
Xiaohai Xu 已提交
300
        target: test create index interface when collection name not existed
301
        method: create collection and add entities in it, create index
X
Xiaohai Xu 已提交
302
            , make sure the collection name not in index
J
JinHai-CN 已提交
303 304
        expected: return code not equals to 0, create index failed
        '''
305
        collection_name = gen_unique_str(collection_id)
306
        default_index["metric_type"] = "IP"
J
JinHai-CN 已提交
307
        with pytest.raises(Exception) as e:
308
            connect.create_index(collection_name, field_name, default_index)
J
JinHai-CN 已提交
309 310

    @pytest.mark.timeout(BUILD_TIMEOUT)
311
    def test_create_index_no_vectors_insert_ip(self, connect, collection, get_simple_index):
J
JinHai-CN 已提交
312
        '''
X
Xiaohai Xu 已提交
313
        target: test create index interface when there is no vectors in collection, and does not affect the subsequent process
314
        method: create collection and add no vectors in it, and then create index, add entities in it
J
JinHai-CN 已提交
315 316
        expected: return code equals to 0
        '''
317
        default_index["metric_type"] = "IP"
318 319 320 321
        connect.create_index(collection, field_name, get_simple_index)
        ids = connect.insert(collection, entities)
        connect.flush([collection])
        count = connect.count_entities(collection)
322
        assert count == nb
J
JinHai-CN 已提交
323

D
del-zhenwu 已提交
324
    @pytest.mark.level(2)
J
JinHai-CN 已提交
325
    @pytest.mark.timeout(BUILD_TIMEOUT)
326
    def test_create_same_index_repeatedly_ip(self, connect, collection, get_simple_index):
J
JinHai-CN 已提交
327 328 329 330 331
        '''
        target: check if index can be created repeatedly, with the same create_index params
        method: create index after index have been built
        expected: return code success, and search ok
        '''
332
        default_index["metric_type"] = "IP"
333 334
        connect.create_index(collection, field_name, get_simple_index)
        connect.create_index(collection, field_name, get_simple_index)
J
JinHai-CN 已提交
335

336
    # TODO:
D
del-zhenwu 已提交
337
    @pytest.mark.level(2)
J
JinHai-CN 已提交
338
    @pytest.mark.timeout(BUILD_TIMEOUT)
339
    def test_create_different_index_repeatedly_ip(self, connect, collection):
J
JinHai-CN 已提交
340 341 342 343 344
        '''
        target: check if index can be created repeatedly, with the different create_index params
        method: create another index with different index_params after index have been built
        expected: return code 0, and describe index result equals with the second index params
        '''
345 346
        ids = connect.insert(collection, entities)
        indexs = [default_index, {"index_type": "FLAT", "params": {"nlist": 1024}, "metric_type": "IP"}]
347
        for index in indexs:
348 349
            connect.create_index(collection, field_name, index)
            stats = connect.get_collection_stats(collection)
350
            # assert stats["partitions"][0]["segments"][0]["index_name"] == index["index_type"]
351
            assert stats["row_count"] == nb
J
JinHai-CN 已提交
352 353 354

    """
    ******************************************************************
355
      The following cases are used to test `drop_index` function
J
JinHai-CN 已提交
356 357 358
    ******************************************************************
    """

359
    def test_drop_index(self, connect, collection, get_simple_index):
J
JinHai-CN 已提交
360
        '''
361 362 363
        target: test drop index interface
        method: create collection and add entities in it, create index, call drop index
        expected: return code 0, and default index param
J
JinHai-CN 已提交
364
        '''
365
        # ids = connect.insert(collection, entities)
366 367
        connect.create_index(collection, field_name, get_simple_index)
        connect.drop_index(collection, field_name)
368 369 370
        stats = connect.get_collection_stats(collection)
        # assert stats["partitions"][0]["segments"][0]["index_name"] == default_index_type
        assert not stats["partitions"][0]["segments"]
J
JinHai-CN 已提交
371

372 373
    @pytest.mark.level(2)
    def test_drop_index_repeatly(self, connect, collection, get_simple_index):
J
JinHai-CN 已提交
374
        '''
375 376 377
        target: test drop index repeatly
        method: create index, call drop index, and drop again
        expected: return code 0
J
JinHai-CN 已提交
378
        '''
379
        connect.create_index(collection, field_name, get_simple_index)
380
        stats = connect.get_collection_stats(collection)
381 382
        connect.drop_index(collection, field_name)
        connect.drop_index(collection, field_name)
383 384 385 386
        stats = connect.get_collection_stats(collection)
        logging.getLogger().info(stats)
        # assert stats["partitions"][0]["segments"][0]["index_name"] == default_index_type
        assert not stats["partitions"][0]["segments"]
J
JinHai-CN 已提交
387

388 389
    @pytest.mark.level(2)
    def test_drop_index_without_connect(self, dis_connect, collection):
J
JinHai-CN 已提交
390
        '''
391 392 393
        target: test drop index without connection
        method: drop index, and check if drop successfully
        expected: raise exception
J
JinHai-CN 已提交
394
        '''
395
        with pytest.raises(Exception) as e:
396
            dis_connect.drop_index(collection, field_name)
J
JinHai-CN 已提交
397

398
    def test_drop_index_collection_not_existed(self, connect):
J
JinHai-CN 已提交
399
        '''
400 401 402 403
        target: test drop index interface when collection name not existed
        method: create collection and add entities in it, create index
            , make sure the collection name not in index, and then drop it
        expected: return code not equals to 0, drop index failed
J
JinHai-CN 已提交
404
        '''
405
        collection_name = gen_unique_str(collection_id)
J
JinHai-CN 已提交
406
        with pytest.raises(Exception) as e:
407
            connect.drop_index(collection_name, field_name)
J
JinHai-CN 已提交
408

409
    def test_drop_index_collection_not_create(self, connect, collection):
J
JinHai-CN 已提交
410
        '''
411 412 413
        target: test drop index interface when index not created
        method: create collection and add entities in it, create index
        expected: return code not equals to 0, drop index failed
J
JinHai-CN 已提交
414
        '''
415 416
        # ids = connect.insert(collection, entities)
        # no create index
417
        connect.drop_index(collection, field_name)
J
JinHai-CN 已提交
418

419 420 421 422 423 424 425 426
    @pytest.mark.level(2)
    def test_create_drop_index_repeatly(self, connect, collection, get_simple_index):
        '''
        target: test create / drop index repeatly, use the same index params
        method: create index, drop index, four times
        expected: return code 0
        '''
        for i in range(4):
427 428
            connect.create_index(collection, field_name, get_simple_index)
            connect.drop_index(collection, field_name)
J
JinHai-CN 已提交
429

430
    def test_drop_index_ip(self, connect, collection, get_simple_index):
J
JinHai-CN 已提交
431 432
        '''
        target: test drop index interface
433
        method: create collection and add entities in it, create index, call drop index
J
JinHai-CN 已提交
434 435
        expected: return code 0, and default index param
        '''
436
        # ids = connect.insert(collection, entities)
437
        get_simple_index["metric_type"] = "IP"
438 439 440
        connect.create_index(collection, field_name, get_simple_index)
        connect.drop_index(collection, field_name)
        stats = connect.get_collection_stats(collection)
441 442
        # assert stats["partitions"][0]["segments"][0]["index_name"] == default_index_type
        assert not stats["partitions"][0]["segments"]
J
JinHai-CN 已提交
443

D
del-zhenwu 已提交
444
    @pytest.mark.level(2)
445
    def test_drop_index_repeatly_ip(self, connect, collection, get_simple_index):
J
JinHai-CN 已提交
446 447 448 449 450
        '''
        target: test drop index repeatly
        method: create index, call drop index, and drop again
        expected: return code 0
        '''
451
        get_simple_index["metric_type"] = "IP"
452 453 454 455 456
        connect.create_index(collection, field_name, get_simple_index)
        stats = connect.get_collection_stats(collection)
        connect.drop_index(collection, field_name)
        connect.drop_index(collection, field_name)
        stats = connect.get_collection_stats(collection)
457 458 459
        logging.getLogger().info(stats)
        # assert stats["partitions"][0]["segments"][0]["index_name"] == default_index_type
        assert not stats["partitions"][0]["segments"]
J
JinHai-CN 已提交
460

461
    @pytest.mark.level(2)
462
    def test_drop_index_without_connect_ip(self, dis_connect, collection):
J
JinHai-CN 已提交
463
        '''
464 465 466
        target: test drop index without connection
        method: drop index, and check if drop successfully
        expected: raise exception
J
JinHai-CN 已提交
467 468
        '''
        with pytest.raises(Exception) as e:
469
            dis_connect.drop_index(collection, field_name)
J
JinHai-CN 已提交
470

471
    def test_drop_index_collection_not_create_ip(self, connect, collection):
J
JinHai-CN 已提交
472 473
        '''
        target: test drop index interface when index not created
474
        method: create collection and add entities in it, create index
J
JinHai-CN 已提交
475 476
        expected: return code not equals to 0, drop index failed
        '''
477
        # ids = connect.insert(collection, entities)
J
JinHai-CN 已提交
478
        # no create index
479
        connect.drop_index(collection, field_name)
J
JinHai-CN 已提交
480

D
del-zhenwu 已提交
481
    @pytest.mark.level(2)
482
    def test_create_drop_index_repeatly_ip(self, connect, collection, get_simple_index):
J
JinHai-CN 已提交
483 484 485 486 487
        '''
        target: test create / drop index repeatly, use the same index params
        method: create index, drop index, four times
        expected: return code 0
        '''
488
        get_simple_index["metric_type"] = "IP"
489
        for i in range(4):
490 491
            connect.create_index(collection, field_name, get_simple_index)
            connect.drop_index(collection, field_name)
J
JinHai-CN 已提交
492 493


D
del-zhenwu 已提交
494
class TestIndexBinary:
J
JinHai-CN 已提交
495 496
    @pytest.fixture(
        scope="function",
497
        params=gen_simple_index()
J
JinHai-CN 已提交
498
    )
499 500 501
    def get_simple_index(self, request, connect):
        if str(connect._cmd("mode")) == "CPU":
            if request.param["index_type"] in index_cpu_not_support():
502
                pytest.skip("sq8h not support in CPU mode")
503
        return request.param
J
JinHai-CN 已提交
504 505 506

    @pytest.fixture(
        scope="function",
507
        params=gen_binary_index()
J
JinHai-CN 已提交
508
    )
509
    def get_jaccard_index(self, request, connect):
510 511 512 513 514
        if request.param["index_type"] in binary_support():
            request.param["metric_type"] = "JACCARD"
            return request.param
        else:
            pytest.skip("Skip index")
515

516 517 518 519 520 521 522 523
    @pytest.fixture(
        scope="function",
        params=gen_binary_index()
    )
    def get_l2_index(self, request, connect):
        request.param["metric_type"] = "L2"
        return request.param

524 525 526 527 528 529 530 531 532 533 534
    @pytest.fixture(
        scope="function",
        params=[
            1,
            10,
            1500
        ],
    )
    def get_nq(self, request):
        yield request.param

J
JinHai-CN 已提交
535 536 537 538 539
    """
    ******************************************************************
      The following cases are used to test `create_index` function
    ******************************************************************
    """
540

J
JinHai-CN 已提交
541
    @pytest.mark.timeout(BUILD_TIMEOUT)
D
del-zhenwu 已提交
542
    def test_create_index(self, connect, binary_collection, get_jaccard_index):
J
JinHai-CN 已提交
543 544
        '''
        target: test create index interface
545 546
        method: create collection and add entities in it, create index
        expected: return search success
J
JinHai-CN 已提交
547
        '''
D
del-zhenwu 已提交
548 549
        ids = connect.insert(binary_collection, binary_entities)
        connect.create_index(binary_collection, binary_field_name, get_jaccard_index)
J
JinHai-CN 已提交
550

Z
zhenwu 已提交
551
    @pytest.mark.timeout(BUILD_TIMEOUT)
D
del-zhenwu 已提交
552
    def test_create_index_partition(self, connect, binary_collection, get_jaccard_index):
Z
zhenwu 已提交
553 554
        '''
        target: test create index interface
555 556
        method: create collection, create partition, and add entities in it, create index
        expected: return search success
Z
zhenwu 已提交
557
        '''
D
del-zhenwu 已提交
558 559 560
        connect.create_partition(binary_collection, tag)
        ids = connect.insert(binary_collection, binary_entities, partition_tag=tag)
        connect.create_index(binary_collection, binary_field_name, get_jaccard_index)
J
JinHai-CN 已提交
561 562

    @pytest.mark.timeout(BUILD_TIMEOUT)
D
del-zhenwu 已提交
563
    def test_create_index_search_with_query_vectors(self, connect, binary_collection, get_jaccard_index, get_nq):
J
JinHai-CN 已提交
564 565
        '''
        target: test create index interface, search with more query vectors
566 567
        method: create collection and add entities in it, create index
        expected: return search success
J
JinHai-CN 已提交
568
        '''
569
        nq = get_nq
D
del-zhenwu 已提交
570 571
        ids = connect.insert(binary_collection, binary_entities)
        connect.create_index(binary_collection, binary_field_name, get_jaccard_index)
572
        query, vecs = gen_query_vectors(binary_field_name, binary_entities, top_k, nq, metric_type="JACCARD")
D
del-zhenwu 已提交
573
        search_param = get_search_param(get_jaccard_index["index_type"], metric_type="JACCARD")
574
        logging.getLogger().info(search_param)
D
del-zhenwu 已提交
575
        res = connect.search(binary_collection, query, search_params=search_param)
576
        assert len(res) == nq
J
JinHai-CN 已提交
577

578 579 580 581 582 583 584 585 586 587 588 589 590 591
    @pytest.mark.timeout(BUILD_TIMEOUT)
    def test_create_index_invalid_metric_type_binary(self, connect, binary_collection, get_l2_index):
        '''
        target: test create index interface with invalid metric type
        method: add entitys into binary connection, flash, create index with L2 metric type.
        expected: return create_index failure
        '''
        # insert 6000 vectors
        ids = connect.insert(binary_collection, binary_entities)
        connect.flush([binary_collection])

        with pytest.raises(Exception) as e:
            res = connect.create_index(binary_collection, binary_field_name, get_l2_index)

592 593 594 595 596 597
    """
    ******************************************************************
      The following cases are used to test `get_index_info` function
    ******************************************************************
    """

D
del-zhenwu 已提交
598
    def test_get_index_info(self, connect, binary_collection, get_jaccard_index):
J
JinHai-CN 已提交
599
        '''
600 601 602
        target: test describe index interface
        method: create collection and add entities in it, create index, call describe index
        expected: return code 0, and index instructure
J
JinHai-CN 已提交
603
        '''
D
del-zhenwu 已提交
604 605 606 607
        ids = connect.insert(binary_collection, binary_entities)
        connect.flush([binary_collection])
        connect.create_index(binary_collection, binary_field_name, get_jaccard_index)
        stats = connect.get_collection_stats(binary_collection)
D
del-zhenwu 已提交
608 609 610 611 612 613 614 615
        assert stats["row_count"] == nb
        for partition in stats["partitions"]:
            segments = partition["segments"]
            if segments:
                for segment in segments:
                    for file in segment["files"]:
                        if "index_type" in file:
                            assert file["index_type"] == get_jaccard_index["index_type"]
J
JinHai-CN 已提交
616

D
del-zhenwu 已提交
617
    def test_get_index_info_partition(self, connect, binary_collection, get_jaccard_index):
618 619 620 621 622
        '''
        target: test describe index interface
        method: create collection, create partition and add entities in it, create index, call describe index
        expected: return code 0, and index instructure
        '''
D
del-zhenwu 已提交
623 624 625 626 627
        connect.create_partition(binary_collection, tag)
        ids = connect.insert(binary_collection, binary_entities, partition_tag=tag)
        connect.flush([binary_collection])
        connect.create_index(binary_collection, binary_field_name, get_jaccard_index)
        stats = connect.get_collection_stats(binary_collection)
628
        logging.getLogger().info(stats)
D
del-zhenwu 已提交
629 630 631 632 633 634 635 636 637
        assert stats["row_count"] == nb
        assert len(stats["partitions"]) == 2
        for partition in stats["partitions"]:
            segments = partition["segments"]
            if segments:
                for segment in segments:
                    for file in segment["files"]:
                        if "index_type" in file:
                            assert file["index_type"] == get_jaccard_index["index_type"]
J
JinHai-CN 已提交
638

639 640 641 642 643 644
    """
    ******************************************************************
      The following cases are used to test `drop_index` function
    ******************************************************************
    """

645
    def test_drop_index(self, connect, binary_collection, get_jaccard_index):
646 647 648 649 650
        '''
        target: test drop index interface
        method: create collection and add entities in it, create index, call drop index
        expected: return code 0, and default index param
        '''
651 652
        connect.create_index(binary_collection, binary_field_name, get_jaccard_index)
        stats = connect.get_collection_stats(binary_collection)
653
        logging.getLogger().info(stats)
654 655
        connect.drop_index(binary_collection, binary_field_name)
        stats = connect.get_collection_stats(binary_collection)
656 657 658
        # assert stats["partitions"][0]["segments"][0]["index_name"] == default_index_type
        assert not stats["partitions"][0]["segments"]

659
    def test_drop_index_partition(self, connect, binary_collection, get_jaccard_index):
660 661 662 663 664
        '''
        target: test drop index interface
        method: create collection, create partition and add entities in it, create index on collection, call drop collection index
        expected: return code 0, and default index param
        '''
665 666 667 668 669 670 671
        connect.create_partition(binary_collection, tag)
        ids = connect.insert(binary_collection, binary_entities, partition_tag=tag)
        connect.flush([binary_collection])
        connect.create_index(binary_collection, binary_field_name, get_jaccard_index)
        stats = connect.get_collection_stats(binary_collection)
        connect.drop_index(binary_collection, binary_field_name)
        stats = connect.get_collection_stats(binary_collection)
D
del-zhenwu 已提交
672 673 674 675 676 677 678 679 680 681
        assert stats["row_count"] == nb
        for partition in stats["partitions"]:
            segments = partition["segments"]
            if segments:
                for segment in segments:
                    for file in segment["files"]:
                        if "index_type" not in file:
                            continue
                        if file["index_type"] == get_jaccard_index["index_type"]:
                            assert False
J
JinHai-CN 已提交
682

D
del-zhenwu 已提交
683

684
class TestIndexInvalid(object):
J
JinHai-CN 已提交
685
    """
X
Xiaohai Xu 已提交
686
    Test create / describe / drop index interfaces with invalid collection names
J
JinHai-CN 已提交
687
    """
688

J
JinHai-CN 已提交
689 690
    @pytest.fixture(
        scope="function",
691
        params=gen_invalid_strs()
J
JinHai-CN 已提交
692
    )
X
Xiaohai Xu 已提交
693
    def get_collection_name(self, request):
J
JinHai-CN 已提交
694 695
        yield request.param

Z
zhenwu 已提交
696
    @pytest.mark.level(1)
X
Xiaohai Xu 已提交
697 698
    def test_create_index_with_invalid_collectionname(self, connect, get_collection_name):
        collection_name = get_collection_name
699
        with pytest.raises(Exception) as e:
700
            connect.create_index(collection_name, field_name, default_index)
J
JinHai-CN 已提交
701

Z
zhenwu 已提交
702
    @pytest.mark.level(1)
X
Xiaohai Xu 已提交
703 704
    def test_drop_index_with_invalid_collectionname(self, connect, get_collection_name):
        collection_name = get_collection_name
705 706
        with pytest.raises(Exception) as e:
            connect.drop_index(collection_name)
J
JinHai-CN 已提交
707 708 709

    @pytest.fixture(
        scope="function",
710
        params=gen_invalid_index()
J
JinHai-CN 已提交
711
    )
712
    def get_index(self, request):
J
JinHai-CN 已提交
713 714
        yield request.param

Z
zhenwu 已提交
715
    @pytest.mark.level(1)
X
Xiaohai Xu 已提交
716
    def test_create_index_with_invalid_index_params(self, connect, collection, get_index):
717
        logging.getLogger().info(get_index)
718
        with pytest.raises(Exception) as e:
719
            connect.create_index(collection, field_name, get_simple_index)
720

D
del-zhenwu 已提交
721 722

class TestIndexAsync:
Z
zw 已提交
723 724 725 726 727
    @pytest.fixture(scope="function", autouse=True)
    def skip_http_check(self, args):
        if args["handler"] == "HTTP":
            pytest.skip("skip in http mode")

D
del-zhenwu 已提交
728 729 730 731 732 733 734 735 736 737 738
    """
    ******************************************************************
      The following cases are used to test `create_index` function
    ******************************************************************
    """

    @pytest.fixture(
        scope="function",
        params=gen_simple_index()
    )
    def get_simple_index(self, request, connect):
739 740
        if str(connect._cmd("mode")) == "CPU":
            if request.param["index_type"] in index_cpu_not_support():
D
del-zhenwu 已提交
741 742 743
                pytest.skip("sq8h not support in CPU mode")
        return request.param

744 745 746
    def check_result(self, res):
        logging.getLogger().info("In callback check search result")
        logging.getLogger().info(res)
D
del-zhenwu 已提交
747 748 749 750 751 752 753 754 755 756 757

    """
    ******************************************************************
      The following cases are used to test `create_index` function
    ******************************************************************
    """

    @pytest.mark.timeout(BUILD_TIMEOUT)
    def test_create_index(self, connect, collection, get_simple_index):
        '''
        target: test create index interface
758 759
        method: create collection and add entities in it, create index
        expected: return search success
D
del-zhenwu 已提交
760
        '''
761
        ids = connect.insert(collection, entities)
D
del-zhenwu 已提交
762
        logging.getLogger().info("start index")
763
        future = connect.create_index(collection, field_name, get_simple_index, _async=True)
D
del-zhenwu 已提交
764
        logging.getLogger().info("before result")
765 766 767
        res = future.result()
        # TODO:
        logging.getLogger().info(res)
D
del-zhenwu 已提交
768 769 770

    def test_create_index_with_invalid_collectionname(self, connect):
        collection_name = " "
771
        future = connect.create_index(collection_name, field_name, default_index, _async=True)
772 773
        with pytest.raises(Exception) as e:
            res = future.result()
D
del-zhenwu 已提交
774

775 776 777 778 779 780 781 782 783
    @pytest.mark.timeout(BUILD_TIMEOUT)
    def test_create_index_callback(self, connect, collection, get_simple_index):
        '''
        target: test create index interface
        method: create collection and add entities in it, create index
        expected: return search success
        '''
        ids = connect.insert(collection, entities)
        logging.getLogger().info("start index")
784
        future = connect.create_index(collection, field_name, get_simple_index, _async=True,
785 786 787 788 789
                                      _callback=self.check_result)
        logging.getLogger().info("before result")
        res = future.result()
        # TODO:
        logging.getLogger().info(res)