TestAddIP.test_add_vector_create_collection()   A
last analyzed

Complexity

Conditions 1

Size

Total Lines 14
Code Lines 9

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 1
eloc 9
nop 3
dl 0
loc 14
rs 9.95
c 0
b 0
f 0
1
import time
2
import pdb
3
import threading
4
import logging
5
import threading
6
from multiprocessing import Pool, Process
7
import pytest
8
from milvus import IndexType, MetricType
9
from utils import *
10
11
12
dim = 128
13
index_file_size = 10
14
collection_id = "test_add"
15
ADD_TIMEOUT = 60
16
tag = "1970-01-01"
17
add_interval_time = 5
18
nb = 6000
19
20
21
class TestAddBase:
22
    """
23
    ******************************************************************
24
      The following cases are used to test `insert` function
25
    ******************************************************************
26
    """
27 View Code Duplication
    @pytest.fixture(
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
28
        scope="function",
29
        params=gen_simple_index()
30
    )
31
    def get_simple_index(self, request, connect):
32
        if str(connect._cmd("mode")[1]) == "CPU":
33
            if request.param["index_type"] == IndexType.IVF_SQ8H:
34
                pytest.skip("sq8h not support in cpu mode")
35
        if request.param["index_type"] == IndexType.IVF_PQ:
36
            pytest.skip("Skip PQ Temporary")
37
        return request.param
38
39
    def test_add_vector_create_collection(self, connect, collection):
40
        '''
41
        target: test add vector, then create collection again
42
        method: add vector and create collection
43
        expected: status not ok
44
        '''
45
        vector = gen_single_vector(dim)
46
        status, ids = connect.insert(collection, vector)
47
        param = {'collection_name': collection,
48
                 'dimension': dim,
49
                 'index_file_size': index_file_size,
50
                 'metric_type': MetricType.L2}
51
        status = connect.create_collection(param)
52
        assert not status.OK()
53
54
    def test_add_vector_has_collection(self, connect, collection):
55
        '''
56
        target: test add vector, then check collection existence
57
        method: add vector and call Hascollection
58
        expected: collection exists, status ok
59
        '''
60
        vector = gen_single_vector(dim)
61
        status, ids = connect.insert(collection, vector)
62
        assert assert_has_collection(connect, collection)
63
64
    def test_add_vector_with_empty_vector(self, connect, collection):
65
        '''
66
        target: test add vectors with empty vectors list
67
        method: set empty vectors list as add method params
68
        expected: raises a Exception
69
        '''
70
        vector = []
71
        with pytest.raises(Exception) as e:
72
            status, ids = connect.insert(collection, vector)
73
74
    def test_add_vector_with_None(self, connect, collection):
75
        '''
76
        target: test add vectors with None
77
        method: set None as add method params
78
        expected: raises a Exception
79
        '''
80
        vector = None
81
        with pytest.raises(Exception) as e:
82
            status, ids = connect.insert(collection, vector)
83
84
    @pytest.mark.timeout(ADD_TIMEOUT)
85
    def test_drop_collection_add_vector(self, connect, collection):
86
        '''
87
        target: test add vector after collection deleted
88
        method: delete collection and add vector
89
        expected: status not ok
90
        '''
91
        status = connect.drop_collection(collection)
92
        vector = gen_single_vector(dim)
93
        status, ids = connect.insert(collection, vector)
94
        assert not status.OK()
95
96 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
97
    def test_drop_collection_add_vector_another(self, connect, collection):
98
        '''
99
        target: test add vector to collection_1 after collection_2 deleted
100
        method: delete collection_2 and add vector to collection_1
101
        expected: status ok
102
        '''
103
        param = {'collection_name': gen_unique_str(),
104
                 'dimension': dim,
105
                 'index_file_size': index_file_size,
106
                 'metric_type': MetricType.L2}
107
        status = connect.create_collection(param)
108
        status = connect.drop_collection(collection)
109
        vector = gen_single_vector(dim)
110
        status, ids = connect.insert(param['collection_name'], vector)
111
        assert status.OK()
112
113
    @pytest.mark.timeout(ADD_TIMEOUT)
114
    def test_add_vector_drop_collection(self, connect, collection):
115
        '''
116
        target: test delete collection after add vector
117
        method: add vector and delete collection
118
        expected: status ok
119
        '''
120
        vector = gen_single_vector(dim)
121
        status, ids = connect.insert(collection, vector)
122
        status = connect.drop_collection(collection)
123
        assert status.OK()
124
125
    @pytest.mark.timeout(ADD_TIMEOUT)
126
    def test_add_vector_delete_another_collection(self, connect, collection):
127
        '''
128
        target: test delete collection_1 collection after add vector to collection_2
129
        method: add vector and delete collection
130
        expected: status ok
131
        '''
132
        param = {'collection_name': gen_unique_str(),
133
                 'dimension': dim,
134
                 'index_file_size': index_file_size,
135
                 'metric_type': MetricType.L2}
136
        status = connect.create_collection(param)
137
        vector = gen_single_vector(dim)
138
        status, ids = connect.insert(collection, vector)
139
        assert status.OK()
140
141
    @pytest.mark.timeout(ADD_TIMEOUT)
142
    def test_add_vector_sleep_drop_collection(self, connect, collection):
143
        '''
144
        target: test delete collection after add vector for a while
145
        method: add vector, sleep, and delete collection
146
        expected: status ok
147
        '''
148
        vector = gen_single_vector(dim)
149
        status, ids = connect.insert(collection, vector)
150
        assert status.OK()
151
        connect.flush([collection])
152
        status = connect.drop_collection(collection)
153
        assert status.OK()
154
155 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
156
    def test_add_vector_sleep_delete_another_collection(self, connect, collection):
157
        '''
158
        target: test delete collection_1 collection after add vector to collection_2 for a while
159
        method: add vector , sleep, and delete collection
160
        expected: status ok
161
        '''
162
        param = {'collection_name': gen_unique_str(),
163
                 'dimension': dim,
164
                 'index_file_size': index_file_size,
165
                 'metric_type': MetricType.L2}
166
        status = connect.create_collection(param)
167
        vector = gen_single_vector(dim)
168
        status, ids = connect.insert(collection, vector)
169
        connect.flush([collection])
170
        status = connect.drop_collection(param['collection_name'])
171
        assert status.OK()
172
173 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
174
    def test_create_index_add_vector(self, connect, collection, get_simple_index):
175
        '''
176
        target: test add vector after build index
177
        method: build index and add vector
178
        expected: status ok
179
        '''
180
        index_param = get_simple_index["index_param"]
181
        index_type = get_simple_index["index_type"]
182
        status = connect.create_index(collection, index_type, index_param)
183
        vector = gen_single_vector(dim)
184
        status, ids = connect.insert(collection, vector)
185
        assert status.OK()
186
187 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
188
    def test_create_index_add_vector_another(self, connect, collection, get_simple_index):
189
        '''
190
        target: test add vector to collection_2 after build index for collection_1
191
        method: build index and add vector
192
        expected: status ok
193
        '''
194
        index_param = get_simple_index["index_param"]
195
        index_type = get_simple_index["index_type"]
