test_index   F
last analyzed

Complexity

Total Complexity 194

Size/Duplication

Total Lines 1887
Duplicated Lines 40.43 %

Importance

Changes 0
Metric Value
eloc 1180
dl 763
loc 1887
rs 0.88
c 0
b 0
f 0
wmc 194

93 Methods

Rating   Name   Duplication   Size   Complexity  
A TestIndexJAC.test_drop_index() 21 21 1
A TestIndexBase.test_create_index_partition() 14 14 1
A TestIndexBase.test_create_index() 0 13 1
A TestIndexBase.test_create_index_partition_flush() 15 15 1
A TestIndexBase.test_create_index_no_vectors() 0 12 1
A TestIndexBase.get_index() 0 12 5
A TestIndexBase.test_create_index_search_with_query_vectors() 20 20 1
A TestIndexBase.get_simple_index() 12 12 5
A TestIndexBase.test_describe_and_drop_index_multi_collections() 38 38 4
A TestIndexBase.test_drop_index_collection_not_create() 0 13 1
A TestIndexBase.test_create_same_index_repeatedly() 0 13 1
A TestIndexBase.test_create_index_a_multithreads() 0 26 4
A TestIndexBase.test_drop_index_collection_not_existed() 0 10 1
A TestIndexBase.test_create_drop_index_repeatly_different_index_params() 0 20 2
A TestIndexIP.test_create_index_collection() 0 14 1
B TestIndexBase._test_create_index_multiprocessing_multicollection() 49 49 5
A TestIndexBase.test_create_index_collection_None() 0 12 2
A TestIndexIP.test_create_index() 0 14 1
A TestIndexBase.test_create_different_index_repeatedly() 21 21 2
B TestIndexBase.test_create_index_multithread_multicollection() 47 47 5
A TestIndexBase._test_create_index_multiprocessing() 32 32 3
A TestIndexBase.test_drop_index() 0 19 1
A TestIndexBase.test_drop_index_collection_None() 0 9 2
A TestIndexBase.test_create_index_collection_not_existed() 0 13 1
A TestIndexBase.test_create_index_multithread() 32 32 3
B TestIndexIP.get_simple_index() 0 14 6
B TestIndexIP.get_index() 0 14 6
A TestIndexBase.test_get_index_info_collection_None() 0 9 2
A TestIndexBase.test_get_index_info_not_create() 0 11 1
A TestIndexBase.test_get_index_info_collection_not_existed() 0 10 1
A TestIndexBase.test_get_index_info() 0 17 2
A TestIndexBase.test_create_index_no_vectors_then_insert() 0 12 1
A TestIndexBase.test_drop_index_repeatly() 22 22 1
A TestIndexBase.test_create_drop_index_repeatly() 21 21 2
A TestIndexBinary.test_create_index_partition() 16 16 1
A TestIndexIP.test_create_different_index_repeatedly() 20 20 2
A TestIndexIP.test_create_drop_index_repeatly() 21 21 2
A TestIndexIP.test_drop_index_repeatly() 28 28 3
A TestIndexIP.test_drop_index() 24 24 3
A TestIndexAsync.test_create_index() 0 18 1
A TestIndexAsync.check_status() 0 3 1
A TestIndexJAC.get_simple_index() 12 12 5
B TestIndexIP._test_create_index_multiprocessing_multicollection() 0 47 5
A TestIndexIP.test_create_index_search_with_query_vectors() 0 22 1
A TestIndexAsync.skip_http_check() 0 4 2
A TestIndexJAC.test_get_index_info_partition() 17 17 1
A TestIndexIP.test_drop_index_collection_not_create() 0 13 1
A TestIndexIP.test_create_same_index_repeatedly() 0 20 1
A TestIndexJAC.test_get_index_info() 0 16 1
A TestIndexCollectionInvalid.get_collection_name() 0 6 1
A TestIndexJAC.test_create_index_partition() 0 14 1
A TestCreateIndexParamsInvalid.test_create_index_with_empty_param() 0 9 2
A TestIndexBinary.get_hamming_index() 10 10 3
A TestIndexIP._test_create_index_multiprocessing() 32 32 3
A TestIndexBinary.test_get_index_info_partition() 17 17 1
A TestIndexIP.test_drop_index_partition_C() 0 20 1
A TestIndexBinary.get_substructure_index() 0 10 2
A TestIndexIP.test_describe_and_drop_index_multi_collections() 36 36 4
A TestIndexBinary.test_create_index_partition_structure() 16 16 1
A TestCreateIndexParamsInvalid.get_index() 0 6 1
A TestIndexJAC.test_drop_index_partition() 0 20 1
A TestIndexIP.test_create_drop_index_repeatly_different_index_params() 0 24 2
A TestCreateIndexParamsInvalid.test_create_index_with_invalid_index_params() 0 12 4
A TestIndexBinary.test_get_index_info() 0 16 1
A TestIndexIP.test_drop_index_partition() 0 20 1
A TestCreateIndexParamsInvalid.test_create_index_with_invalid_nlist() 0 5 2
A TestIndexCollectionInvalid.test_create_index_with_invalid_collectionname() 0 7 1
A TestIndexIP.test_get_index_info_not_create() 0 11 1
A TestIndexBinary.test_create_index_search_with_query_vectors() 20 20 1
A TestIndexIP.test_get_index_info_partition_A() 0 20 1
A TestIndexBinary.test_create_index_search_with_query_vectors_superstructure() 20 20 1
A TestIndexBinary.test_create_index() 0 16 3
A TestIndexIP.test_create_index_no_vectors_then_insert() 0 12 1
A TestIndexIP.test_get_index_info() 0 21 3
A TestIndexJAC.get_jaccard_index() 10 10 3
A TestIndexBinary.test_drop_index() 21 21 1
A TestIndexCollectionInvalid.test_get_index_info_with_invalid_collectionname() 0 5 1
A TestCreateIndexParamsInvalid.get_index_type() 0 6 1
A TestIndexBinary.test_drop_index_substructure() 20 20 1
A TestIndexBinary.test_get_index_info_partition_superstructrue() 17 17 1
A TestIndexIP.test_create_index_no_vectors() 0 11 1
A TestIndexCollectionInvalid.test_drop_index_with_invalid_collectionname() 0 5 1
A TestIndexJAC.test_create_index() 0 16 3
A TestIndexBinary.test_drop_index_partition() 0 20 1
A TestIndexBinary.get_superstructure_index() 0 10 2
A TestIndexIP.test_get_index_info_partition() 17 17 1
A TestIndexJAC.test_create_index_search_with_query_vectors() 20 20 1
A TestIndexBinary.get_simple_index() 11 11 5
A TestIndexAsync.test_create_index_with_invalid_collectionname() 0 7 1
A TestIndexBinary.get_index() 0 11 5
A TestIndexAsync.get_simple_index() 14 14 5
A TestIndexJAC.get_index() 0 12 5
A TestIndexAsync.get_index() 0 12 5

How to fix   Duplicated Code    Complexity   

Duplicated Code

Duplicate code is one of the most pungent code smells. A rule that is often used is to re-structure code once it is duplicated in three or more places.

Common duplication problems, and corresponding solutions are:

Complexity

 Tip:   Before tackling complexity, make sure that you eliminate any duplication first. This often can reduce the size of classes significantly.

Complex classes like test_index often do a lot of different things. To break such a class down, we need to identify a cohesive component within that class. A common approach to find such a component is to look for fields/methods that share the same prefixes, or suffixes.

Once you have determined the fields that belong together, you can apply the Extract Class refactoring. If the component makes sense as a sub-class, Extract Subclass is also a candidate, and is often faster.

1
"""
2
   For testing index operations, including `create_index`, `get_index_info` and `drop_index` interfaces
3
"""
4
import logging
5
import pytest
6
import time
7
import pdb
8
import threading
9
from multiprocessing import Pool, Process
10
import numpy
11
import sklearn.preprocessing
12
from milvus import IndexType, MetricType
13
from utils import *
14
15
nb = 6000
16
dim = 128
17
index_file_size = 10
18
vectors = gen_vectors(nb, dim)
19
vectors = sklearn.preprocessing.normalize(vectors, axis=1, norm='l2')
20
vectors = vectors.tolist()
21
BUILD_TIMEOUT = 300
22
nprobe = 1
23
tag = "1970-01-01"
24
NLIST = 4046
25
INVALID_NLIST = 100000000
26
27
28
class TestIndexBase:
29
    @pytest.fixture(
30
        scope="function",
31
        params=gen_index()
32
    )
33
    def get_index(self, request, connect):
34
        if str(connect._cmd("mode")[1]) == "CPU":
35
            if request.param["index_type"] == IndexType.IVF_SQ8H:
36
                pytest.skip("sq8h not support in CPU mode")
37
        if str(connect._cmd("mode")[1]) == "GPU":
38
            if request.param["index_type"] == IndexType.IVF_PQ:
39
                pytest.skip("ivfpq not support in GPU mode")
40
        return request.param
41
42 View Code Duplication
    @pytest.fixture(
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
43
        scope="function",
44
        params=gen_simple_index()
45
    )
46
    def get_simple_index(self, request, connect):
47
        if str(connect._cmd("mode")[1]) == "CPU":
48
            if request.param["index_type"] == IndexType.IVF_SQ8H:
49
                pytest.skip("sq8h not support in CPU mode")
50
        if str(connect._cmd("mode")[1]) == "GPU":
51
            if request.param["index_type"] == IndexType.IVF_PQ:
52
                pytest.skip("ivfpq not support in GPU mode")
53
        return request.param
54
55
    """
56
    ******************************************************************
57
      The following cases are used to test `create_index` function
58
    ******************************************************************
59
    """
60
61
    @pytest.mark.timeout(BUILD_TIMEOUT)
62
    def test_create_index(self, connect, collection, get_simple_index):
63
        '''
64
        target: test create index interface
65
        method: create collection and add vectors in it, create index
66
        expected: return code equals to 0, and search success
67
        '''
68
        index_param = get_simple_index["index_param"]
69
        index_type = get_simple_index["index_type"]
70
        logging.getLogger().info(get_simple_index)
71
        status, ids = connect.insert(collection, vectors)
72
        status = connect.create_index(collection, index_type, index_param)
73
        assert status.OK()
74
75
    @pytest.mark.timeout(BUILD_TIMEOUT)
76
    def test_create_index_no_vectors(self, connect, collection, get_simple_index):
77
        '''
78
        target: test create index interface
79
        method: create collection and add vectors in it, create index
80
        expected: return code equals to 0, and search success
81
        '''
82
        index_param = get_simple_index["index_param"]
83
        index_type = get_simple_index["index_type"]
84
        logging.getLogger().info(get_simple_index)
85
        status = connect.create_index(collection, index_type, index_param)
86
        assert status.OK()
87
88 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
89
    def test_create_index_partition(self, connect, collection, get_simple_index):
90
        '''
91
        target: test create index interface
92
        method: create collection, create partition, and add vectors in it, create index
93
        expected: return code equals to 0, and search success
94
        '''
95
        index_param = get_simple_index["index_param"]
