Total Complexity | 134 |
Total Lines | 1463 |
Duplicated Lines | 30.62 % |
Changes | 0 |
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:
Complex classes like test_insert 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 | 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( |
|
|
|||
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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): |
|
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( |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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) |
|
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): |
|
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 |