196
        param = {'collection_name': gen_unique_str(),
197
                 'dimension': dim,
198
                 'index_file_size': index_file_size,
199
                 'metric_type': MetricType.L2}
200
        status = connect.create_collection(param)
201
        status = connect.create_index(collection, index_type, index_param)
202
        vector = gen_single_vector(dim)
203
        status, ids = connect.insert(collection, vector)
204
        connect.drop_collection(param['collection_name'])
205
        assert status.OK()
206
207
    @pytest.mark.timeout(ADD_TIMEOUT)
208
    def test_add_vector_create_index(self, connect, collection, get_simple_index):
209
        '''
210
        target: test build index add after vector
211
        method: add vector and build index
212
        expected: status ok
213
        '''
214
        index_param = get_simple_index["index_param"]
215
        index_type = get_simple_index["index_type"]
216
        logging.getLogger().info(index_param)
217
        vector = gen_single_vector(dim)
218
        status, ids = connect.insert(collection, vector)
219
        status = connect.create_index(collection, index_type, index_param)
220
        assert status.OK()
221
222 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
223
    def test_add_vector_create_index_another(self, connect, collection, get_simple_index):
224
        '''
225
        target: test add vector to collection_2 after build index for collection_1
226
        method: build index and add vector
227
        expected: status ok
228
        '''
229
        index_param = get_simple_index["index_param"]
230
        index_type = get_simple_index["index_type"]
231
        param = {'collection_name': gen_unique_str(),
232
                 'dimension': dim,
233
                 'index_file_size': index_file_size,
234
                 'metric_type': MetricType.L2}
235
        status = connect.create_collection(param)
236
        vector = gen_single_vector(dim)
237
        status, ids = connect.insert(collection, vector)
238
        status = connect.create_index(
239
            param['collection_name'], index_type, index_param)
240
        assert status.OK()
241
242 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
243
    def test_add_vector_sleep_create_index(self, connect, collection, get_simple_index):
244
        '''
245
        target: test build index add after vector for a while
246
        method: add vector and build index
247
        expected: status ok
248
        '''
249
        index_param = get_simple_index["index_param"]
250
        index_type = get_simple_index["index_type"]
251
        vector = gen_single_vector(dim)
252
        status, ids = connect.insert(collection, vector)
253
        connect.flush([collection])
254
        status = connect.create_index(collection, index_type, index_param)
255
        assert status.OK()
256
257 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
258
    def test_add_vector_sleep_create_index_another(self, connect, collection, get_simple_index):
259
        '''
260
        target: test add vector to collection_2 after build index for collection_1 for a while
261
        method: build index and add vector
262
        expected: status ok
263
        '''
264
        index_param = get_simple_index["index_param"]
265
        index_type = get_simple_index["index_type"]
266
        param = {'collection_name': gen_unique_str(),
267
                 'dimension': dim,
268
                 'index_file_size': index_file_size,
269
                 'metric_type': MetricType.L2}
270
        status = connect.create_collection(param)
271
        vector = gen_single_vector(dim)
272
        status, ids = connect.insert(collection, vector)
273
        connect.flush([collection])
274
        status = connect.create_index(
275
            param['collection_name'], index_type, index_param)
276
        assert status.OK()
277
278
    @pytest.mark.timeout(ADD_TIMEOUT)
279
    def test_search_vector_add_vector(self, connect, collection):
280
        '''
281
        target: test add vector after search collection
282
        method: search collection and add vector
283
        expected: status ok
284
        '''
285
        vector = gen_single_vector(dim)
286
        status, result = connect.search(collection, 1, vector)
287
        status, ids = connect.insert(collection, vector)
288
        assert status.OK()
289
290 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
291
    def test_search_vector_add_vector_another(self, connect, collection):
292
        '''
293
        target: test add vector to collection_1 after search collection_2
294
        method: search collection and add vector
295
        expected: status ok
296
        '''
297
        param = {'collection_name': gen_unique_str(),
298
                 'dimension': dim,
299
                 'index_file_size': index_file_size,
300
                 'metric_type': MetricType.L2}
301
        status = connect.create_collection(param)
302
        vector = gen_single_vector(dim)
303
        status, result = connect.search(collection, 1, vector)
304
        status, ids = connect.insert(param['collection_name'], vector)
305
        assert status.OK()
306
307
    @pytest.mark.timeout(ADD_TIMEOUT)
308
    def test_add_vector_search_vector(self, connect, collection):
309
        '''
310
        target: test search vector after add vector
311
        method: add vector and search collection
312
        expected: status ok
313
        '''
314
        vector = gen_single_vector(dim)
315
        status, ids = connect.insert(collection, vector)
316
        assert status.OK()
317
        connect.flush([collection])
318
        status, result = connect.search(collection, 1, vector)
319
        assert status.OK()
320
321 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
322
    def test_add_vector_search_vector_another(self, connect, collection):
323
        '''
324
        target: test add vector to collection_1 after search collection_2
325
        method: search collection and add vector
326
        expected: status ok
327
        '''
328
        param = {'collection_name': gen_unique_str(),
329
                 'dimension': dim,
330
                 'index_file_size': index_file_size,
331
                 'metric_type': MetricType.L2}
332
        status = connect.create_collection(param)
333
        vector = gen_single_vector(dim)
334
        status, ids = connect.insert(collection, vector)
335
        status, result = connect.search(param['collection_name'], 1, vector)
336
        assert status.OK()
337
338
    @pytest.mark.timeout(ADD_TIMEOUT)
339
    def test_add_vector_sleep_search_vector(self, connect, collection):
340
        '''
341
        target: test search vector after add vector after a while
342
        method: add vector, sleep, and search collection
343
        expected: status ok
344
        '''
345
        vector = gen_single_vector(dim)
346
        status, ids = connect.insert(collection, vector)
347
        connect.flush([collection])
348
        status, result = connect.search(collection, 1, vector)
349
        assert status.OK()
350
351 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
352
    def test_add_vector_sleep_search_vector_another(self, connect, collection):
353
        '''
354
        target: test add vector to collection_1 after search collection_2 a while
355
        method: search collection , sleep, and add vector
356
        expected: status ok
357
        '''
358
        param = {'collection_name': gen_unique_str(),
359
                 'dimension': dim,
360
                 'index_file_size': index_file_size,
361
                 'metric_type': MetricType.L2}
362
        status = connect.create_collection(param)
363
        vector = gen_single_vector(dim)
364
        status, ids = connect.insert(collection, vector)
365
        connect.flush([collection])
366
        status, result = connect.search(param['collection_name'], 1, vector)
367
        assert status.OK()
368
369
    """
370
    ******************************************************************
371
      The following cases are used to test `insert` function
372
    ******************************************************************
373
    """
374
375 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
376
    def test_insert_ids(self, connect, collection):
377
        '''
378
        target: test add vectors in collection, use customize ids
379
        method: create collection and add vectors in it, check the ids returned and the collection length after vectors added
380
        expected: the length of ids and the collection row count
381
        '''
382
        nq = 5
383
        top_k = 1
384
        vectors = gen_vectors(nq, dim)