96
        index_type = get_simple_index["index_type"]
97
        logging.getLogger().info(get_simple_index)
98
        status = connect.create_partition(collection, tag)
99
        status, ids = connect.insert(collection, vectors, partition_tag=tag)
100
        status = connect.create_index(collection, index_type, index_param)
101
        assert status.OK()
102
103 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
104
    def test_create_index_partition_flush(self, connect, collection, get_simple_index):
105
        '''
106
        target: test create index interface
107
        method: create collection, create partition, and add vectors in it, create index
108
        expected: return code equals to 0, and search success
109
        '''
110
        index_param = get_simple_index["index_param"]
111
        index_type = get_simple_index["index_type"]
112
        logging.getLogger().info(get_simple_index)
113
        status = connect.create_partition(collection, tag)
114
        status, ids = connect.insert(collection, vectors, partition_tag=tag)
115
        connect.flush()
116
        status = connect.create_index(collection, index_type, index_param)
117
        assert status.OK()
118
119
    # @pytest.mark.level(2)
120
    # def test_create_index_without_connect(self, dis_connect, collection):
121
    #     '''
122
    #     target: test create index without connection
123
    #     method: create collection and add vectors in it, check if added successfully
124
    #     expected: raise exception
125
    #     '''
126
    #     nlist = NLIST
127
    #     index_type = IndexType.IVF_SQ8
128
    #     index_param = {"nlist": nlist}
129
    #     with pytest.raises(Exception) as e:
130
    #         status = dis_connect.create_index(collection, index_type, index_param)
131
132 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
133
    def test_create_index_search_with_query_vectors(self, connect, collection, get_simple_index):
134
        '''
135
        target: test create index interface, search with more query vectors
136
        method: create collection and add vectors in it, create index
137
        expected: return code equals to 0, and search success
138
        '''
139
        index_param = get_simple_index["index_param"]
140
        index_type = get_simple_index["index_type"]
141
        logging.getLogger().info(get_simple_index)
142
        status, ids = connect.insert(collection, vectors)
143
        status = connect.create_index(collection, index_type, index_param)
144
        logging.getLogger().info(connect.get_index_info(collection))
145
        query_vecs = [vectors[0], vectors[1], vectors[2]]
146
        top_k = 5
147
        search_param = get_search_param(index_type)
148
        status, result = connect.search(collection, top_k, query_vecs, params=search_param)
149
        assert status.OK()
150
        assert len(result) == len(query_vecs)
151
        logging.getLogger().info(result)
152
153 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
154
    @pytest.mark.level(2)
155
    def test_create_index_multithread(self, connect, collection, args):
156
        '''
157
        target: test create index interface with multiprocess
158
        method: create collection and add vectors in it, create index
159
        expected: return code equals to 0, and search success
160
        '''
161
        status, ids = connect.insert(collection, vectors)
162
163
        def build(connect):
164
            status = connect.create_index(collection, IndexType.IVFLAT, {"nlist": NLIST})
165
            assert status.OK()
166
167
        threads_num = 8
168
        threads = []
169
        for i in range(threads_num):
170
            m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"])
171
            t = threading.Thread(target=build, args=(m,))
172
            threads.append(t)
173
            t.start()
174
            time.sleep(0.2)
175
        for t in threads:
176
            t.join()
177
178
        query_vec = [vectors[0]]
179
        top_k = 1
180
        search_param = {"nprobe": nprobe}
181
        status, result = connect.search(collection, top_k, query_vec, params=search_param)
182
        assert len(result) == 1
183
        assert len(result[0]) == top_k
184
        assert result[0][0].distance == 0.0
185
186 View Code Duplication
    @pytest.mark.level(2)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
187
    @pytest.mark.timeout(BUILD_TIMEOUT)
188
    def test_create_index_multithread_multicollection(self, connect, args):
189
        '''
190
        target: test create index interface with multiprocess
191
        method: create collection and add vectors in it, create index
192
        expected: return code equals to 0, and search success
193
        '''
194
        threads_num = 8
195
        loop_num = 8
196
        threads = []
197
        collection = []
198
        j = 0
199
        while j < (threads_num*loop_num):
200
            collection_name = gen_unique_str("test_create_index_multiprocessing")
201
            collection.append(collection_name)
202
            param = {'collection_name': collection_name,
203
                     'dimension': dim,
204
                     'index_type': IndexType.FLAT,
205
                     'store_raw_vector': False}
206
            connect.create_collection(param)
207
            j = j + 1
208
209
        def create_index():
210
            i = 0
211
            while i < loop_num:
212
                # assert connect.has_collection(collection[ids*process_num+i])
213
                status, ids = connect.insert(collection[ids*threads_num+i], vectors)
0 ignored issues
show
introduced by
The variable ids does not seem to be defined in case the while loop on line 211 is not entered. Are you sure this can never be the case?
Loading history...
214
                status = connect.create_index(collection[ids*threads_num+i], IndexType.IVFLAT, {"nlist": NLIST})
215
                assert status.OK()
216
                query_vec = [vectors[0]]
217
                top_k = 1
218
                search_param = {"nprobe": nprobe}
219
                status, result = connect.search(collection[ids*threads_num+i], top_k, query_vec, params=search_param)
220
                assert len(result) == 1
221
                assert len(result[0]) == top_k
222
                assert result[0][0].distance == 0.0
223
                i = i + 1
224
        for i in range(threads_num):
225
            m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"])
226
            ids = i
227
            t = threading.Thread(target=create_index, args=(m, ids))
228
            threads.append(t)
229
            t.start()
230
            time.sleep(0.2)
231
        for t in threads:
232
            t.join()
233
234
    @pytest.mark.timeout(BUILD_TIMEOUT)
235
    @pytest.mark.level(2)
236
    def test_create_index_a_multithreads(self, connect, collection, args):
237
        status, ids = connect.insert(collection, vectors)
238
        def build(connect):
239
            status = connect.create_index(collection, IndexType.IVFLAT, {"nlist": NLIST})
240
            assert status.OK()
241
        def count(connect):
242
            status, count = connect.count_entities(collection)
243
            assert status.OK()
244
            assert count == nb
245
246
        threads_num = 8
247
        threads = []
248
        uri = "tcp://%s:%s" % (args["ip"], args["port"])
249
        for i in range(threads_num):
250
            m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"])
251
            if(i % 2 == 0):
252
                p = threading.Thread(target=build, args=(m,))
253
            else:
254
                p = threading.Thread(target=count, args=(m,))
255
            threads.append(p)
256
            p.start()
257
            time.sleep(0.2)
258
        for p in threads:
259
            p.join()
260
261
262
    # TODO: enable
263 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
264
    @pytest.mark.level(2)
265
    def _test_create_index_multiprocessing(self, connect, collection, args):
266
        '''
267
        target: test create index interface with multiprocess
268
        method: create collection and add vectors in it, create index
269
        expected: return code equals to 0, and search success
270
        '''
271
        status, ids = connect.insert(collection, vectors)
272
273
        def build(connect):
274
            status = connect.create_index(collection, IndexType.IVFLAT, {"nlist": NLIST})
275
            assert status.OK()
276
277
        process_num = 8
278
        processes = []
279
        for i in range(process_num):
280
            m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"])
281
            p = Process(target=build, args=(m,))
282
            processes.append(p)
283
            p.start()
284
            time.sleep(0.2)
285
        for p in processes:
286
            p.join()
287
288
        query_vec = [vectors[0]]
289
        top_k = 1
290
        search_param = {"nprobe": nprobe}
291
        status, result = connect.search(collection, top_k, query_vec, params=search_param)
292
        assert len(result) == 1
293
        assert len(result[0]) == top_k
294
        assert result[0][0].distance == 0.0
295
296
    # TODO: enable
297 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
298
    def _test_create_index_multiprocessing_multicollection(self, connect, args):
299
        '''
300
        target: test create index interface with multiprocess
301
        method: create collection and add vectors in it, create index
302
        expected: return code equals to 0, and search success
303
        '''
304
        process_num = 8
305
        loop_num = 8
306
        processes = []
307
308
        collection = []
309
        j = 0
310
        while j < (process_num*loop_num):
311
            collection_name = gen_unique_str("test_create_index_multiprocessing")
312
            collection.append(collection_name)
313
            param = {'collection_name': collection_name,
314
                    'dimension': dim,
315
                    'index_type': IndexType.FLAT,
316
                    'store_raw_vector': False}
317
            connect.create_collection(param)
318
            j = j + 1
319
320
        def create_index():
321
            i = 0
322
            while i < loop_num:
323
                # assert connect.has_collection(collection[ids*process_num+i])
324
                status, ids = connect.insert(collection[ids*process_num+i], vectors)
0 ignored issues
show
introduced by
The variable ids does not seem to be defined in case the while loop on line 322 is not entered. Are you sure this can never be the case?
Loading history...
325
326
                status = connect.create_index(collection[ids*process_num+i], IndexType.IVFLAT, {"nlist": NLIST})
327
                assert status.OK()
328
                query_vec = [vectors[0]]
329
                top_k = 1
330
                search_param = {"nprobe": nprobe}
331
                status, result = connect.search(collection[ids*process_num+i], top_k, query_vec, params=search_param)
332
                assert len(result) == 1
333
                assert len(result[0]) == top_k
334
                assert result[0][0].distance == 0.0
335
                i = i + 1
336
337
        for i in range(process_num):
338
            m = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"])
339
            ids = i
340
            p = Process(target=create_index, args=(m,ids))
341
            processes.append(p)
342
            p.start()
343
            time.sleep(0.2)
344
        for p in processes:
345
            p.join()
346
347
    def test_create_index_collection_not_existed(self, connect):
348
        '''
349
        target: test create index interface when collection name not existed
350
        method: create collection and add vectors in it, create index
351
            , make sure the collection name not in index
352
        expected: return code not equals to 0, create index failed
353
        '''
354
        collection_name = gen_unique_str(self.__class__.__name__)
355
        nlist = NLIST
356
        index_type = IndexType.IVF_SQ8
357
        index_param = {"nlist": nlist}
358
        status = connect.create_index(collection_name, index_type, index_param)
359
        assert not status.OK()
360
361
    def test_create_index_collection_None(self, connect):
362
        '''
363
        target: test create index interface when collection name is None
364
        method: create collection and add vectors in it, create index with an collection_name: None
365
        expected: return code not equals to 0, create index failed
366
        '''
367
        collection_name = None
368
        nlist = NLIST
369
        index_type = IndexType.IVF_SQ8
370
        index_param = {"nlist": nlist}
371
        with pytest.raises(Exception) as e:
372
            status = connect.create_index(collection_name, index_type, index_param)
373
374
    @pytest.mark.timeout(BUILD_TIMEOUT)
375
    def test_create_index_no_vectors_then_insert(self, connect, collection, get_simple_index):
376
        '''
377
        target: test create index interface when there is no vectors in collection, and does not affect the subsequent process
378
        method: create collection and add no vectors in it, and then create index, add vectors in it
379
        expected: return code equals to 0
380
        '''
381
        index_param = get_simple_index["index_param"]
382
        index_type = get_simple_index["index_type"]
383
        status = connect.create_index(collection, index_type, index_param)
384
        status, ids = connect.insert(collection, vectors)
385
        assert status.OK()
386
387
    @pytest.mark.level(2)
388
    @pytest.mark.timeout(BUILD_TIMEOUT)
389
    def test_create_same_index_repeatedly(self, connect, collection, get_simple_index):
390
        '''
391
        target: check if index can be created repeatedly, with the same create_index params
392
        method: create index after index have been built
393
        expected: return code success, and search ok
394
        '''
395
        index_param = get_simple_index["index_param"]
396
        index_type = get_simple_index["index_type"]
397
        status = connect.create_index(collection, index_type, index_param)
398
        status = connect.create_index(collection, index_type, index_param)
399
        assert status.OK()
400
401 View Code Duplication
    @pytest.mark.level(2)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
402
    @pytest.mark.timeout(BUILD_TIMEOUT)
403
    def test_create_different_index_repeatedly(self, connect, collection):
404
        '''
405
        target: check if index can be created repeatedly, with the different create_index params
406
        method: create another index with different index_params after index have been built
407
        expected: return code 0, and describe index result equals with the second index params
408
        '''
409
        nlist = NLIST
410
        status, ids = connect.insert(collection, vectors)
411
        index_type_1 = IndexType.IVF_SQ8
412
        index_type_2 = IndexType.IVFLAT
413
        indexs = [{"index_type": index_type_1, "index_param": {"nlist": nlist}}, {"index_type": index_type_2, "index_param": {"nlist": nlist}}]
414
        logging.getLogger().info(indexs)
415
        for index in indexs:
416
            status = connect.create_index(collection, index["index_type"], index["index_param"])
417
            assert status.OK()
418
        status, result = connect.get_index_info(collection)
419
        assert result._params["nlist"] == nlist
420
        assert result._collection_name == collection
421
        assert result._index_type == index_type_2
422
423
    """
424
    ******************************************************************
425
      The following cases are used to test `get_index_info` function
426
    ******************************************************************
427
    """
428
429
    def test_get_index_info(self, connect, collection, get_index):
430
        '''
431
        target: test describe index interface
432
        method: create collection and add vectors in it, create index, call describe index
433
        expected: return code 0, and index instructure
434
        '''
435
        index_param = get_index["index_param"]
436
        index_type = get_index["index_type"]
437
        logging.getLogger().info(get_index)
438
        # status, ids = connect.insert(collection, vectors)
439
        status = connect.create_index(collection, index_type, index_param)
440
        if status.OK():
441
            status, result = connect.get_index_info(collection)
442
            logging.getLogger().info(result)
443
            assert result._params == index_param
444
            assert result._collection_name == collection
445
            assert result._index_type == index_type