385
        ids = [i for i in range(nq)]
386
        status, ids = connect.insert(collection, vectors, ids)
387
        connect.flush([collection])
388
        assert status.OK()
389
        assert len(ids) == nq
390
        status, result = connect.search(
391
            collection, top_k, query_records=vectors)
392
        logging.getLogger().info(result)
393
        assert len(result) == nq
394
        for i in range(nq):
395
            assert result[i][0].id == i
396
397
    @pytest.mark.timeout(ADD_TIMEOUT)
398
    def test_insert_twice_ids_no_ids(self, connect, collection):
399
        '''
400
        target: check the result of insert, with params ids and no ids
401
        method: test add vectors twice, use customize ids first, and then use no ids
402
        expected: status not OK 
403
        '''
404
        nq = 5
405
        top_k = 1
406
        vectors = gen_vectors(nq, dim)
407
        ids = [i for i in range(nq)]
408
        status, ids = connect.insert(collection, vectors, ids)
409
        assert status.OK()
410
        status, ids = connect.insert(collection, vectors)
411
        logging.getLogger().info(status)
412
        logging.getLogger().info(ids)
413
        assert not status.OK()
414
415 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
416
    def test_insert_twice_not_ids_ids(self, connect, collection):
417
        '''
418
        target: check the result of insert, with params ids and no ids
419
        method: test add vectors twice, use not ids first, and then use customize ids
420
        expected: status not OK 
421
        '''
422
        nq = 5
423
        top_k = 1
424
        vectors = gen_vectors(nq, dim)
425
        ids = [i for i in range(nq)]
426
        status, ids = connect.insert(collection, vectors)
427
        assert status.OK()
428
        status, ids = connect.insert(collection, vectors, ids)
429
        logging.getLogger().info(status)
430
        logging.getLogger().info(ids)
431
        assert not status.OK()
432
433 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
434
    def test_insert_ids_length_not_match(self, connect, collection):
435
        '''
436
        target: test add vectors in collection, use customize ids, len(ids) != len(vectors)
437
        method: create collection and add vectors in it
438
        expected: raise an exception
439
        '''
440
        nq = 5
441
        vectors = gen_vectors(nq, dim)
442
        ids = [i for i in range(1, nq)]
443
        with pytest.raises(Exception) as e:
444
            status, ids = connect.insert(collection, vectors, ids)
445
446
    @pytest.fixture(
447
        scope="function",
448
        params=gen_invalid_vector_ids()
449
    )
450
    def get_vector_id(self, request):
451
        yield request.param
452
453
    @pytest.mark.level(2)
454
    def test_insert_ids_invalid(self, connect, collection, get_vector_id):
455
        '''
456
        target: test add vectors in collection, use customize ids, which are not int64
457
        method: create collection and add vectors in it
458
        expected: raise an exception
459
        '''
460
        nq = 5
461
        vectors = gen_vectors(nq, dim)
462
        vector_id = get_vector_id
463
        ids = [vector_id for _ in range(nq)]
464
        with pytest.raises(Exception):
465
            connect.insert(collection, vectors, ids)
466
467
    @pytest.mark.timeout(ADD_TIMEOUT)
468
    def test_insert(self, connect, collection):
469
        '''
470
        target: test add vectors in collection created before
471
        method: create collection and add vectors in it, check the ids returned and the collection length after vectors added
472
        expected: the length of ids and the collection row count
473
        '''
474
        nq = 5
475
        vectors = gen_vectors(nq, dim)
476
        status, ids = connect.insert(collection, vectors)
477
        assert status.OK()
478
        assert len(ids) == nq
479
480 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
481
    def test_insert_tag(self, connect, collection):
482
        '''
483
        target: test add vectors in collection created before
484
        method: create collection and add vectors in it, with the partition_tag param
485
        expected: the collection row count equals to nq
486
        '''
487
        nq = 5
488
        vectors = gen_vectors(nq, dim)
489
        status = connect.create_partition(collection, tag)
490
        status, ids = connect.insert(collection, vectors, partition_tag=tag)
491
        assert status.OK()
492
        assert len(ids) == nq
493
494 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
495
    def test_insert_tag_A(self, connect, collection):
496
        '''
497
        target: test add vectors in collection created before
498
        method: create partition and add vectors in it
499
        expected: the collection row count equals to nq
500
        '''
501
        nq = 5
502
        vectors = gen_vectors(nq, dim)
503
        status = connect.create_partition(collection, tag)
504
        status, ids = connect.insert(collection, vectors, partition_tag=tag)
505
        assert status.OK()
506
        assert len(ids) == nq
507
508
    @pytest.mark.timeout(ADD_TIMEOUT)
509
    def test_insert_tag_not_existed(self, connect, collection):
510
        '''
511
        target: test add vectors in collection created before
512
        method: create collection and add vectors in it, with the not existed partition_tag param
513
        expected: status not ok
514
        '''
515
        nq = 5
516
        vectors = gen_vectors(nq, dim)
517
        status, ids = connect.insert(collection, vectors, partition_tag=tag)
518
        assert not status.OK()
519
520
    @pytest.mark.timeout(ADD_TIMEOUT)
521
    def test_insert_tag_not_existed_A(self, connect, collection):
522
        '''
523
        target: test add vectors in collection created before
524
        method: create partition, add vectors with the not existed partition_tag param
525
        expected: status not ok
526
        '''
527
        nq = 5
528
        vectors = gen_vectors(nq, dim)
529
        new_tag = "new_tag"
530
        status = connect.create_partition(collection, tag)
531
        status, ids = connect.insert(
532
            collection, vectors, partition_tag=new_tag)
533
        assert not status.OK()
534
535
    @pytest.mark.timeout(ADD_TIMEOUT)
536
    def test_insert_tag_existed(self, connect, collection):
537
        '''
538
        target: test add vectors in collection created before
539
        method: create collection and add vectors in it repeatly, with the partition_tag param
540
        expected: the collection row count equals to nq
541
        '''
542
        nq = 5
543
        vectors = gen_vectors(nq, dim)
544
        status = connect.create_partition(collection, tag)
545
        status, ids = connect.insert(collection, vectors, partition_tag=tag)
546
        for i in range(5):
547
            status, ids = connect.insert(
548
                collection, vectors, partition_tag=tag)
549
            assert status.OK()
550
            assert len(ids) == nq
551
552
    @pytest.mark.level(2)
553
    def test_insert_without_connect(self, dis_connect, collection):
554
        '''
555
        target: test add vectors without connection
556
        method: create collection and add vectors in it, check if added successfully
557
        expected: raise exception
558
        '''
559
        nq = 5
560
        vectors = gen_vectors(nq, dim)
561
        with pytest.raises(Exception) as e:
562
            status, ids = dis_connect.insert(collection, vectors)
563
564
    def test_add_collection_not_existed(self, connect):
565
        '''
566
        target: test add vectors in collection, which not existed before
567
        method: add vectors collection not existed, check the status
568
        expected: status not ok
569
        '''
570
        nq = 5
571
        vector = gen_single_vector(dim)
572
        status, ids = connect.insert(
573
            gen_unique_str("not_exist_collection"), vector)