446
447 View Code Duplication
    def test_describe_and_drop_index_multi_collections(self, connect, get_simple_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
448
        '''
449
        target: test create, describe and drop index interface with multiple collections of L2
450
        method: create collections and add vectors in it, create index, call describe index
451
        expected: return code 0, and index instructure
452
        '''
453
        nq = 100
454
        vectors = gen_vectors(nq, dim)
455
        collection_list = []
456
        for i in range(10):
457
            collection_name = gen_unique_str()
458
            collection_list.append(collection_name)
459
            param = {'collection_name': collection_name,
460
                     'dimension': dim,
461
                     'index_file_size': index_file_size,
462
                     'metric_type': MetricType.L2}
463
            connect.create_collection(param)
464
            index_param = get_simple_index["index_param"]
465
            index_type = get_simple_index["index_type"]
466
            logging.getLogger().info(get_simple_index)
467
            status, ids = connect.insert(collection_name=collection_name, records=vectors)
468
            status = connect.create_index(collection_name, index_type, index_param)
469
            assert status.OK()
470
471
        for i in range(10):
472
            status, result = connect.get_index_info(collection_list[i])
473
            logging.getLogger().info(result)
474
            assert result._params == index_param
0 ignored issues
show
introduced by
The variable index_param does not seem to be defined in case the for loop on line 456 is not entered. Are you sure this can never be the case?
Loading history...
475
            assert result._collection_name == collection_list[i]
476
            assert result._index_type == index_type
0 ignored issues
show
introduced by
The variable index_type does not seem to be defined in case the for loop on line 456 is not entered. Are you sure this can never be the case?
Loading history...
477
478
        for i in range(10):
479
            status = connect.drop_index(collection_list[i])
480
            assert status.OK()
481
            status, result = connect.get_index_info(collection_list[i])
482
            logging.getLogger().info(result)
483
            assert result._collection_name == collection_list[i]
484
            assert result._index_type == IndexType.FLAT
485
486
    # @pytest.mark.level(2)
487
    # def test_get_index_info_without_connect(self, dis_connect, collection):
488
    #     '''
489
    #     target: test describe index without connection
490
    #     method: describe index, and check if describe successfully
491
    #     expected: raise exception
492
    #     '''
493
    #     with pytest.raises(Exception) as e:
494
    #         status = dis_connect.get_index_info(collection)
495
496
    def test_get_index_info_collection_not_existed(self, connect):
497
        '''
498
        target: test describe index interface when collection name not existed
499
        method: create collection and add vectors in it, create index
500
            , make sure the collection name not in index
501
        expected: return code not equals to 0, describe index failed
502
        '''
503
        collection_name = gen_unique_str(self.__class__.__name__)
504
        status, result = connect.get_index_info(collection_name)
505
        assert not status.OK()
506
507
    def test_get_index_info_collection_None(self, connect):
508
        '''
509
        target: test describe index interface when collection name is None
510
        method: create collection and add vectors in it, create index with an collection_name: None
511
        expected: return code not equals to 0, describe index failed
512
        '''
513
        collection_name = None
514
        with pytest.raises(Exception) as e:
515
            status = connect.get_index_info(collection_name)
516
517
    def test_get_index_info_not_create(self, connect, collection):
518
        '''
519
        target: test describe index interface when index not created
520
        method: create collection and add vectors in it, create index
521
            , make sure the collection name not in index
522
        expected: return code not equals to 0, describe index failed
523
        '''
524
        status, ids = connect.insert(collection, vectors)
525
        status, result = connect.get_index_info(collection)
526
        logging.getLogger().info(result)
527
        assert status.OK()
528
        # assert result._params["nlist"] == index_params["nlist"]
529
        # assert result._collection_name == collection
530
        # assert result._index_type == index_params["index_type"]
531
532
    """
533
    ******************************************************************
534
      The following cases are used to test `drop_index` function
535
    ******************************************************************
536
    """
537
538
    def test_drop_index(self, connect, collection, get_simple_index):
539
        '''
540
        target: test drop index interface
541
        method: create collection and add vectors in it, create index, call drop index
542
        expected: return code 0, and default index param
543
        '''
544
        index_param = get_simple_index["index_param"]
545
        index_type = get_simple_index["index_type"]
546
        # status, ids = connect.insert(collection, vectors)
547
        status = connect.create_index(collection, index_type, index_param)
548
        assert status.OK()
549
        status, result = connect.get_index_info(collection)
550
        logging.getLogger().info(result)
551
        status = connect.drop_index(collection)
552
        assert status.OK()
553
        status, result = connect.get_index_info(collection)
554
        logging.getLogger().info(result)
555
        assert result._collection_name == collection
556
        assert result._index_type == IndexType.FLAT
557
558 View Code Duplication
    @pytest.mark.level(2)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
559
    def test_drop_index_repeatly(self, connect, collection, get_simple_index):
560
        '''
561
        target: test drop index repeatly
562
        method: create index, call drop index, and drop again
563
        expected: return code 0
564
        '''
565
        index_param = get_simple_index["index_param"]
566
        index_type = get_simple_index["index_type"]
567
        # status, ids = connect.insert(collection, vectors)
568
        status = connect.create_index(collection, index_type, index_param)
569
        assert status.OK()
570
        status, result = connect.get_index_info(collection)
571
        logging.getLogger().info(result)
572
        status = connect.drop_index(collection)
573
        assert status.OK()
574
        status = connect.drop_index(collection)
575
        assert status.OK()
576
        status, result = connect.get_index_info(collection)
577
        logging.getLogger().info(result)
578
        assert result._collection_name == collection
579
        assert result._index_type == IndexType.FLAT
580
581
    # @pytest.mark.level(2)
582
    # def test_drop_index_without_connect(self, dis_connect, collection):
583
    #     '''
584
    #     target: test drop index without connection
585
    #     method: drop index, and check if drop successfully
586
    #     expected: raise exception
587
    #     '''
588
    #     with pytest.raises(Exception) as e:
589
    #         status = dis_connect.drop_index(collection)
590
591
    def test_drop_index_collection_not_existed(self, connect):
592
        '''
593
        target: test drop index interface when collection name not existed
594
        method: create collection and add vectors in it, create index
595
            , make sure the collection name not in index, and then drop it
596
        expected: return code not equals to 0, drop index failed
597
        '''
598
        collection_name = gen_unique_str(self.__class__.__name__)
599
        status = connect.drop_index(collection_name)
600
        assert not status.OK()
601
602
    def test_drop_index_collection_None(self, connect):
603
        '''
604
        target: test drop index interface when collection name is None
605
        method: create collection and add vectors in it, create index with an collection_name: None
606
        expected: return code not equals to 0, drop index failed
607
        '''
608
        collection_name = None
609
        with pytest.raises(Exception) as e:
610
            status = connect.drop_index(collection_name)
611
612
    def test_drop_index_collection_not_create(self, connect, collection):
613
        '''
614
        target: test drop index interface when index not created
615
        method: create collection and add vectors in it, create index
616
        expected: return code not equals to 0, drop index failed
617
        '''
618
        status, ids = connect.insert(collection, vectors)
619
        status, result = connect.get_index_info(collection)
620
        logging.getLogger().info(result)
621
        # no create index
622
        status = connect.drop_index(collection)
623
        logging.getLogger().info(status)
624
        assert status.OK()
625
626 View Code Duplication
    @pytest.mark.level(2)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
627
    def test_create_drop_index_repeatly(self, connect, collection, get_simple_index):
628
        '''
629
        target: test create / drop index repeatly, use the same index params
630
        method: create index, drop index, four times
631
        expected: return code 0
632
        '''
633
        index_param = get_simple_index["index_param"]
634
        index_type = get_simple_index["index_type"]
635
        # status, ids = connect.insert(collection, vectors)
636
        for i in range(2):
637
            status = connect.create_index(collection, index_type, index_param)
638
            assert status.OK()
639
            status, result = connect.get_index_info(collection)
640
            logging.getLogger().info(result)
641
            status = connect.drop_index(collection)
642
            assert status.OK()
643
            status, result = connect.get_index_info(collection)
644
            logging.getLogger().info(result)
645
            assert result._collection_name == collection
646
            assert result._index_type == IndexType.FLAT
647
648
    def test_create_drop_index_repeatly_different_index_params(self, connect, collection):
649
        '''
650
        target: test create / drop index repeatly, use the different index params
651
        method: create index, drop index, four times, each tme use different index_params to create index
652
        expected: return code 0
653
        '''
654
        nlist = NLIST
655
        indexs = [{"index_type": IndexType.IVFLAT, "index_param": {"nlist": nlist}}, {"index_type": IndexType.IVF_SQ8, "index_param": {"nlist": nlist}}]
656
        # status, ids = connect.insert(collection, vectors)
657
        for i in range(2):
658
            status = connect.create_index(collection, indexs[i]["index_type"], indexs[i]["index_param"])
659
            assert status.OK()
660
            status, result = connect.get_index_info(collection)
661
            logging.getLogger().info(result)
662
            status = connect.drop_index(collection)
663
            assert status.OK()
664
            status, result = connect.get_index_info(collection)
665
            logging.getLogger().info(result)
666
            assert result._collection_name == collection
667
            assert result._index_type == IndexType.FLAT
668
669
670
class TestIndexIP:
671
    @pytest.fixture(
672
        scope="function",
673
        params=gen_index()
674
    )
675
    def get_index(self, request, connect):
676
        if str(connect._cmd("mode")[1]) == "CPU":
677
            if request.param["index_type"] == IndexType.IVF_SQ8H:
678
                pytest.skip("sq8h not support in CPU mode")
679
        if str(connect._cmd("mode")[1]) == "GPU":
680
            if request.param["index_type"] == IndexType.IVF_PQ:
681
                pytest.skip("ivfpq not support in GPU mode")
682
        if request.param["index_type"] == IndexType.RNSG:
683
            pytest.skip("rnsg not support in ip")
684
        return request.param
685
686
    @pytest.fixture(
687
        scope="function",
688
        params=gen_simple_index()
689
    )
690
    def get_simple_index(self, request, connect):
691
        if str(connect._cmd("mode")[1]) == "CPU":
692
            if request.param["index_type"] == IndexType.IVF_SQ8H:
693
                pytest.skip("sq8h not support in CPU mode")
694
        if str(connect._cmd("mode")[1]) == "GPU":
695
            if request.param["index_type"] == IndexType.IVF_PQ:
696
                pytest.skip("ivfpq not support in GPU mode")
697
        if request.param["index_type"] == IndexType.RNSG:
698
            pytest.skip("rnsg not support in ip")
699
        return request.param
700
    """
701
    ******************************************************************
702
      The following cases are used to test `create_index` function
703
    ******************************************************************
704
    """
705
    @pytest.mark.level(2)
706
    @pytest.mark.timeout(BUILD_TIMEOUT)
707
    def test_create_index(self, connect, ip_collection, get_simple_index):
708
        '''
709
        target: test create index interface
710
        method: create collection and add vectors in it, create index
711
        expected: return code equals to 0, and search success
712
        '''
713
        index_param = get_simple_index["index_param"]
714
        index_type = get_simple_index["index_type"]
715
        logging.getLogger().info(get_simple_index)
716
        status, ids = connect.insert(ip_collection, vectors)
717
        status = connect.create_index(ip_collection, index_type, index_param)
718
        assert status.OK()
719
720
    @pytest.mark.timeout(BUILD_TIMEOUT)
721
    def test_create_index_collection(self, connect, ip_collection, get_simple_index):
722
        '''
723
        target: test create index interface
724
        method: create collection, create partition, and add vectors in it, create index on collection
725
        expected: return code equals to 0, and search success
726
        '''
727
        index_param = get_simple_index["index_param"]
728
        index_type = get_simple_index["index_type"]
729
        logging.getLogger().info(get_simple_index)
730
        status = connect.create_partition(ip_collection, tag)
731
        status, ids = connect.insert(ip_collection, vectors, partition_tag=tag)
732
        status = connect.create_index(ip_collection, index_type, index_param)
733
        assert status.OK()
734
735
    # @pytest.mark.level(2)
736
    # def test_create_index_without_connect(self, dis_connect, ip_collection):
737
    #     '''
738
    #     target: test create index without connection
739
    #     method: create collection and add vectors in it, check if added successfully
740
    #     expected: raise exception
741
    #     '''
742
    #     nlist = NLIST
743
    #     index_type = IndexType.IVF_SQ8
744
    #     index_param = {"nlist": nlist}
745
    #     with pytest.raises(Exception) as e:
746
    #         status = dis_connect.create_index(ip_collection, index_type, index_param)
747
748
    @pytest.mark.timeout(BUILD_TIMEOUT)
749
    def test_create_index_search_with_query_vectors(self, connect, ip_collection, get_simple_index):
750
        '''
751
        target: test create index interface, search with more query vectors
752
        method: create collection and add vectors in it, create index, with no manual flush 
753
        expected: return code equals to 0, and search success
754
        '''
755
        index_param = get_simple_index["index_param"]
756
        index_type = get_simple_index["index_type"]
757
        logging.getLogger().info(get_simple_index)
758
        logging.getLogger().info(connect.get_collection_info(ip_collection))
759
        status, ids = connect.insert(ip_collection, vectors)
760
        status = connect.create_index(ip_collection, index_type, index_param)
761
        logging.getLogger().info(connect.get_index_info(ip_collection))
762
        logging.getLogger().info(connect.get_collection_stats(ip_collection))
763
        query_vecs = [vectors[0], vectors[1], vectors[2]]
764
        top_k = 5
765
        search_param = get_search_param(index_type)
766
        status, result = connect.search(ip_collection, top_k, query_vecs, params=search_param)
767
        logging.getLogger().info(result)
768
        assert status.OK()
769
        assert len(result) == len(query_vecs)
770
771
    # TODO: enable
772 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
773
    @pytest.mark.level(2)
774
    def _test_create_index_multiprocessing(self, connect, ip_collection, args):
775
        '''
776
        target: test create index interface with multiprocess
777
        method: create collection and add vectors in it, create index
778
        expected: return code equals to 0, and search success
779
        '''
780
        status, ids = connect.insert(ip_collection, vectors)
781
        def build(connect):
782
            status = connect.create_index(ip_collection, IndexType.IVFLAT, {"nlist": NLIST})
783
            assert status.OK()
784
785
        process_num = 8
786
        processes = []
787
788
        for i in range(process_num):
789
            m = get_milvus(args["ip"], args["port"], handler=args["handler"])
790
            p = Process(target=build, args=(m,))
791
            processes.append(p)
792
            p.start()
793
            time.sleep(0.2)
794
        for p in processes:
795
            p.join()
796
797
        query_vec = [vectors[0]]
798
        top_k = 1
799
        search_param = {"nprobe": nprobe}
800
        status, result = connect.search(ip_collection, top_k, query_vec, params=search_param)
801
        assert len(result) == 1
802
        assert len(result[0]) == top_k
803
        assert result[0][0].distance == 0.0
804
805
    # TODO: enable
806
    @pytest.mark.timeout(BUILD_TIMEOUT)
807
    def _test_create_index_multiprocessing_multicollection(self, connect, args):
808
        '''
809
        target: test create index interface with multiprocess
810
        method: create collection and add vectors in it, create index
811
        expected: return code equals to 0, and search success
812
        '''
813
        process_num = 8
814
        loop_num = 8
815
        processes = []
816
817
        collection = []
818
        j = 0
819
        while j < (process_num*loop_num):
820
            collection_name = gen_unique_str("test_create_index_multiprocessing")
821
            collection.append(collection_name)
822
            param = {'collection_name': collection_name,
823
                    'dimension': dim}
824
            connect.create_collection(param)
825
            j = j + 1
826
827
        def create_index():
828
            i = 0
829
            while i < loop_num:
830
                # assert connect.has_collection(collection[ids*process_num+i])
831
                status, ids = connect.insert(collection[ids*process_num+i], vectors)
0 ignored issues
show
introduced by
The variable ids does not seem to be defined in case the while loop on line 829 is not entered. Are you sure this can never be the case?
Loading history...
832
833
                status = connect.create_index(collection[ids*process_num+i], IndexType.IVFLAT, {"nlist": NLIST})
834
                assert status.OK()
835
                query_vec = [vectors[0]]
836
                top_k = 1
837
                search_param = {"nprobe": nprobe}
838
                status, result = connect.search(collection[ids*process_num+i], top_k, query_vec, params=search_param)
839
                assert len(result) == 1
840
                assert len(result[0]) == top_k
841
                assert result[0][0].distance == 0.0
842
                i = i + 1
843
844
        for i in range(process_num):
845
            m = get_milvus(args["ip"], args["port"], handler=args["handler"])
846
            ids = i
847
            p = Process(target=create_index, args=(m,ids))
848
            processes.append(p)
849
            p.start()
850
            time.sleep(0.2)
851
        for p in processes:
852
            p.join()
853
854
    def test_create_index_no_vectors(self, connect, ip_collection):
855
        '''
856
        target: test create index interface when there is no vectors in collection
857
        method: create collection and add no vectors in it, and then create index
858
        expected: return code equals to 0
859
        '''
860
        nlist = NLIST
861
        index_type = IndexType.IVF_SQ8
862
        index_param = {"nlist": nlist}
863
        status = connect.create_index(ip_collection, index_type, index_param)
864
        assert status.OK()
865
866
    @pytest.mark.timeout(BUILD_TIMEOUT)
867
    def test_create_index_no_vectors_then_insert(self, connect, ip_collection, get_simple_index):
868
        '''
869
        target: test create index interface when there is no vectors in collection, and does not affect the subsequent process
870
        method: create collection and add no vectors in it, and then create index, add vectors in it
871
        expected: return code equals to 0
872
        '''
873
        index_param = get_simple_index["index_param"]
874
        index_type = get_simple_index["index_type"]
875
        status = connect.create_index(ip_collection, index_type, index_param)
876
        status, ids = connect.insert(ip_collection, vectors)
877
        assert status.OK()
878
879
    @pytest.mark.timeout(BUILD_TIMEOUT)
880
    def test_create_same_index_repeatedly(self, connect, ip_collection):
881
        '''
882
        target: check if index can be created repeatedly, with the same create_index params
883
        method: create index after index have been built
884
        expected: return code success, and search ok
885
        '''
886
        nlist = NLIST
887
        status, ids = connect.insert(ip_collection, vectors)
888
        index_type = IndexType.IVF_SQ8
889
        index_param = {"nlist": nlist}
890
        status = connect.create_index(ip_collection, index_type, index_param)
891
        status = connect.create_index(ip_collection, index_type, index_param)
892
        assert status.OK()
893
        query_vec = [vectors[0]]
894
        top_k = 1
895
        search_param = {"nprobe": nprobe}
896
        status, result = connect.search(ip_collection, top_k, query_vec, params=search_param)
897
        assert len(result) == 1
898
        assert len(result[0]) == top_k
899
900 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
901
    def test_create_different_index_repeatedly(self, connect, ip_collection):
902
        '''
903
        target: check if index can be created repeatedly, with the different create_index params
904
        method: create another index with different index_params after index have been built
905
        expected: return code 0, and describe index result equals with the second index params
906
        '''
907
        nlist = NLIST
908
        status, ids = connect.insert(ip_collection, vectors)
909
        index_type_1 = IndexType.IVF_SQ8
910
        index_type_2 = IndexType.IVFLAT
911
        indexs = [{"index_type": index_type_1, "index_param": {"nlist": nlist}}, {"index_type": index_type_2, "index_param": {"nlist": nlist}}]
912
        logging.getLogger().info(indexs)
913
        for index in indexs:
914
            status = connect.create_index(ip_collection, index["index_type"], index["index_param"])
915
            assert status.OK()
916
        status, result = connect.get_index_info(ip_collection)
917
        assert result._params["nlist"] == nlist
918
        assert result._collection_name == ip_collection
919
        assert result._index_type == index_type_2
920
921
    """
922
    ******************************************************************
923
      The following cases are used to test `get_index_info` function
924
    ******************************************************************
925
    """
926
927
    def test_get_index_info(self, connect, ip_collection, get_simple_index):
928
        '''
929
        target: test describe index interface
930
        method: create collection and add vectors in it, create index, call describe index
931
        expected: return code 0, and index instructure
932
        '''
933
        index_param = get_simple_index["index_param"]
934
        index_type = get_simple_index["index_type"]
935
        logging.getLogger().info(get_simple_index)
936
        # status, ids = connect.insert(ip_collection, vectors[:5000])
937
        status = connect.create_index(ip_collection, index_type, index_param)
938
        status, result = connect.get_index_info(ip_collection)
939
        logging.getLogger().info(result)
940
        assert result._collection_name == ip_collection
941
        status, mode = connect._cmd("mode")
942
        if str(mode) == "GPU" and index_type == IndexType.IVF_PQ:
943
            assert result._index_type == IndexType.FLAT
944
            assert result._params["nlist"] == NLIST
945
        else:
946
            assert result._index_type == index_type
947
            assert result._params == index_param
948
949 View Code Duplication
    def test_get_index_info_partition(self, connect, ip_collection, get_simple_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
950
        '''
951
        target: test describe index interface
952
        method: create collection, create partition and add vectors in it, create index, call describe index
953
        expected: return code 0, and index instructure
954
        '''
955
        index_param = get_simple_index["index_param"]
956
        index_type = get_simple_index["index_type"]
957
        logging.getLogger().info(get_simple_index)
958
        status = connect.create_partition(ip_collection, tag)
959
        status, ids = connect.insert(ip_collection, vectors, partition_tag=tag)
960
        status = connect.create_index(ip_collection, index_type, index_param)
961
        status, result = connect.get_index_info(ip_collection)
962
        logging.getLogger().info(result)
963
        assert result._params == index_param
964
        assert result._collection_name == ip_collection
965
        assert result._index_type == index_type
966
967
    def test_get_index_info_partition_A(self, connect, ip_collection, get_simple_index):
968
        '''
969
        target: test describe index interface
970
        method: create collection, create partitions and add vectors in it, create index on partitions, call describe index
971
        expected: return code 0, and index instructure
972
        '''
973
        new_tag = "new_tag"
974
        index_param = get_simple_index["index_param"]
975
        index_type = get_simple_index["index_type"]
976
        logging.getLogger().info(get_simple_index)
977
        status = connect.create_partition(ip_collection, tag)
978
        status = connect.create_partition(ip_collection, new_tag)
979
        # status, ids = connect.insert(ip_collection, vectors, partition_tag=tag)
980
        # status, ids = connect.insert(ip_collection, vectors, partition_tag=new_tag)
981
        status = connect.create_index(ip_collection, index_type, index_param)
982
        status, result = connect.get_index_info(ip_collection)
983
        logging.getLogger().info(result)
984
        assert result._params == index_param
985
        assert result._collection_name == ip_collection
986
        assert result._index_type == index_type
987
988 View Code Duplication
    def test_describe_and_drop_index_multi_collections(self, connect, get_simple_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
989
        '''
990
        target: test create, describe and drop index interface with multiple collections of IP
991
        method: create collections and add vectors in it, create index, call describe index
992
        expected: return code 0, and index instructure
993
        '''
994
        nq = 100
995
        vectors = gen_vectors(nq, dim)
996
        collection_list = []
997
        for i in range(10):
998
            collection_name = gen_unique_str()
999
            collection_list.append(collection_name)
1000
            param = {'collection_name': collection_name,
1001
                     'dimension': dim,
1002
                     'index_file_size': index_file_size,
1003
                     'metric_type': MetricType.IP}
1004
            connect.create_collection(param)
1005
            index_param = get_simple_index["index_param"]
1006
            index_type = get_simple_index["index_type"]
1007
            logging.getLogger().info(get_simple_index)
1008
            status, ids = connect.insert(collection_name=collection_name, records=vectors)
1009
            status = connect.create_index(collection_name, index_type, index_param)
1010
            assert status.OK()
1011
        for i in range(10):
1012
            status, result = connect.get_index_info(collection_list[i])
1013
            logging.getLogger().info(result)
1014
            assert result._params == index_param
0 ignored issues
show
introduced by
The variable index_param does not seem to be defined in case the for loop on line 997 is not entered. Are you sure this can never be the case?
Loading history...
1015
            assert result._collection_name == collection_list[i]
1016
            assert result._index_type == index_type
0 ignored issues
show
introduced by
The variable index_type does not seem to be defined in case the for loop on line 997 is not entered. Are you sure this can never be the case?
Loading history...
1017
        for i in range(10):
1018
            status = connect.drop_index(collection_list[i])
1019
            assert status.OK()
1020
            status, result = connect.get_index_info(collection_list[i])
1021
            logging.getLogger().info(result)
1022
            assert result._collection_name == collection_list[i]
1023
            assert result._index_type == IndexType.FLAT
1024
1025
    # @pytest.mark.level(2)
1026
    # def test_get_index_info_without_connect(self, dis_connect, ip_collection):
1027
    #     '''
1028
    #     target: test describe index without connection
1029
    #     method: describe index, and check if describe successfully
1030
    #     expected: raise exception
1031
    #     '''
1032
    #     with pytest.raises(Exception) as e:
1033
    #         status = dis_connect.get_index_info(ip_collection)
1034
1035
    def test_get_index_info_not_create(self, connect, ip_collection):
1036
        '''
1037
        target: test describe index interface when index not created
1038
        method: create collection and add vectors in it, create index
1039
            , make sure the collection name not in index
1040
        expected: return code not equals to 0, describe index failed
1041
        '''
1042
        status, ids = connect.insert(ip_collection, vectors)
1043
        status, result = connect.get_index_info(ip_collection)
1044
        logging.getLogger().info(result)
1045
        assert status.OK()
1046
        # assert result._params["nlist"] == index_params["nlist"]
1047
        # assert result._collection_name == collection
1048
        # assert result._index_type == index_params["index_type"]
1049
1050
    """
1051
    ******************************************************************
1052
      The following cases are used to test `drop_index` function
1053
    ******************************************************************
1054
    """
1055
1056 View Code Duplication
    def test_drop_index(self, connect, ip_collection, get_simple_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1057
        '''
1058
        target: test drop index interface
1059
        method: create collection and add vectors in it, create index, call drop index
1060
        expected: return code 0, and default index param
1061
        '''
1062
        index_param = get_simple_index["index_param"]
1063
        index_type = get_simple_index["index_type"]
1064
        status, mode = connect._cmd("mode")
1065
        assert status.OK()
1066
        # status, ids = connect.insert(ip_collection, vectors)
1067
        status = connect.create_index(ip_collection, index_type, index_param)
1068
        if str(mode) == "GPU" and (index_type == IndexType.IVF_PQ):
1069
            assert not status.OK()
1070
        else:
1071
            assert status.OK()
1072
        status, result = connect.get_index_info(ip_collection)
1073
        logging.getLogger().info(result)
1074
        status = connect.drop_index(ip_collection)
1075
        assert status.OK()
1076
        status, result = connect.get_index_info(ip_collection)
1077
        logging.getLogger().info(result)
1078
        assert result._collection_name == ip_collection
1079
        assert result._index_type == IndexType.FLAT
1080
1081
    def test_drop_index_partition(self, connect, ip_collection, get_simple_index):
1082
        '''
1083
        target: test drop index interface
1084
        method: create collection, create partition and add vectors in it, create index on collection, call drop collection index
1085
        expected: return code 0, and default index param
1086
        '''
1087
        index_param = get_simple_index["index_param"]
1088
        index_type = get_simple_index["index_type"]
1089
        status = connect.create_partition(ip_collection, tag)
1090
        status, ids = connect.insert(ip_collection, vectors, partition_tag=tag)
1091
        status = connect.create_index(ip_collection, index_type, index_param)
1092
        assert status.OK()
1093
        status, result = connect.get_index_info(ip_collection)
1094
        logging.getLogger().info(result)
1095
        status = connect.drop_index(ip_collection)
1096
        assert status.OK()
1097
        status, result = connect.get_index_info(ip_collection)
1098
        logging.getLogger().info(result)
1099
        assert result._collection_name == ip_collection
1100
        assert result._index_type == IndexType.FLAT
1101
1102
    def test_drop_index_partition_C(self, connect, ip_collection, get_simple_index):
1103
        '''
1104
        target: test drop index interface
1105
        method: create collection, create partitions and add vectors in it, create index on partitions, call drop partition index
1106
        expected: return code 0, and default index param
1107
        '''
1108
        new_tag = "new_tag"
1109
        index_param = get_simple_index["index_param"]
1110
        index_type = get_simple_index["index_type"]
1111
        status = connect.create_partition(ip_collection, tag)
1112
        status = connect.create_partition(ip_collection, new_tag)
1113
        status, ids = connect.insert(ip_collection, vectors)
1114
        status = connect.create_index(ip_collection, index_type, index_param)
1115
        assert status.OK()
1116
        status = connect.drop_index(ip_collection)
1117
        assert status.OK()
1118
        status, result = connect.get_index_info(ip_collection)
1119
        logging.getLogger().info(result)
1120
        assert result._collection_name == ip_collection
1121
        assert result._index_type == IndexType.FLAT
1122
1123 View Code Duplication
    @pytest.mark.level(2)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1124
    def test_drop_index_repeatly(self, connect, ip_collection, get_simple_index):
1125
        '''
1126
        target: test drop index repeatly
1127
        method: create index, call drop index, and drop again
1128
        expected: return code 0
1129
        '''
1130
        index_param = get_simple_index["index_param"]
1131
        index_type = get_simple_index["index_type"]
1132
        # status, ids = connect.insert(ip_collection, vectors)
1133
        status, mode = connect._cmd("mode")
1134
        assert status.OK()
1135
        # status, ids = connect.insert(ip_collection, vectors)
1136
        status = connect.create_index(ip_collection, index_type, index_param)
1137
        if str(mode) == "GPU" and (index_type == IndexType.IVF_PQ):
1138
            assert not status.OK()
1139
        else:
1140
            assert status.OK()        
1141
        status, result = connect.get_index_info(ip_collection)
1142
        logging.getLogger().info(result)
1143
        status = connect.drop_index(ip_collection)
1144
        assert status.OK()
1145
        status = connect.drop_index(ip_collection)
1146
        assert status.OK()
1147
        status, result = connect.get_index_info(ip_collection)
1148
        logging.getLogger().info(result)
1149
        assert result._collection_name == ip_collection
1150
        assert result._index_type == IndexType.FLAT
1151
1152
    # @pytest.mark.level(2)
1153
    # def test_drop_index_without_connect(self, dis_connect, ip_collection):
1154
    #     '''
1155
    #     target: test drop index without connection
1156
    #     method: drop index, and check if drop successfully
1157
    #     expected: raise exception
1158
    #     '''
1159
    #     nlist = NLIST
1160
    #     index_type = IndexType.IVFLAT
1161
    #     index_param = {"nlist": nlist}
1162
    #     with pytest.raises(Exception) as e:
1163
    #         status = dis_connect.drop_index(ip_collection, index_type, index_param)
1164
1165
    def test_drop_index_collection_not_create(self, connect, ip_collection):
1166
        '''
1167
        target: test drop index interface when index not created
1168
        method: create collection and add vectors in it, create index
1169
        expected: return code not equals to 0, drop index failed
1170
        '''
1171
        status, ids = connect.insert(ip_collection, vectors)
1172
        status, result = connect.get_index_info(ip_collection)
1173
        logging.getLogger().info(result)
1174
        # no create index
1175
        status = connect.drop_index(ip_collection)
1176
        logging.getLogger().info(status)
1177
        assert status.OK()
1178
1179 View Code Duplication
    @pytest.mark.level(2)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1180
    def test_create_drop_index_repeatly(self, connect, ip_collection, get_simple_index):
1181
        '''
1182
        target: test create / drop index repeatly, use the same index params
1183
        method: create index, drop index, four times
1184
        expected: return code 0
1185
        '''
1186
        index_param = get_simple_index["index_param"]
1187
        index_type = get_simple_index["index_type"]
1188
        status, ids = connect.insert(ip_collection, vectors)
1189
        for i in range(2):
1190
            status = connect.create_index(ip_collection, index_type, index_param)
1191
            assert status.OK()
1192
            status, result = connect.get_index_info(ip_collection)
1193
            logging.getLogger().info(result)
1194
            status = connect.drop_index(ip_collection)
1195
            assert status.OK()
1196
            status, result = connect.get_index_info(ip_collection)
1197
            logging.getLogger().info(result)
1198
            assert result._collection_name == ip_collection
1199
            assert result._index_type == IndexType.FLAT
1200
1201
    def test_create_drop_index_repeatly_different_index_params(self, connect, ip_collection):
1202
        '''
1203
        target: test create / drop index repeatly, use the different index params
1204
        method: create index, drop index, four times, each tme use different index_params to create index
1205
        expected: return code 0
1206
        '''
1207
        nlist = NLIST
1208
        indexs = [{"index_type": IndexType.IVFLAT, "index_param": {"nlist": nlist}}, {"index_type": IndexType.IVF_SQ8, "index_param": {"nlist": nlist}}]
1209
        status, ids = connect.insert(ip_collection, vectors)
1210
        for i in range(2):
1211
            status = connect.create_index(ip_collection, indexs[i]["index_type"], indexs[i]["index_param"])
1212
            assert status.OK()
1213
            status, result = connect.get_index_info(ip_collection)
1214
            assert result._params == indexs[i]["index_param"]
1215
            assert result._collection_name == ip_collection
1216
            assert result._index_type == indexs[i]["index_type"]
1217
            status, result = connect.get_index_info(ip_collection)
1218
            logging.getLogger().info(result)
1219
            status = connect.drop_index(ip_collection)
1220
            assert status.OK()
1221
            status, result = connect.get_index_info(ip_collection)
1222
            logging.getLogger().info(result)
1223
            assert result._collection_name == ip_collection
1224
            assert result._index_type == IndexType.FLAT
1225
1226
1227
class TestIndexJAC:
1228
    tmp, vectors = gen_binary_vectors(nb, dim)
1229
1230
    @pytest.fixture(
1231
        scope="function",
1232
        params=gen_index()
1233
    )
1234
    def get_index(self, request, connect):
1235
        if str(connect._cmd("mode")[1]) == "CPU":
1236
            if request.param["index_type"] == IndexType.IVF_SQ8H:
1237
                pytest.skip("sq8h not support in CPU mode")
1238
        if str(connect._cmd("mode")[1]) == "GPU":
1239
            if request.param["index_type"] == IndexType.IVF_PQ:
1240
                pytest.skip("ivfpq not support in GPU mode")
1241
        return request.param
1242
1243 View Code Duplication
    @pytest.fixture(
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1244
        scope="function",
1245
        params=gen_simple_index()
1246
    )
1247
    def get_simple_index(self, request, connect):
1248
        if str(connect._cmd("mode")[1]) == "CPU":
1249
            if request.param["index_type"] == IndexType.IVF_SQ8H:
1250
                pytest.skip("sq8h not support in CPU mode")
1251
        if str(connect._cmd("mode")[1]) == "GPU":
1252
            if request.param["index_type"] == IndexType.IVF_PQ:
1253
                pytest.skip("ivfpq not support in GPU mode")
1254
        return request.param
1255
1256 View Code Duplication
    @pytest.fixture(
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1257
        scope="function",
1258
        params=gen_simple_index()
1259
    )
1260
    def get_jaccard_index(self, request, connect):
1261
        logging.getLogger().info(request.param)
1262
        if request.param["index_type"] == IndexType.IVFLAT or request.param["index_type"] == IndexType.FLAT:
1263
            return request.param
1264
        else:
1265
            pytest.skip("Skip index Temporary")
1266
1267
    """
1268
    ******************************************************************
1269
      The following cases are used to test `create_index` function
1270
    ******************************************************************
1271
    """
1272
    @pytest.mark.timeout(BUILD_TIMEOUT)
1273
    def test_create_index(self, connect, jac_collection, get_jaccard_index):
1274
        '''
1275
        target: test create index interface
1276
        method: create collection and add vectors in it, create index
1277
        expected: return code equals to 0, and search success
1278
        '''
1279
        index_param = get_jaccard_index["index_param"]
1280
        index_type = get_jaccard_index["index_type"]
1281
        logging.getLogger().info(get_jaccard_index)
1282
        status, ids = connect.insert(jac_collection, self.vectors)
1283
        status = connect.create_index(jac_collection, index_type, index_param)
1284
        if index_type != IndexType.FLAT and index_type != IndexType.IVFLAT:
1285
            assert not status.OK()
1286
        else:
1287
            assert status.OK()
1288
1289
    @pytest.mark.timeout(BUILD_TIMEOUT)
1290
    def test_create_index_partition(self, connect, jac_collection, get_jaccard_index):
1291
        '''
1292
        target: test create index interface
1293
        method: create collection, create partition, and add vectors in it, create index
1294
        expected: return code equals to 0, and search success
1295
        '''
1296
        index_param = get_jaccard_index["index_param"]
1297
        index_type = get_jaccard_index["index_type"]
1298
        logging.getLogger().info(get_jaccard_index)
1299
        status = connect.create_partition(jac_collection, tag)
1300
        status, ids = connect.insert(jac_collection, self.vectors, partition_tag=tag)
1301
        status = connect.create_index(jac_collection, index_type, index_param)
1302
        assert status.OK()
1303
1304
    # @pytest.mark.level(2)
1305
    # def test_create_index_without_connect(self, dis_connect, jac_collection):
1306
    #     '''
1307
    #     target: test create index without connection
1308
    #     method: create collection and add vectors in it, check if added successfully
1309
    #     expected: raise exception
1310
    #     '''
1311
    #     nlist = NLIST
1312
    #     index_param = {"nlist": nlist}
1313
    #     with pytest.raises(Exception) as e:
1314
    #         status = dis_connect.create_index(jac_collection, IndexType.IVF_SQ8, index_param)
1315
1316 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1317
    def test_create_index_search_with_query_vectors(self, connect, jac_collection, get_jaccard_index):
1318
        '''
1319
        target: test create index interface, search with more query vectors
1320
        method: create collection and add vectors in it, create index
1321
        expected: return code equals to 0, and search success
1322
        '''
1323
        index_param = get_jaccard_index["index_param"]
1324
        index_type = get_jaccard_index["index_type"]
1325
        logging.getLogger().info(get_jaccard_index)
1326
        status, ids = connect.insert(jac_collection, self.vectors)
1327
        status = connect.create_index(jac_collection, index_type, index_param)
1328
        logging.getLogger().info(connect.get_index_info(jac_collection))
1329
        query_vecs = [self.vectors[0], self.vectors[1], self.vectors[2]]
1330
        top_k = 5
1331
        search_param = get_search_param(index_type)
1332
        status, result = connect.search(jac_collection, top_k, query_vecs, params=search_param)
1333
        logging.getLogger().info(result)
1334
        assert status.OK()
1335
        assert len(result) == len(query_vecs)
1336
1337
    """
1338
    ******************************************************************
1339
      The following cases are used to test `get_index_info` function
1340
    ******************************************************************
1341
    """
1342
1343
    def test_get_index_info(self, connect, jac_collection, get_jaccard_index):
1344
        '''
1345
        target: test describe index interface
1346
        method: create collection and add vectors in it, create index, call describe index
1347
        expected: return code 0, and index instructure
1348
        '''
1349
        index_param = get_jaccard_index["index_param"]
1350
        index_type = get_jaccard_index["index_type"]
1351
        logging.getLogger().info(get_jaccard_index)
1352
        # status, ids = connect.insert(jac_collection, vectors[:5000])
1353
        status = connect.create_index(jac_collection, index_type, index_param)
1354
        status, result = connect.get_index_info(jac_collection)
1355
        logging.getLogger().info(result)
1356
        assert result._collection_name == jac_collection
1357
        assert result._index_type == index_type
1358
        assert result._params == index_param
1359
1360 View Code Duplication
    def test_get_index_info_partition(self, connect, jac_collection, get_jaccard_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1361
        '''
1362
        target: test describe index interface
1363
        method: create collection, create partition and add vectors in it, create index, call describe index
1364
        expected: return code 0, and index instructure
1365
        '''
1366
        index_param = get_jaccard_index["index_param"]
1367
        index_type = get_jaccard_index["index_type"]
1368
        logging.getLogger().info(get_jaccard_index)
1369
        status = connect.create_partition(jac_collection, tag)
1370
        status, ids = connect.insert(jac_collection, vectors, partition_tag=tag)
1371
        status = connect.create_index(jac_collection, index_type, index_param)
1372
        status, result = connect.get_index_info(jac_collection)
1373
        logging.getLogger().info(result)
1374
        assert result._params == index_param
1375
        assert result._collection_name == jac_collection
1376
        assert result._index_type == index_type
1377
1378
    """
1379
    ******************************************************************
1380
      The following cases are used to test `drop_index` function
1381
    ******************************************************************
1382
    """
1383
1384 View Code Duplication
    def test_drop_index(self, connect, jac_collection, get_jaccard_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1385
        '''
1386
        target: test drop index interface
1387
        method: create collection and add vectors in it, create index, call drop index
1388
        expected: return code 0, and default index param
1389
        '''
1390
        index_param = get_jaccard_index["index_param"]
1391
        index_type = get_jaccard_index["index_type"]
1392
        status, mode = connect._cmd("mode")
1393
        assert status.OK()
1394
        # status, ids = connect.insert(ip_collection, vectors)
1395
        status = connect.create_index(jac_collection, index_type, index_param)
1396
        assert status.OK()
1397
        status, result = connect.get_index_info(jac_collection)
1398
        logging.getLogger().info(result)
1399
        status = connect.drop_index(jac_collection)
1400
        assert status.OK()
1401
        status, result = connect.get_index_info(jac_collection)
1402
        logging.getLogger().info(result)
1403
        assert result._collection_name == jac_collection
1404
        assert result._index_type == IndexType.FLAT
1405
1406
    def test_drop_index_partition(self, connect, jac_collection, get_jaccard_index):
1407
        '''
1408
        target: test drop index interface
1409
        method: create collection, create partition and add vectors in it, create index on collection, call drop collection index
1410
        expected: return code 0, and default index param
1411
        '''
1412
        index_param = get_jaccard_index["index_param"]
1413
        index_type = get_jaccard_index["index_type"]
1414
        status = connect.create_partition(jac_collection, tag)
1415
        status, ids = connect.insert(jac_collection, vectors, partition_tag=tag)
1416
        status = connect.create_index(jac_collection, index_type, index_param)
1417
        assert status.OK()
1418
        status, result = connect.get_index_info(jac_collection)
1419
        logging.getLogger().info(result)
1420
        status = connect.drop_index(jac_collection)
1421
        assert status.OK()
1422
        status, result = connect.get_index_info(jac_collection)
1423
        logging.getLogger().info(result)
1424
        assert result._collection_name == jac_collection
1425
        assert result._index_type == IndexType.FLAT
1426
1427
1428
class TestIndexBinary:
1429
    tmp, vectors = gen_binary_vectors(nb, dim)
1430
1431
    @pytest.fixture(
1432
        scope="function",
1433
        params=gen_index()
1434
    )
1435
    def get_index(self, request, connect):
1436
        if str(connect._cmd("mode")[1]) == "CPU":
1437
            if request.param["index_type"] == IndexType.IVF_SQ8H:
1438
                pytest.skip("sq8h not support in CPU mode")
1439
        if request.param["index_type"] == IndexType.IVF_PQ or request.param["index_type"] == IndexType.HNSW:
1440
            pytest.skip("Skip PQ Temporary")
1441
        return request.param
1442
1443 View Code Duplication
    @pytest.fixture(
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1444
        scope="function",
1445
        params=gen_simple_index()
1446
    )
1447
    def get_simple_index(self, request, connect):
1448
        if str(connect._cmd("mode")[1]) == "CPU":
1449
            if request.param["index_type"] == IndexType.IVF_SQ8H:
1450
                pytest.skip("sq8h not support in CPU mode")
1451
        if request.param["index_type"] == IndexType.IVF_PQ or request.param["index_type"] == IndexType.HNSW:
1452
            pytest.skip("Skip PQ Temporary")
1453
        return request.param
1454
1455 View Code Duplication
    @pytest.fixture(
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1456
        scope="function",
1457
        params=gen_simple_index()
1458
    )
1459
    def get_hamming_index(self, request, connect):
1460
        logging.getLogger().info(request.param)
1461
        if request.param["index_type"] == IndexType.IVFLAT or request.param["index_type"] == IndexType.FLAT:
1462
            return request.param
1463
        else:
1464
            pytest.skip("Skip index Temporary")
1465
1466
    @pytest.fixture(
1467
        scope="function",
1468
        params=gen_simple_index()
1469
    )
1470
    def get_substructure_index(self, request, connect):
1471
        logging.getLogger().info(request.param)
1472
        if request.param["index_type"] == IndexType.FLAT:
1473
            return request.param
1474
        else:
1475
            pytest.skip("Skip index Temporary")
1476
1477
    @pytest.fixture(
1478
        scope="function",
1479
        params=gen_simple_index()
1480
    )
1481
    def get_superstructure_index(self, request, connect):
1482
        logging.getLogger().info(request.param)
1483
        if request.param["index_type"] == IndexType.FLAT:
1484
            return request.param
1485
        else:
1486
            pytest.skip("Skip index Temporary")
1487
1488
    """
1489
    ******************************************************************
1490
      The following cases are used to test `create_index` function
1491
    ******************************************************************
1492
    """
1493
    @pytest.mark.timeout(BUILD_TIMEOUT)
1494
    def test_create_index(self, connect, ham_collection, get_hamming_index):
1495
        '''
1496
        target: test create index interface
1497
        method: create collection and add vectors in it, create index
1498
        expected: return code equals to 0, and search success
1499
        '''
1500
        index_param = get_hamming_index["index_param"]
1501
        index_type = get_hamming_index["index_type"]
1502
        logging.getLogger().info(get_hamming_index)
1503
        status, ids = connect.insert(ham_collection, self.vectors)
1504
        status = connect.create_index(ham_collection, index_type, index_param)
1505
        if index_type != IndexType.FLAT and index_type != IndexType.IVFLAT:
1506
            assert not status.OK()
1507
        else:
1508
            assert status.OK()
1509
1510 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1511
    def test_create_index_partition(self, connect, ham_collection, get_hamming_index):
1512
        '''
1513
        target: test create index interface
1514
        method: create collection, create partition, and add vectors in it, create index
1515
        expected: return code equals to 0, and search success
1516
        '''
1517
        index_param = get_hamming_index["index_param"]
1518
        index_type = get_hamming_index["index_type"]
1519
        logging.getLogger().info(get_hamming_index)
1520
        status = connect.create_partition(ham_collection, tag)
1521
        status, ids = connect.insert(ham_collection, self.vectors, partition_tag=tag)
1522
        status = connect.create_index(ham_collection, index_type, index_param)
1523
        assert status.OK()
1524
        status, res = connect.count_entities(ham_collection)
1525
        assert res == len(self.vectors)
1526
1527 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1528
    def test_create_index_partition_structure(self, connect, substructure_collection, get_substructure_index):
1529
        '''
1530
        target: test create index interface
1531
        method: create collection, create partition, and add vectors in it, create index
1532
        expected: return code equals to 0, and search success
1533
        '''
1534
        index_param = get_substructure_index["index_param"]
1535
        index_type = get_substructure_index["index_type"]
1536
        logging.getLogger().info(get_substructure_index)
1537
        status = connect.create_partition(substructure_collection, tag)
1538
        status, ids = connect.insert(substructure_collection, self.vectors, partition_tag=tag)
1539
        status = connect.create_index(substructure_collection, index_type, index_param)
1540
        assert status.OK()
1541
        status, res = connect.count_entities(substructure_collection,)
1542
        assert res == len(self.vectors)
1543
1544
    # @pytest.mark.level(2)
1545
    # def test_create_index_without_connect(self, dis_connect, ham_collection):
1546
    #     '''
1547
    #     target: test create index without connection
1548
    #     method: create collection and add vectors in it, check if added successfully
1549
    #     expected: raise exception
1550
    #     '''
1551
    #     nlist = NLIST
1552
    #     index_param = {"nlist": nlist}
1553
    #     with pytest.raises(Exception) as e:
1554
    #         status = dis_connect.create_index(ham_collection, IndexType.IVF_SQ8, index_param)
1555
1556 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1557
    def test_create_index_search_with_query_vectors(self, connect, ham_collection, get_hamming_index):
1558
        '''
1559
        target: test create index interface, search with more query vectors
1560
        method: create collection and add vectors in it, create index
1561
        expected: return code equals to 0, and search success
1562
        '''
1563
        index_param = get_hamming_index["index_param"]
1564
        index_type = get_hamming_index["index_type"]
1565
        logging.getLogger().info(get_hamming_index)
1566
        status, ids = connect.insert(ham_collection, self.vectors)
1567
        status = connect.create_index(ham_collection,  index_type, index_param)
1568
        logging.getLogger().info(connect.get_index_info(ham_collection))
1569
        query_vecs = [self.vectors[0], self.vectors[1], self.vectors[2]]
1570
        top_k = 5
1571
        search_param = get_search_param(index_type)
1572
        status, result = connect.search(ham_collection, top_k, query_vecs, params=search_param)
1573
        logging.getLogger().info(result)
1574
        assert status.OK()
1575
        assert len(result) == len(query_vecs)
1576
1577 View Code Duplication
    @pytest.mark.timeout(BUILD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1578
    def test_create_index_search_with_query_vectors_superstructure(self, connect, superstructure_collection, get_superstructure_index):
1579
        '''
1580
        target: test create index interface, search with more query vectors
1581
        method: create collection and add vectors in it, create index
1582
        expected: return code equals to 0, and search success
1583
        '''
1584
        index_param = get_superstructure_index["index_param"]
1585
        index_type = get_superstructure_index["index_type"]
1586
        logging.getLogger().info(get_superstructure_index)
1587
        status, ids = connect.insert(superstructure_collection, self.vectors)
1588
        status = connect.create_index(superstructure_collection, index_type, index_param)
1589
        logging.getLogger().info(connect.get_index_info(superstructure_collection))
1590
        query_vecs = [self.vectors[0], self.vectors[1], self.vectors[2]]
1591
        top_k = 5
1592
        search_param = get_search_param(index_type)
1593
        status, result = connect.search(superstructure_collection, top_k, query_vecs, params=search_param)
1594
        logging.getLogger().info(result)
1595
        assert status.OK()
1596
        assert len(result) == len(query_vecs)
1597
1598
    """
1599
    ******************************************************************
1600
      The following cases are used to test `get_index_info` function
1601
    ******************************************************************
1602
    """
1603
1604
    def test_get_index_info(self, connect, ham_collection, get_hamming_index):
1605
        '''
1606
        target: test describe index interface
1607
        method: create collection and add vectors in it, create index, call describe index
1608
        expected: return code 0, and index instructure
1609
        '''
1610
        index_param = get_hamming_index["index_param"]
1611
        index_type = get_hamming_index["index_type"]
1612
        logging.getLogger().info(get_hamming_index)
1613
        # status, ids = connect.insert(jac_collection, vectors[:5000])
1614
        status = connect.create_index(ham_collection, index_type, index_param)
1615
        status, result = connect.get_index_info(ham_collection)
1616
        logging.getLogger().info(result)
1617
        assert result._collection_name == ham_collection
1618
        assert result._index_type == index_type
1619
        assert result._params == index_param
1620
1621 View Code Duplication
    def test_get_index_info_partition(self, connect, ham_collection, get_hamming_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1622
        '''
1623
        target: test describe index interface
1624
        method: create collection, create partition and add vectors in it, create index, call describe index
1625
        expected: return code 0, and index instructure
1626
        '''
1627
        index_param = get_hamming_index["index_param"]
1628
        index_type = get_hamming_index["index_type"]
1629
        logging.getLogger().info(get_hamming_index)
1630
        status = connect.create_partition(ham_collection, tag)
1631
        status, ids = connect.insert(ham_collection, vectors, partition_tag=tag)
1632
        status = connect.create_index(ham_collection, index_type, index_param)
1633
        status, result = connect.get_index_info(ham_collection)
1634
        logging.getLogger().info(result)
1635
        assert result._params == index_param
1636
        assert result._collection_name == ham_collection
1637
        assert result._index_type == index_type
1638
1639 View Code Duplication
    def test_get_index_info_partition_superstructrue(self, connect, superstructure_collection, get_superstructure_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1640
        '''
1641
        target: test describe index interface
1642
        method: create collection, create partition and add vectors in it, create index, call describe index
1643
        expected: return code 0, and index instructure
1644
        '''
1645
        index_param = get_superstructure_index["index_param"]
1646
        index_type = get_superstructure_index["index_type"]
1647
        logging.getLogger().info(get_superstructure_index)
1648
        status = connect.create_partition(superstructure_collection, tag)
1649
        status, ids = connect.insert(superstructure_collection, vectors, partition_tag=tag)
1650
        status = connect.create_index(superstructure_collection, index_type, index_param)
1651
        status, result = connect.get_index_info(superstructure_collection)
1652
        logging.getLogger().info(result)
1653
        assert result._params == index_param
1654
        assert result._collection_name == superstructure_collection
1655
        assert result._index_type == index_type
1656
1657
    """
1658
    ******************************************************************
1659
      The following cases are used to test `drop_index` function
1660
    ******************************************************************
1661
    """
1662
1663 View Code Duplication
    def test_drop_index(self, connect, ham_collection, get_hamming_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1664
        '''
1665
        target: test drop index interface
1666
        method: create collection and add vectors in it, create index, call drop index
1667
        expected: return code 0, and default index param
1668
        '''
1669
        index_param = get_hamming_index["index_param"]
1670
        index_type = get_hamming_index["index_type"]
1671
        status, mode = connect._cmd("mode")
1672
        assert status.OK()
1673
        # status, ids = connect.insert(ip_collection, vectors)
1674
        status = connect.create_index(ham_collection, index_type, index_param)
1675
        assert status.OK()
1676
        status, result = connect.get_index_info(ham_collection)
1677
        logging.getLogger().info(result)
1678
        status = connect.drop_index(ham_collection)
1679
        assert status.OK()
1680
        status, result = connect.get_index_info(ham_collection)
1681
        logging.getLogger().info(result)
1682
        assert result._collection_name == ham_collection
1683
        assert result._index_type == IndexType.FLAT
1684
1685 View Code Duplication
    def test_drop_index_substructure(self, connect, substructure_collection, get_substructure_index):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1686
        '''
1687
        target: test drop index interface
1688
        method: create collection and add vectors in it, create index, call drop index
1689
        expected: return code 0, and default index param
1690
        '''
1691
        index_param = get_substructure_index["index_param"]
1692
        index_type = get_substructure_index["index_type"]
1693
        status, mode = connect._cmd("mode")
1694
        assert status.OK()
1695
        status = connect.create_index(substructure_collection, index_type, index_param)
1696
        assert status.OK()
1697
        status, result = connect.get_index_info(substructure_collection)
1698
        logging.getLogger().info(result)
1699
        status = connect.drop_index(substructure_collection)
1700
        assert status.OK()
1701
        status, result = connect.get_index_info(substructure_collection)
1702
        logging.getLogger().info(result)
1703
        assert result._collection_name == substructure_collection
1704
        assert result._index_type == IndexType.FLAT
1705
1706
    def test_drop_index_partition(self, connect, ham_collection, get_hamming_index):
1707
        '''
1708
        target: test drop index interface
1709
        method: create collection, create partition and add vectors in it, create index on collection, call drop collection index
1710
        expected: return code 0, and default index param
1711
        '''
1712
        index_param = get_hamming_index["index_param"]
1713
        index_type = get_hamming_index["index_type"]
1714
        status = connect.create_partition(ham_collection, tag)
1715
        status, ids = connect.insert(ham_collection, vectors, partition_tag=tag)
1716
        status = connect.create_index(ham_collection, index_type, index_param)
1717
        assert status.OK()
1718
        status, result = connect.get_index_info(ham_collection)
1719
        logging.getLogger().info(result)
1720
        status = connect.drop_index(ham_collection)
1721
        assert status.OK()
1722
        status, result = connect.get_index_info(ham_collection)
1723
        logging.getLogger().info(result)
1724
        assert result._collection_name == ham_collection
1725
        assert result._index_type == IndexType.FLAT
1726
1727
class TestIndexCollectionInvalid(object):
1728
    """
1729
    Test create / describe / drop index interfaces with invalid collection names
1730
    """
1731
    @pytest.fixture(
1732
        scope="function",
1733
        params=gen_invalid_collection_names()
1734
    )
1735
    def get_collection_name(self, request):
1736
        yield request.param
1737
1738
    @pytest.mark.level(1)
1739
    def test_create_index_with_invalid_collectionname(self, connect, get_collection_name):
1740
        collection_name = get_collection_name
1741
        nlist = NLIST
1742
        index_param = {"nlist": nlist}
1743
        status = connect.create_index(collection_name, IndexType.IVF_SQ8, index_param)
1744
        assert not status.OK()
1745
1746
    @pytest.mark.level(1)
1747
    def test_get_index_info_with_invalid_collectionname(self, connect, get_collection_name):
1748
        collection_name = get_collection_name
1749
        status, result = connect.get_index_info(collection_name)
1750
        assert not status.OK()   
1751
1752
    @pytest.mark.level(1)
1753
    def test_drop_index_with_invalid_collectionname(self, connect, get_collection_name):
1754
        collection_name = get_collection_name
1755
        status = connect.drop_index(collection_name)
1756
        assert not status.OK()
1757
1758
1759
class TestCreateIndexParamsInvalid(object):
1760
    """
1761
    Test Building index with invalid collection names, collection names not in db
1762
    """
1763
    @pytest.fixture(
1764
        scope="function",
1765
        params=gen_invalid_index()
1766
    )
1767
    def get_index(self, request):
1768
        yield request.param
1769
1770
    @pytest.mark.level(1)
1771
    def test_create_index_with_invalid_index_params(self, connect, collection, get_index):
1772
        index_param = get_index["index_param"]
1773
        index_type = get_index["index_type"]
1774
        logging.getLogger().info(get_index)
1775
        # status, ids = connect.insert(collection, vectors)
1776
        if (not index_type) or (not isinstance(index_type, IndexType)):
1777
            with pytest.raises(Exception) as e:
1778
                status = connect.create_index(collection, index_type, index_param)
1779
        else:
1780
            status = connect.create_index(collection, index_type, index_param)
1781
            assert not status.OK()
1782
1783
    """
1784
    Test Building index with invalid nlist
1785
    """
1786
    @pytest.fixture(
1787
        scope="function",
1788
        params=[IndexType.FLAT,IndexType.IVFLAT,IndexType.IVF_SQ8,IndexType.IVF_SQ8H]
1789
    )
1790
    def get_index_type(self, request):
1791
        yield request.param
1792
1793
    def test_create_index_with_invalid_nlist(self, connect, collection, get_index_type):
1794
        status, ids = connect.insert(collection, vectors)
1795
        status = connect.create_index(collection, get_index_type, {"nlist": INVALID_NLIST})
1796
        if get_index_type != IndexType.FLAT:
1797
            assert not status.OK()
1798
1799
    '''
1800
    Test Building index with empty params
1801
    '''
1802
    def test_create_index_with_empty_param(self, connect, collection, get_index_type):
1803
        logging.getLogger().info(get_index_type)
1804
        status = connect.create_index(collection, get_index_type, {})
1805
        if get_index_type != IndexType.FLAT :
1806
            assert not status.OK()
1807
        status, result = connect.get_index_info(collection)
1808
        logging.getLogger().info(result)
1809
        assert result._collection_name == collection
1810
        assert result._index_type == IndexType.FLAT
1811
1812
class TestIndexAsync:
1813
    @pytest.fixture(scope="function", autouse=True)
1814
    def skip_http_check(self, args):
1815
        if args["handler"] == "HTTP":
1816
            pytest.skip("skip in http mode")
1817
1818
    """
1819
    ******************************************************************
1820
      The following cases are used to test `create_index` function
1821
    ******************************************************************
1822
    """
1823
    @pytest.fixture(
1824
        scope="function",
1825
        params=gen_index()
1826
    )
1827
    def get_index(self, request, connect):
1828
        if str(connect._cmd("mode")[1]) == "CPU":
1829
            if request.param["index_type"] == IndexType.IVF_SQ8H:
1830
                pytest.skip("sq8h not support in CPU mode")
1831
        if str(connect._cmd("mode")[1]) == "GPU":
1832
            if request.param["index_type"] == IndexType.IVF_PQ:
1833
                pytest.skip("ivfpq not support in GPU mode")
1834
        return request.param
1835
1836 View Code Duplication
    @pytest.fixture(
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1837
        scope="function",
1838
        params=gen_simple_index()
1839
    )
1840
    def get_simple_index(self, request, connect):
1841
        if str(connect._cmd("mode")[1]) == "CPU":
1842
            if request.param["index_type"] == IndexType.IVF_SQ8H:
1843
                pytest.skip("sq8h not support in CPU mode")
1844
        if str(connect._cmd("mode")[1]) == "GPU":
1845
            # if request.param["index_type"] == IndexType.IVF_PQ:
1846
            if request.param["index_type"] not in [IndexType.IVF_FLAT]:
1847
                # pytest.skip("ivfpq not support in GPU mode")
1848
                pytest.skip("debug ivf_flat in GPU mode")
1849
        return request.param
1850
1851
    def check_status(self, status):
1852
        logging.getLogger().info("In callback check status")
1853
        assert status.OK()
1854
1855
    """
1856
    ******************************************************************
1857
      The following cases are used to test `create_index` function
1858
    ******************************************************************
1859
    """
1860
1861
    @pytest.mark.timeout(BUILD_TIMEOUT)
1862
    def test_create_index(self, connect, collection, get_simple_index):
1863
        '''
1864
        target: test create index interface
1865
        method: create collection and add vectors in it, create index
1866
        expected: return code equals to 0, and search success
1867
        '''
1868
        index_param = get_simple_index["index_param"]
1869
        index_type = get_simple_index["index_type"]
1870
        logging.getLogger().info(get_simple_index)
1871
        vectors = gen_vectors(nb, dim)
1872
        status, ids = connect.insert(collection, vectors)
1873
        logging.getLogger().info("start index")
1874
        # future = connect.create_index(collection, index_type, index_param, _async=True, _callback=self.check_status) 
1875
        future = connect.create_index(collection, index_type, index_param, _async=True) 
1876
        logging.getLogger().info("before result")
1877
        status = future.result()
1878
        assert status.OK()
1879
1880
    def test_create_index_with_invalid_collectionname(self, connect):
1881
        collection_name = " "
1882
        nlist = NLIST
1883
        index_param = {"nlist": nlist}
1884
        future = connect.create_index(collection_name, IndexType.IVF_SQ8, index_param, _async=True)
1885
        status = future.result()
1886
        assert not status.OK()
1887
1888