574
        assert not status.OK()
575
        assert not ids
576
577
    def test_add_vector_dim_not_matched(self, connect, collection):
578
        '''
579
        target: test add vector, the vector dimension is not equal to the collection dimension
580
        method: the vector dimension is half of the collection dimension, check the status
581
        expected: status not ok
582
        '''
583
        vector = gen_single_vector(int(dim)//2)
584
        status, ids = connect.insert(collection, vector)
585
        assert not status.OK()
586
587
    def test_insert_dim_not_matched(self, connect, collection):
588
        '''
589
        target: test add vectors, the vector dimension is not equal to the collection dimension
590
        method: the vectors dimension is half of the collection dimension, check the status
591
        expected: status not ok
592
        '''
593
        nq = 5
594
        vectors = gen_vectors(nq, int(dim)//2)
595
        status, ids = connect.insert(collection, vectors)
596
        assert not status.OK()
597
598
    def test_add_vector_query_after_sleep(self, connect, collection):
599
        '''
600
        target: test add vectors, and search it after sleep
601
        method: set vector[0][1] as query vectors
602
        expected: status ok and result length is 1
603
        '''
604
        nq = 5
605
        vectors = gen_vectors(nq, dim)
606
        status, ids = connect.insert(collection, vectors)
607
        connect.flush([collection])
608
        status, result = connect.search(collection, 1, [vectors[0]])
609
        assert status.OK()
610
        assert len(result) == 1
611
612
    @pytest.mark.level(2)
613
    @pytest.mark.timeout(30)
614
    def test_collection_add_rows_count_multi_threading(self, args, collection):
615
        '''
616
        target: test collection rows_count is correct or not with multi threading
617
        method: create collection and add vectors in it(idmap),
618
            assert the value returned by count_entities method is equal to length of vectors
619
        expected: the count is equal to the length of vectors
620
        '''
621
        if args["handler"] == "HTTP":
622
            pytest.skip("Skip test in http mode")
623
        thread_num = 8
624
        threads = []
625
        milvus = get_milvus(
626
            host=args["ip"], port=args["port"], handler=args["handler"], try_connect=False)
627
        vectors = gen_vectors(nb, dim)
628
629
        def add(thread_i):
630
            logging.getLogger().info("In thread-%d" % thread_i)
631
            # milvus = get_milvus(host=args["ip"], port=args["port"], handler=args["handler"])
632
            # assert milvus
633
            status, result = milvus.insert(collection, records=vectors)
634
            assert status.OK()
635
            status = milvus.flush([collection])
636
            assert status.OK()
637
        for i in range(thread_num):
638
            x = threading.Thread(target=add, args=(i, ))
639
            threads.append(x)
640
            x.start()
641
        for th in threads:
642
            th.join()
643
        status, res = milvus.count_entities(collection)
644
        assert res == thread_num * nb
645
646 View Code Duplication
    def test_add_vector_multi_collections(self, connect):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
647
        '''
648
        target: test add vectors is correct or not with multiple collections of L2
649
        method: create 50 collections and add vectors into them in turn
650
        expected: status ok
651
        '''
652
        nq = 100
653
        vectors = gen_vectors(nq, dim)
654
        collection_list = []
655
        for i in range(20):
656
            collection_name = gen_unique_str(
657
                'test_add_vector_multi_collections')
658
            collection_list.append(collection_name)
659
            param = {'collection_name': collection_name,
660
                     'dimension': dim,
661
                     'index_file_size': index_file_size,
662
                     'metric_type': MetricType.L2}
663
            connect.create_collection(param)
664
        for j in range(5):
665
            for i in range(20):
666
                status, ids = connect.insert(
667
                    collection_name=collection_list[i], records=vectors)
668
                assert status.OK()
669
670
671
class TestAddAsync:
672
    @pytest.fixture(scope="function", autouse=True)
673
    def skip_http_check(self, args):
674
        if args["handler"] == "HTTP":
675
            pytest.skip("skip in http mode")
676
677
    @pytest.fixture(
678
        scope="function",
679
        params=[
680
            1,
681
            1000
682
        ],
683
    )
684
    def insert_count(self, request):
685
        yield request.param
686
687
    def check_status(self, status, result):
688
        logging.getLogger().info("In callback check status")
689
        assert status.OK()
690
691
    def check_status_not_ok(self, status, result):
692
        logging.getLogger().info("In callback check status")
693
        assert not status.OK()
694
695
    def test_insert_async(self, connect, collection, insert_count):
696
        '''
697
        target: test add vectors with different length of vectors
698
        method: set different vectors as add method params
699
        expected: length of ids is equal to the length of vectors
700
        '''
701
        nb = insert_count
702
        insert_vec_list = gen_vectors(nb, dim)
703
        future = connect.insert(collection, insert_vec_list, _async=True)
704
        status, ids = future.result()
705
        connect.flush([collection])
706
        assert len(ids) == nb
707
        assert status.OK()
708
709
    @pytest.mark.level(2)
710
    def test_insert_async_false(self, connect, collection, insert_count):
711
        '''
712
        target: test add vectors with different length of vectors
713
        method: set different vectors as add method params
714
        expected: length of ids is equal to the length of vectors
715
        '''
716
        nb = insert_count
717
        insert_vec_list = gen_vectors(nb, dim)
718
        status, ids = connect.insert(collection, insert_vec_list, _async=False)
719
        connect.flush([collection])
720
        assert len(ids) == nb
721
        assert status.OK()
722
723
    def test_insert_async_callback(self, connect, collection, insert_count):
724
        '''
725
        target: test add vectors with different length of vectors
726
        method: set different vectors as add method params
727
        expected: length of ids is equal to the length of vectors
728
        '''
729
        nb = insert_count
730
        insert_vec_list = gen_vectors(nb, dim)
731
        future = connect.insert(
732
            collection, insert_vec_list, _async=True, _callback=self.check_status)
733
        future.done()
734
735
    @pytest.mark.level(2)
736
    def test_insert_async_long(self, connect, collection):
737
        '''
738
        target: test add vectors with different length of vectors
739
        method: set different vectors as add method params
740
        expected: length of ids is equal to the length of vectors
741
        '''
742
        nb = 50000
743
        insert_vec_list = gen_vectors(nb, dim)
744
        future = connect.insert(
745
            collection, insert_vec_list, _async=True, _callback=self.check_status)
746
        status, result = future.result()
747
        assert status.OK()
748
        assert len(result) == nb
749
        connect.flush([collection])
750
        status, count = connect.count_entities(collection)
751
        assert status.OK()
752
        logging.getLogger().info(status)
753
        logging.getLogger().info(count)
754
        assert count == nb
755
756
    def test_insert_async_callback_timeout(self, connect, collection):
757
        '''
758
        target: test add vectors with different length of vectors
759
        method: set different vectors as add method params
760
        expected: length of ids is equal to the length of vectors
761
        '''
762
        nb = 100000
763
        insert_vec_list = gen_vectors(nb, dim)
764
        future = connect.insert(collection, insert_vec_list,
765
                                _async=True, _callback=self.check_status, timeout=1)
766
        future.done()
767
768
    def test_insert_async_invalid_params(self, connect, collection):
769
        '''
770
        target: test add vectors with different length of vectors
771
        method: set different vectors as add method params
772
        expected: length of ids is equal to the length of vectors
773
        '''
774
        insert_vec_list = gen_vectors(nb, dim)
775
        collection_new = gen_unique_str()
776
        future = connect.insert(collection_new, insert_vec_list, _async=True)
777
        status, result = future.result()
778
        assert not status.OK()
779
780
    # TODO: add assertion
781
    def test_insert_async_invalid_params_raise_exception(self, connect, collection):
782
        '''
783
        target: test add vectors with different length of vectors
784
        method: set different vectors as add method params
785
        expected: length of ids is equal to the length of vectors
786
        '''
787
        insert_vec_list = []
788
        collection_new = gen_unique_str()
789
        with pytest.raises(Exception) as e:
790
            future = connect.insert(
791
                collection_new, insert_vec_list, _async=True)
792
793
794
795
class TestAddIP:
796
    """
797
    ******************************************************************
798
      The following cases are used to test `insert / index / search / delete` mixed function
799
    ******************************************************************
800
    """
801 View Code Duplication
    @pytest.fixture(
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
802
        scope="function",
803
        params=gen_simple_index()
804
    )
805
    def get_simple_index(self, request, connect):
806
        if str(connect._cmd("mode")[1]) == "CPU":
807
            if request.param["index_type"] == IndexType.IVF_SQ8H:
808
                pytest.skip("sq8h not support in cpu mode")
809
        if request.param["index_type"] == IndexType.IVF_PQ:
810
            pytest.skip("Skip PQ Temporary")
811
        return request.param
812
813
    def test_add_vector_create_collection(self, connect, ip_collection):
814
        '''
815
        target: test add vector, then create collection again
816
        method: add vector and create collection
817
        expected: status not ok
818
        '''
819
        vector = gen_single_vector(dim)
820
        status, ids = connect.insert(ip_collection, vector)
821
        param = {'collection_name': ip_collection,
822
                 'dimension': dim,
823
                 'index_file_size': index_file_size,
824
                 'metric_type': MetricType.L2}
825
        status = connect.create_collection(param)
826
        assert not status.OK()
827
828
    def test_add_vector_has_collection(self, connect, ip_collection):
829
        '''
830
        target: test add vector, then check collection existence
831
        method: add vector and call Hascollection
832
        expected: collection exists, status ok
833
        '''
834
        vector = gen_single_vector(dim)
835
        status, ids = connect.insert(ip_collection, vector)
836
        assert assert_has_collection(connect, ip_collection)
837
838
    @pytest.mark.timeout(ADD_TIMEOUT)
839
    def test_drop_collection_add_vector(self, connect, ip_collection):
840
        '''
841
        target: test add vector after collection deleted
842
        method: delete collection and add vector
843
        expected: status not ok
844
        '''
845
        status = connect.drop_collection(ip_collection)
846
        vector = gen_single_vector(dim)
847
        status, ids = connect.insert(ip_collection, vector)
848
        assert not status.OK()
849
850 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
851
    def test_drop_collection_add_vector_another(self, connect, ip_collection):
852
        '''
853
        target: test add vector to collection_1 after collection_2 deleted
854
        method: delete collection_2 and add vector to collection_1
855
        expected: status ok
856
        '''
857
        param = {'collection_name': gen_unique_str(),
858
                 'dimension': dim,
859
                 'index_file_size': index_file_size,
860
                 'metric_type': MetricType.L2}
861
        status = connect.create_collection(param)
862
        status = connect.drop_collection(ip_collection)
863
        vector = gen_single_vector(dim)
864
        status, ids = connect.insert(param['collection_name'], vector)
865
        assert status.OK()
866
867
    @pytest.mark.timeout(ADD_TIMEOUT)
868
    def test_add_vector_drop_collection(self, connect, ip_collection):
869
        '''
870
        target: test delete collection after add vector
871
        method: add vector and delete collection
872
        expected: status ok
873
        '''
874
        vector = gen_single_vector(dim)
875
        status, ids = connect.insert(ip_collection, vector)
876
        status = connect.drop_collection(ip_collection)
877
        assert status.OK()
878
879
    @pytest.mark.timeout(ADD_TIMEOUT)
880
    def test_add_vector_delete_another_collection(self, connect, ip_collection):
881
        '''
882
        target: test delete collection_1 collection after add vector to collection_2
883
        method: add vector and delete collection
884
        expected: status ok
885
        '''
886
        param = {'collection_name': 'test_add_vector_delete_another_collection',
887
                 'dimension': dim,
888
                 'index_file_size': index_file_size,
889
                 'metric_type': MetricType.L2}
890
        status = connect.create_collection(param)
891
        vector = gen_single_vector(dim)
892
        status, ids = connect.insert(ip_collection, vector)
893
        status = connect.drop_collection(param['collection_name'])
894
        assert status.OK()
895
896
    @pytest.mark.timeout(ADD_TIMEOUT)
897
    def test_add_vector_sleep_drop_collection(self, connect, ip_collection):
898
        '''
899
        target: test delete collection after add vector for a while
900
        method: add vector, sleep, and delete collection
901
        expected: status ok
902
        '''
903
        vector = gen_single_vector(dim)
904
        status, ids = connect.insert(ip_collection, vector)
905
        connect.flush([ip_collection])
906
        status = connect.drop_collection(ip_collection)
907
        assert status.OK()
908
909
    @pytest.mark.timeout(ADD_TIMEOUT)
910
    def test_add_vector_sleep_delete_another_collection(self, connect, ip_collection):
911
        '''
912
        target: test delete collection_1 collection after add vector to collection_2 for a while
913
        method: add vector , sleep, and delete collection
914
        expected: status ok
915
        '''
916
        param = {'collection_name': gen_unique_str(),
917
                 'dimension': dim,
918
                 'index_file_size': index_file_size,
919
                 'metric_type': MetricType.L2}
920
        status = connect.create_collection(param)
921
        vector = gen_single_vector(dim)
922
        status, ids = connect.insert(ip_collection, vector)
923
        connect.flush([ip_collection])
924
        status = connect.drop_collection(param['collection_name'])
925
        assert status.OK()
926
927 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
928
    def test_create_index_add_vector(self, connect, ip_collection, get_simple_index):
929
        '''
930
        target: test add vector after build index
931
        method: build index and add vector
932
        expected: status ok
933
        '''
934
        index_param = get_simple_index["index_param"]
935
        index_type = get_simple_index["index_type"]
936
        status = connect.create_index(ip_collection, index_type, index_param)
937
        vector = gen_single_vector(dim)
938
        status, ids = connect.insert(ip_collection, vector)
939
        assert status.OK()
940
941
    @pytest.mark.timeout(ADD_TIMEOUT)
942
    def test_create_index_add_vector_another(self, connect, ip_collection, get_simple_index):
943
        '''
944
        target: test add vector to collection_2 after build index for collection_1
945
        method: build index and add vector
946
        expected: status ok
947
        '''
948
        index_param = get_simple_index["index_param"]
949
        index_type = get_simple_index["index_type"]
950
        param = {'collection_name': gen_unique_str(),
951
                 'dimension': dim,
952
                 'index_file_size': index_file_size,
953
                 'metric_type': MetricType.L2}
954
        status = connect.create_collection(param)
955
        status = connect.create_index(ip_collection, index_type, index_param)
956
        vector = gen_single_vector(dim)
957
        status, ids = connect.insert(ip_collection, vector)
958
        assert status.OK()
959
960
    @pytest.mark.timeout(ADD_TIMEOUT)
961
    def test_add_vector_create_index(self, connect, ip_collection, get_simple_index):
962
        '''
963
        target: test build index add after vector
964
        method: add vector and build index
965
        expected: status ok
966
        '''
967
        index_param = get_simple_index["index_param"]
968
        index_type = get_simple_index["index_type"]
969
        vector = gen_single_vector(dim)
970
        status, ids = connect.insert(ip_collection, vector)
971
        status, mode = connect._cmd("mode")
972
        assert status.OK()
973
        status = connect.create_index(ip_collection, index_type, index_param)
974
        if str(mode) == "GPU" and (index_type == IndexType.IVF_PQ):
975
            assert not status.OK()
976
        else:
977
            assert status.OK()
978
979 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
980
    def test_add_vector_create_index_another(self, connect, ip_collection, get_simple_index):
981
        '''
982
        target: test add vector to collection_2 after build index for collection_1
983
        method: build index and add vector
984
        expected: status ok
985
        '''
986
        index_param = get_simple_index["index_param"]
987
        index_type = get_simple_index["index_type"]
988
        param = {'collection_name': gen_unique_str(),
989
                 'dimension': dim,
990
                 'index_file_size': index_file_size,
991
                 'metric_type': MetricType.L2}
992
        status = connect.create_collection(param)
993
        vector = gen_single_vector(dim)
994
        status, ids = connect.insert(ip_collection, vector)
995
        status = connect.create_index(
996
            param['collection_name'], index_type, index_param)
997
        assert status.OK()
998
999
    @pytest.mark.timeout(ADD_TIMEOUT)
1000
    def test_add_vector_sleep_create_index(self, connect, ip_collection, get_simple_index):
1001
        '''
1002
        target: test build index add after vector for a while
1003
        method: add vector and build index
1004
        expected: status ok
1005
        '''
1006
        index_param = get_simple_index["index_param"]
1007
        index_type = get_simple_index["index_type"]
1008
        if index_type == IndexType.IVF_PQ:
1009
            pytest.skip("Skip some PQ cases")
1010
        vector = gen_single_vector(dim)
1011
        status, ids = connect.insert(ip_collection, vector)
1012
        assert status.OK()
1013
        time.sleep(add_interval_time)
1014
        status = connect.create_index(ip_collection, index_type, index_param)
1015
        assert status.OK()
1016
1017 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1018
    def test_add_vector_sleep_create_index_another(self, connect, ip_collection, get_simple_index):
1019
        '''
1020
        target: test add vector to collection_2 after build index for collection_1 for a while
1021
        method: build index and add vector
1022
        expected: status ok
1023
        '''
1024
        index_param = get_simple_index["index_param"]
1025
        index_type = get_simple_index["index_type"]
1026
        param = {'collection_name': gen_unique_str(),
1027
                 'dimension': dim,
1028
                 'index_file_size': index_file_size,
1029
                 'metric_type': MetricType.L2}
1030
        status = connect.create_collection(param)
1031
        vector = gen_single_vector(dim)
1032
        status, ids = connect.insert(ip_collection, vector)
1033
        connect.flush([ip_collection])
1034
        status = connect.create_index(
1035
            param['collection_name'], index_type, index_param)
1036
        assert status.OK()
1037
1038
    @pytest.mark.timeout(ADD_TIMEOUT)
1039
    def test_search_vector_add_vector(self, connect, ip_collection):
1040
        '''
1041
        target: test add vector after search collection
1042
        method: search collection and add vector
1043
        expected: status ok
1044
        '''
1045
        vector = gen_single_vector(dim)
1046
        status, result = connect.search(ip_collection, 1, vector)
1047
        status, ids = connect.insert(ip_collection, vector)
1048
        assert status.OK()
1049
1050 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1051
    def test_search_vector_add_vector_another(self, connect, ip_collection):
1052
        '''
1053
        target: test add vector to collection_1 after search collection_2
1054
        method: search collection and add vector
1055
        expected: status ok
1056
        '''
1057
        param = {'collection_name': gen_unique_str(),
1058
                 'dimension': dim,
1059
                 'index_file_size': index_file_size,
1060
                 'metric_type': MetricType.L2}
1061
        status = connect.create_collection(param)
1062
        vector = gen_single_vector(dim)
1063
        status, result = connect.search(ip_collection, 1, vector)
1064
        status, ids = connect.insert(param['collection_name'], vector)
1065
        assert status.OK()
1066
1067
    @pytest.mark.timeout(ADD_TIMEOUT)
1068
    def test_add_vector_search_vector(self, connect, ip_collection):
1069
        '''
1070
        target: test search vector after add vector
1071
        method: add vector and search collection
1072
        expected: status ok
1073
        '''
1074
        vector = gen_single_vector(dim)
1075
        status, ids = connect.insert(ip_collection, vector)
1076
        assert status.OK()
1077
        connect.flush([ip_collection])
1078
        status, result = connect.search(ip_collection, 1, vector)
1079
        assert status.OK()
1080
1081
    @pytest.mark.timeout(ADD_TIMEOUT)
1082
    def test_add_vector_search_vector_another(self, connect, ip_collection):
1083
        '''
1084
        target: test add vector to collection_1 after search collection_2
1085
        method: search collection and add vector
1086
        expected: status ok
1087
        '''
1088
        param = {'collection_name': gen_unique_str(),
1089
                 'dimension': dim,
1090
                 'index_file_size': index_file_size,
1091
                 'metric_type': MetricType.L2}
1092
        status = connect.create_collection(param)
1093
        vector = gen_single_vector(dim)
1094
        status, ids = connect.insert(ip_collection, vector)
1095
        connect.flush([ip_collection])
1096
        status, result = connect.search(param['collection_name'], 1, vector)
1097
        assert status.OK()
1098
1099
    @pytest.mark.timeout(ADD_TIMEOUT)
1100
    def test_add_vector_sleep_search_vector(self, connect, ip_collection):
1101
        '''
1102
        target: test search vector after add vector after a while
1103
        method: add vector, sleep, and search collection
1104
        expected: status ok
1105
        '''
1106
        vector = gen_single_vector(dim)
1107
        status, ids = connect.insert(ip_collection, vector)
1108
        time.sleep(add_interval_time)
1109
        status, result = connect.search(ip_collection, 1, vector)
1110
        assert status.OK()
1111
1112
    @pytest.mark.timeout(ADD_TIMEOUT)
1113
    def test_add_vector_sleep_search_vector_another(self, connect, ip_collection):
1114
        '''
1115
        target: test add vector to collection_1 after search collection_2 a while
1116
        method: search collection , sleep, and add vector
1117
        expected: status ok
1118
        '''
1119
        param = {'collection_name': gen_unique_str(),
1120
                 'dimension': dim,
1121
                 'index_file_size': index_file_size,
1122
                 'metric_type': MetricType.L2}
1123
        status = connect.create_collection(param)
1124
        vector = gen_single_vector(dim)
1125
        status, ids = connect.insert(ip_collection, vector)
1126
        assert status.OK()
1127
        time.sleep(add_interval_time)
1128
        status, result = connect.search(param['collection_name'], 1, vector)
1129
        assert status.OK()
1130
1131
    """
1132
    ******************************************************************
1133
      The following cases are used to test `insert` function
1134
    ******************************************************************
1135
    """
1136
1137 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1138
    def test_insert_ids(self, connect, ip_collection):
1139
        '''
1140
        target: test add vectors in collection, use customize ids
1141
        method: create collection and add vectors in it, check the ids returned and the collection length after vectors added
1142
        expected: the length of ids and the collection row count
1143
        '''
1144
        nq = 5
1145
        top_k = 1
1146
        vectors = gen_vectors(nq, dim)
1147
        ids = [i for i in range(nq)]
1148
        status, ids = connect.insert(ip_collection, vectors, ids)
1149
        assert status.OK()
1150
        connect.flush([ip_collection])
1151
        assert len(ids) == nq
1152
        # check search result
1153
        status, result = connect.search(ip_collection, top_k, vectors)
1154
        logging.getLogger().info(result)
1155
        assert len(result) == nq
1156
        for i in range(nq):
1157
            assert result[i][0].id == i
1158
1159 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1160
    def test_insert_twice_ids_no_ids(self, connect, ip_collection):
1161
        '''
1162
        target: check the result of insert, with params ids and no ids
1163
        method: test add vectors twice, use customize ids first, and then use no ids
1164
        expected: status not OK 
1165
        '''
1166
        nq = 5
1167
        top_k = 1
1168
        vectors = gen_vectors(nq, dim)
1169
        ids = [i for i in range(nq)]
1170
        status, ids = connect.insert(ip_collection, vectors, ids)
1171
        assert status.OK()
1172
        status, ids = connect.insert(ip_collection, vectors)
1173
        logging.getLogger().info(status)
1174
        logging.getLogger().info(ids)
1175
        assert not status.OK()
1176
1177 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1178
    def test_insert_twice_not_ids_ids(self, connect, ip_collection):
1179
        '''
1180
        target: check the result of insert, with params ids and no ids
1181
        method: test add vectors twice, use not ids first, and then use customize ids
1182
        expected: status not OK 
1183
        '''
1184
        nq = 5
1185
        top_k = 1
1186
        vectors = gen_vectors(nq, dim)
1187
        ids = [i for i in range(nq)]
1188
        status, ids = connect.insert(ip_collection, vectors)
1189
        assert status.OK()
1190
        status, ids = connect.insert(ip_collection, vectors, ids)
1191
        logging.getLogger().info(status)
1192
        logging.getLogger().info(ids)
1193
        assert not status.OK()
1194
1195 View Code Duplication
    @pytest.mark.timeout(ADD_TIMEOUT)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1196
    def test_insert_ids_length_not_match(self, connect, ip_collection):
1197
        '''
1198
        target: test add vectors in collection, use customize ids, len(ids) != len(vectors)
1199
        method: create collection and add vectors in it
1200
        expected: raise an exception
1201
        '''
1202
        nq = 5
1203
        vectors = gen_vectors(nq, dim)
1204
        ids = [i for i in range(1, nq)]
1205
        with pytest.raises(Exception) as e:
1206
            status, ids = connect.insert(ip_collection, vectors, ids)
1207
1208
    @pytest.fixture(
1209
        scope="function",
1210
        params=gen_invalid_vector_ids()
1211
    )
1212
    def get_vector_id(self, request):
1213
        yield request.param
1214
1215
    @pytest.mark.level(2)
1216
    def test_insert_ids_invalid(self, connect, ip_collection, get_vector_id):
1217
        '''
1218
        target: test add vectors in collection, use customize ids, which are not int64
1219
        method: create collection and add vectors in it
1220
        expected: raise an exception
1221
        '''
1222
        nq = 5
1223
        vectors = gen_vectors(nq, dim)
1224
        vector_id = get_vector_id
1225
        ids = [vector_id for i in range(nq)]
1226
        with pytest.raises(Exception) as e:
1227
            status, ids = connect.insert(ip_collection, vectors, ids)
1228
1229
    @pytest.mark.timeout(ADD_TIMEOUT)
1230
    def test_insert(self, connect, ip_collection):
1231
        '''
1232
        target: test add vectors in collection created before
1233
        method: create collection and add vectors in it, check the ids returned and the collection length after vectors added
1234
        expected: the length of ids and the collection row count
1235
        '''
1236
        nq = 5
1237
        vectors = gen_vectors(nq, dim)
1238
        status, ids = connect.insert(ip_collection, vectors)
1239
        assert status.OK()
1240
        assert len(ids) == nq
1241
1242
    # @pytest.mark.level(2)
1243
    # def test_insert_without_connect(self, dis_connect, ip_collection):
1244
    #     '''
1245
    #     target: test add vectors without connection
1246
    #     method: create collection and add vectors in it, check if added successfully
1247
    #     expected: raise exception
1248
    #     '''
1249
    #     nq = 5
1250
    #     vectors = gen_vectors(nq, dim)
1251
    #     with pytest.raises(Exception) as e:
1252
    #         status, ids = dis_connect.insert(ip_collection, vectors)
1253
1254
    def test_add_vector_dim_not_matched(self, connect, ip_collection):
1255
        '''
1256
        target: test add vector, the vector dimension is not equal to the collection dimension
1257
        method: the vector dimension is half of the collection dimension, check the status
1258
        expected: status not ok
1259
        '''
1260
        vector = gen_single_vector(int(dim)//2)
1261
        status, ids = connect.insert(ip_collection, vector)
1262
        assert not status.OK()
1263
1264
    def test_insert_dim_not_matched(self, connect, ip_collection):
1265
        '''
1266
        target: test add vectors, the vector dimension is not equal to the collection dimension
1267
        method: the vectors dimension is half of the collection dimension, check the status
1268
        expected: status not ok
1269
        '''
1270
        nq = 5
1271
        vectors = gen_vectors(nq, int(dim)//2)
1272
        status, ids = connect.insert(ip_collection, vectors)
1273
        assert not status.OK()
1274
1275
    def test_add_vector_query_after_sleep(self, connect, ip_collection):
1276
        '''
1277
        target: test add vectors, and search it after sleep
1278
        method: set vector[0][1] as query vectors
1279
        expected: status ok and result length is 1
1280
        '''
1281
        nq = 5
1282
        vectors = gen_vectors(nq, dim)
1283
        status, ids = connect.insert(ip_collection, vectors)
1284
        time.sleep(add_interval_time)
1285
        status, result = connect.search(ip_collection, 1, [vectors[0]])
1286
        assert status.OK()
1287
        assert len(result) == 1
1288
1289 View Code Duplication
    def test_add_vector_multi_collections(self, connect):
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
1290
        '''
1291
        target: test add vectors is correct or not with multiple collections of IP
1292
        method: create 50 collections and add vectors into them in turn
1293
        expected: status ok
1294
        '''
1295
        nq = 100
1296
        vectors = gen_vectors(nq, dim)
1297
        collection_list = []
1298
        for i in range(20):
1299
            collection_name = gen_unique_str(
1300
                'test_add_vector_multi_collections')
1301
            collection_list.append(collection_name)
1302
            param = {'collection_name': collection_name,
1303
                     'dimension': dim,
1304
                     'index_file_size': index_file_size,
1305
                     'metric_type': MetricType.IP}
1306
            connect.create_collection(param)
1307
        for j in range(10):
1308
            for i in range(20):
1309
                status, ids = connect.insert(
1310
                    collection_name=collection_list[i], records=vectors)
1311
                assert status.OK()
1312
1313
1314
class TestAddAdvance:
1315
    @pytest.fixture(
1316
        scope="function",
1317
        params=[
1318
            1,
1319
            1000,
1320
            6000
1321
        ],
1322
    )
1323
    def insert_count(self, request):
1324
        yield request.param
1325
1326
    def test_insert_much(self, connect, collection, insert_count):
1327
        '''
1328
        target: test add vectors with different length of vectors
1329
        method: set different vectors as add method params
1330
        expected: length of ids is equal to the length of vectors
1331
        '''
1332
        nb = insert_count
1333
        insert_vec_list = gen_vectors(nb, dim)
1334
        status, ids = connect.insert(collection, insert_vec_list)
1335
        assert len(ids) == nb
1336
        assert status.OK()
1337
1338
    def test_insert_much_ip(self, connect, ip_collection, insert_count):
1339
        '''
1340
        target: test add vectors with different length of vectors
1341
        method: set different vectors as add method params
1342
        expected: length of ids is equal to the length of vectors
1343
        '''
1344
        nb = insert_count
1345
        insert_vec_list = gen_vectors(nb, dim)
1346
        status, ids = connect.insert(ip_collection, insert_vec_list)
1347
        assert len(ids) == nb
1348
        assert status.OK()
1349
1350
    def test_insert_much_jaccard(self, connect, jac_collection, insert_count):
1351
        '''
1352
        target: test add vectors with different length of vectors
1353
        method: set different vectors as add method params
1354
        expected: length of ids is equal to the length of vectors
1355
        '''
1356
        nb = insert_count
1357
        tmp, insert_vec_list = gen_binary_vectors(nb, dim)
1358
        status, ids = connect.insert(jac_collection, insert_vec_list)
1359
        assert len(ids) == nb
1360
        assert status.OK()
1361
1362
    def test_insert_much_hamming(self, connect, ham_collection, insert_count):
1363
        '''
1364
        target: test add vectors with different length of vectors
1365
        method: set different vectors as add method params
1366
        expected: length of ids is equal to the length of vectors
1367
        '''
1368
        nb = insert_count
1369
        tmp, insert_vec_list = gen_binary_vectors(nb, dim)
1370
        status, ids = connect.insert(ham_collection, insert_vec_list)
1371
        assert len(ids) == nb
1372
        assert status.OK()
1373
1374
    def test_insert_much_tanimoto(self, connect, tanimoto_collection, insert_count):
1375
        '''
1376
        target: test add vectors with different length of vectors
1377
        method: set different vectors as add method params
1378
        expected: length of ids is equal to the length of vectors
1379
        '''
1380
        nb = insert_count
1381
        tmp, insert_vec_list = gen_binary_vectors(nb, dim)
1382
        status, ids = connect.insert(tanimoto_collection, insert_vec_list)
1383
        assert len(ids) == nb
1384
        assert status.OK()
1385
1386
1387
class TestNameInvalid(object):
1388
    """
1389
    Test adding vectors with invalid collection names
1390
    """
1391
    @pytest.fixture(
1392
        scope="function",
1393
        params=gen_invalid_collection_names()
1394
    )
1395
    def get_collection_name(self, request):
1396
        yield request.param
1397
1398
    @pytest.fixture(
1399
        scope="function",
1400
        params=gen_invalid_collection_names()
1401
    )
1402
    def get_tag_name(self, request):
1403
        yield request.param
1404
1405
    @pytest.mark.level(2)
1406
    def test_insert_with_invalid_collection_name(self, connect, get_collection_name):
1407
        collection_name = get_collection_name
1408
        vectors = gen_vectors(1, dim)
1409
        status, result = connect.insert(collection_name, vectors)
1410
        assert not status.OK()
1411
1412
    @pytest.mark.level(2)
1413
    def test_insert_with_invalid_tag_name(self, connect, get_collection_name, get_tag_name):
1414
        collection_name = get_collection_name
1415
        tag_name = get_tag_name
1416
        vectors = gen_vectors(1, dim)
1417
        status, result = connect.insert(
1418
            collection_name, vectors, partition_tag=tag_name)
1419
        assert not status.OK()
1420
1421
1422
class TestAddCollectionVectorsInvalid(object):
1423
    single_vector = gen_single_vector(dim)
1424
    vectors = gen_vectors(2, dim)
1425
1426
    """
1427
    Test adding vectors with invalid vectors
1428
    """
1429
    @pytest.fixture(
1430
        scope="function",
1431
        params=gen_invalid_vectors()
1432
    )
1433
    def gen_vector(self, request):
1434
        yield request.param
1435
1436
    @pytest.mark.level(2)
1437
    def test_add_vector_with_invalid_vectors(self, connect, collection, gen_vector):
1438
        tmp_single_vector = copy.deepcopy(self.single_vector)
1439
        tmp_single_vector[0][1] = gen_vector
1440
        with pytest.raises(Exception) as e:
1441
            status, result = connect.insert(collection, tmp_single_vector)
1442
1443
    @pytest.mark.level(2)
1444
    def test_insert_with_invalid_vectors(self, connect, collection, gen_vector):
1445
        tmp_vectors = copy.deepcopy(self.vectors)
1446
        tmp_vectors[1][1] = gen_vector
1447
        with pytest.raises(Exception) as e:
1448
            status, result = connect.insert(collection, tmp_vectors)
1449
1450
    @pytest.mark.level(2)
1451
    def test_insert_with_invalid_vectors_jaccard(self, connect, jac_collection, gen_vector):
1452
        tmp_vectors = copy.deepcopy(self.vectors)
1453
        tmp_vectors[1][1] = gen_vector
1454
        with pytest.raises(Exception) as e:
1455
            status, result = connect.insert(jac_collection, tmp_vectors)
1456
1457
    @pytest.mark.level(2)
1458
    def test_insert_with_invalid_vectors_hamming(self, connect, ham_collection, gen_vector):
1459
        tmp_vectors = copy.deepcopy(self.vectors)
1460
        tmp_vectors[1][1] = gen_vector
1461
        with pytest.raises(Exception) as e:
1462
            status, result = connect.insert(ham_collection, tmp_vectors)
1463