Total Complexity | 65 |
Total Lines | 723 |
Duplicated Lines | 35.82 % |
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_collection_count 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 pdb |
||
2 | import pytest |
||
3 | import logging |
||
4 | import itertools |
||
5 | from time import sleep |
||
6 | import threading |
||
7 | from multiprocessing import Process |
||
8 | from milvus import IndexType, MetricType |
||
9 | from utils import * |
||
10 | |||
11 | dim = 128 |
||
12 | index_file_size = 10 |
||
13 | add_time_interval = 3 |
||
14 | tag = "1970-01-01" |
||
15 | nb = 6000 |
||
16 | |||
17 | class TestCollectionCount: |
||
18 | """ |
||
19 | params means different nb, the nb value may trigger merge, or not |
||
20 | """ |
||
21 | @pytest.fixture( |
||
22 | scope="function", |
||
23 | params=[ |
||
24 | 1, |
||
25 | 5000, |
||
26 | 20000, |
||
27 | ], |
||
28 | ) |
||
29 | def insert_nb(self, request): |
||
30 | yield request.param |
||
31 | |||
32 | """ |
||
33 | generate valid create_index params |
||
34 | """ |
||
35 | View Code Duplication | @pytest.fixture( |
|
|
|||
36 | scope="function", |
||
37 | params=gen_simple_index() |
||
38 | ) |
||
39 | def get_simple_index(self, request, connect): |
||
40 | if str(connect._cmd("mode")[1]) == "CPU": |
||
41 | if request.param["index_type"] == IndexType.IVF_SQ8H: |
||
42 | pytest.skip("sq8h not support in cpu mode") |
||
43 | if request.param["index_type"] == IndexType.IVF_PQ: |
||
44 | pytest.skip("Skip PQ Temporary") |
||
45 | return request.param |
||
46 | |||
47 | def test_collection_rows_count(self, connect, collection, insert_nb): |
||
48 | ''' |
||
49 | target: test collection rows_count is correct or not |
||
50 | method: create collection and add vectors in it, |
||
51 | assert the value returned by count_entities method is equal to length of vectors |
||
52 | expected: the count is equal to the length of vectors |
||
53 | ''' |
||
54 | nb = insert_nb |
||
55 | vectors = gen_vectors(nb, dim) |
||
56 | res = connect.insert(collection_name=collection, records=vectors) |
||
57 | connect.flush([collection]) |
||
58 | status, res = connect.count_entities(collection) |
||
59 | assert res == nb |
||
60 | |||
61 | def test_collection_rows_count_partition(self, connect, collection, insert_nb): |
||
62 | ''' |
||
63 | target: test collection rows_count is correct or not |
||
64 | method: create collection, create partition and add vectors in it, |
||
65 | assert the value returned by count_entities method is equal to length of vectors |
||
66 | expected: the count is equal to the length of vectors |
||
67 | ''' |
||
68 | nb = insert_nb |
||
69 | vectors = gen_vectors(nb, dim) |
||
70 | status = connect.create_partition(collection, tag) |
||
71 | assert status.OK() |
||
72 | res = connect.insert(collection_name=collection, records=vectors, partition_tag=tag) |
||
73 | connect.flush([collection]) |
||
74 | status, res = connect.count_entities(collection) |
||
75 | assert res == nb |
||
76 | |||
77 | View Code Duplication | def test_collection_rows_count_multi_partitions_A(self, connect, collection, insert_nb): |
|
78 | ''' |
||
79 | target: test collection rows_count is correct or not |
||
80 | method: create collection, create partitions and add vectors in it, |
||
81 | assert the value returned by count_entities method is equal to length of vectors |
||
82 | expected: the count is equal to the length of vectors |
||
83 | ''' |
||
84 | new_tag = "new_tag" |
||
85 | nb = insert_nb |
||
86 | vectors = gen_vectors(nb, dim) |
||
87 | status = connect.create_partition(collection, tag) |
||
88 | status = connect.create_partition(collection, new_tag) |
||
89 | assert status.OK() |
||
90 | res = connect.insert(collection_name=collection, records=vectors) |
||
91 | connect.flush([collection]) |
||
92 | status, res = connect.count_entities(collection) |
||
93 | assert res == nb |
||
94 | |||
95 | View Code Duplication | def test_collection_rows_count_multi_partitions_B(self, connect, collection, insert_nb): |
|
96 | ''' |
||
97 | target: test collection rows_count is correct or not |
||
98 | method: create collection, create partitions and add vectors in one of the partitions, |
||
99 | assert the value returned by count_entities method is equal to length of vectors |
||
100 | expected: the count is equal to the length of vectors |
||
101 | ''' |
||
102 | new_tag = "new_tag" |
||
103 | nb = insert_nb |
||
104 | vectors = gen_vectors(nb, dim) |
||
105 | status = connect.create_partition(collection, tag) |
||
106 | status = connect.create_partition(collection, new_tag) |
||
107 | assert status.OK() |
||
108 | res = connect.insert(collection_name=collection, records=vectors, partition_tag=tag) |
||
109 | connect.flush([collection]) |
||
110 | status, res = connect.count_entities(collection) |
||
111 | assert res == nb |
||
112 | |||
113 | def test_collection_rows_count_multi_partitions_C(self, connect, collection, insert_nb): |
||
114 | ''' |
||
115 | target: test collection rows_count is correct or not |
||
116 | method: create collection, create partitions and add vectors in one of the partitions, |
||
117 | assert the value returned by count_entities method is equal to length of vectors |
||
118 | expected: the collection count is equal to the length of vectors |
||
119 | ''' |
||
120 | new_tag = "new_tag" |
||
121 | nb = insert_nb |
||
122 | vectors = gen_vectors(nb, dim) |
||
123 | status = connect.create_partition(collection, tag) |
||
124 | status = connect.create_partition(collection, new_tag) |
||
125 | assert status.OK() |
||
126 | res = connect.insert(collection_name=collection, records=vectors, partition_tag=tag) |
||
127 | res = connect.insert(collection_name=collection, records=vectors, partition_tag=new_tag) |
||
128 | connect.flush([collection]) |
||
129 | status, res = connect.count_entities(collection) |
||
130 | assert res == nb * 2 |
||
131 | |||
132 | View Code Duplication | def test_collection_rows_count_after_index_created(self, connect, collection, get_simple_index): |
|
133 | ''' |
||
134 | target: test count_entities, after index have been created |
||
135 | method: add vectors in db, and create index, then calling count_entities with correct params |
||
136 | expected: count_entities raise exception |
||
137 | ''' |
||
138 | index_param = get_simple_index["index_param"] |
||
139 | index_type = get_simple_index["index_type"] |
||
140 | nb = 100 |
||
141 | vectors = gen_vectors(nb, dim) |
||
142 | res = connect.insert(collection_name=collection, records=vectors) |
||
143 | connect.flush([collection]) |
||
144 | connect.create_index(collection, index_type, index_param) |
||
145 | status, res = connect.count_entities(collection) |
||
146 | assert res == nb |
||
147 | |||
148 | # @pytest.mark.level(2) |
||
149 | # def test_count_without_connection(self, collection, dis_connect): |
||
150 | # ''' |
||
151 | # target: test count_entities, without connection |
||
152 | # method: calling count_entities with correct params, with a disconnected instance |
||
153 | # expected: count_entities raise exception |
||
154 | # ''' |
||
155 | # with pytest.raises(Exception) as e: |
||
156 | # status = dis_connect.count_entities(collection) |
||
157 | |||
158 | def test_collection_rows_count_no_vectors(self, connect, collection): |
||
159 | ''' |
||
160 | target: test collection rows_count is correct or not, if collection is empty |
||
161 | method: create collection and no vectors in it, |
||
162 | assert the value returned by count_entities method is equal to 0 |
||
163 | expected: the count is equal to 0 |
||
164 | ''' |
||
165 | collection_name = gen_unique_str() |
||
166 | param = {'collection_name': collection_name, |
||
167 | 'dimension': dim, |
||
168 | 'index_file_size': index_file_size} |
||
169 | connect.create_collection(param) |
||
170 | status, res = connect.count_entities(collection) |
||
171 | assert res == 0 |
||
172 | |||
173 | # TODO: enable |
||
174 | View Code Duplication | @pytest.mark.level(2) |
|
175 | @pytest.mark.timeout(20) |
||
176 | def _test_collection_rows_count_multiprocessing(self, connect, collection, args): |
||
177 | ''' |
||
178 | target: test collection rows_count is correct or not with multiprocess |
||
179 | method: create collection and add vectors in it, |
||
180 | assert the value returned by count_entities method is equal to length of vectors |
||
181 | expected: the count is equal to the length of vectors |
||
182 | ''' |
||
183 | nq = 2 |
||
184 | vectors = gen_vectors(nq, dim) |
||
185 | res = connect.insert(collection_name=collection, records=vectors) |
||
186 | time.sleep(add_time_interval) |
||
187 | |||
188 | def rows_count(milvus): |
||
189 | status, res = milvus.count_entities(collection) |
||
190 | logging.getLogger().info(status) |
||
191 | assert res == nq |
||
192 | |||
193 | process_num = 8 |
||
194 | processes = [] |
||
195 | for i in range(process_num): |
||
196 | milvus = get_milvus(args["ip"], args["port"], handler=args["handler"]) |
||
197 | p = Process(target=rows_count, args=(milvus, )) |
||
198 | processes.append(p) |
||
199 | p.start() |
||
200 | logging.getLogger().info(p) |
||
201 | for p in processes: |
||
202 | p.join() |
||
203 | |||
204 | View Code Duplication | def test_collection_rows_count_multi_collections(self, connect): |
|
205 | ''' |
||
206 | target: test collection rows_count is correct or not with multiple collections of L2 |
||
207 | method: create collection and add vectors in it, |
||
208 | assert the value returned by count_entities method is equal to length of vectors |
||
209 | expected: the count is equal to the length of vectors |
||
210 | ''' |
||
211 | nq = 100 |
||
212 | vectors = gen_vectors(nq, dim) |
||
213 | collection_list = [] |
||
214 | for i in range(20): |
||
215 | collection_name = gen_unique_str() |
||
216 | collection_list.append(collection_name) |
||
217 | param = {'collection_name': collection_name, |
||
218 | 'dimension': dim, |
||
219 | 'index_file_size': index_file_size, |
||
220 | 'metric_type': MetricType.L2} |
||
221 | connect.create_collection(param) |
||
222 | res = connect.insert(collection_name=collection_name, records=vectors) |
||
223 | connect.flush(collection_list) |
||
224 | for i in range(20): |
||
225 | status, res = connect.count_entities(collection_list[i]) |
||
226 | assert status.OK() |
||
227 | assert res == nq |
||
228 | |||
229 | |||
230 | class TestCollectionCountIP: |
||
231 | """ |
||
232 | params means different nb, the nb value may trigger merge, or not |
||
233 | """ |
||
234 | |||
235 | @pytest.fixture( |
||
236 | scope="function", |
||
237 | params=[ |
||
238 | 1, |
||
239 | 5000, |
||
240 | 20000, |
||
241 | ], |
||
242 | ) |
||
243 | def insert_nb(self, request): |
||
244 | yield request.param |
||
245 | |||
246 | """ |
||
247 | generate valid create_index params |
||
248 | """ |
||
249 | |||
250 | View Code Duplication | @pytest.fixture( |
|
251 | scope="function", |
||
252 | params=gen_simple_index() |
||
253 | ) |
||
254 | def get_simple_index(self, request, connect): |
||
255 | if str(connect._cmd("mode")[1]) == "CPU": |
||
256 | if request.param["index_type"] == IndexType.IVF_SQ8H: |
||
257 | pytest.skip("sq8h not support in CPU mode") |
||
258 | if request.param["index_type"] == IndexType.IVF_PQ: |
||
259 | pytest.skip("Skip PQ Temporary") |
||
260 | return request.param |
||
261 | |||
262 | def test_collection_rows_count(self, connect, ip_collection, insert_nb): |
||
263 | ''' |
||
264 | target: test collection rows_count is correct or not |
||
265 | method: create collection and add vectors in it, |
||
266 | assert the value returned by count_entities method is equal to length of vectors |
||
267 | expected: the count is equal to the length of vectors |
||
268 | ''' |
||
269 | nb = insert_nb |
||
270 | vectors = gen_vectors(nb, dim) |
||
271 | res = connect.insert(collection_name=ip_collection, records=vectors) |
||
272 | connect.flush([ip_collection]) |
||
273 | status, res = connect.count_entities(ip_collection) |
||
274 | assert res == nb |
||
275 | |||
276 | View Code Duplication | def test_collection_rows_count_after_index_created(self, connect, ip_collection, get_simple_index): |
|
277 | ''' |
||
278 | target: test count_entities, after index have been created |
||
279 | method: add vectors in db, and create index, then calling count_entities with correct params |
||
280 | expected: count_entities raise exception |
||
281 | ''' |
||
282 | index_param = get_simple_index["index_param"] |
||
283 | index_type = get_simple_index["index_type"] |
||
284 | nb = 100 |
||
285 | vectors = gen_vectors(nb, dim) |
||
286 | res = connect.insert(collection_name=ip_collection, records=vectors) |
||
287 | connect.flush([ip_collection]) |
||
288 | connect.create_index(ip_collection, index_type, index_param) |
||
289 | status, res = connect.count_entities(ip_collection) |
||
290 | assert res == nb |
||
291 | |||
292 | # @pytest.mark.level(2) |
||
293 | # def test_count_without_connection(self, ip_collection, dis_connect): |
||
294 | # ''' |
||
295 | # target: test count_entities, without connection |
||
296 | # method: calling count_entities with correct params, with a disconnected instance |
||
297 | # expected: count_entities raise exception |
||
298 | # ''' |
||
299 | # with pytest.raises(Exception) as e: |
||
300 | # status = dis_connect.count_entities(ip_collection) |
||
301 | |||
302 | def test_collection_rows_count_no_vectors(self, connect, ip_collection): |
||
303 | ''' |
||
304 | target: test collection rows_count is correct or not, if collection is empty |
||
305 | method: create collection and no vectors in it, |
||
306 | assert the value returned by count_entities method is equal to 0 |
||
307 | expected: the count is equal to 0 |
||
308 | ''' |
||
309 | collection_name = gen_unique_str("test_collection") |
||
310 | param = {'collection_name': collection_name, |
||
311 | 'dimension': dim, |
||
312 | 'index_file_size': index_file_size} |
||
313 | connect.create_collection(param) |
||
314 | status, res = connect.count_entities(ip_collection) |
||
315 | assert res == 0 |
||
316 | |||
317 | # TODO: enable |
||
318 | View Code Duplication | @pytest.mark.timeout(60) |
|
319 | def _test_collection_rows_count_multiprocessing(self, connect, ip_collection, args): |
||
320 | ''' |
||
321 | target: test collection rows_count is correct or not with multiprocess |
||
322 | method: create collection and add vectors in it, |
||
323 | assert the value returned by count_entities method is equal to length of vectors |
||
324 | expected: the count is equal to the length of vectors |
||
325 | ''' |
||
326 | nq = 2 |
||
327 | vectors = gen_vectors(nq, dim) |
||
328 | res = connect.insert(collection_name=ip_collection, records=vectors) |
||
329 | time.sleep(add_time_interval) |
||
330 | |||
331 | def rows_count(milvus): |
||
332 | status, res = milvus.count_entities(ip_collection) |
||
333 | logging.getLogger().info(status) |
||
334 | assert res == nq |
||
335 | |||
336 | process_num = 8 |
||
337 | processes = [] |
||
338 | for i in range(process_num): |
||
339 | milvus = get_milvus(args["ip"], args["port"], handler=args["handler"]) |
||
340 | p = Process(target=rows_count, args=(milvus,)) |
||
341 | processes.append(p) |
||
342 | p.start() |
||
343 | logging.getLogger().info(p) |
||
344 | for p in processes: |
||
345 | p.join() |
||
346 | |||
347 | View Code Duplication | def test_collection_rows_count_multi_collections(self, connect): |
|
348 | ''' |
||
349 | target: test collection rows_count is correct or not with multiple collections of IP |
||
350 | method: create collection and add vectors in it, |
||
351 | assert the value returned by count_entities method is equal to length of vectors |
||
352 | expected: the count is equal to the length of vectors |
||
353 | ''' |
||
354 | nq = 100 |
||
355 | vectors = gen_vectors(nq, dim) |
||
356 | collection_list = [] |
||
357 | for i in range(20): |
||
358 | collection_name = gen_unique_str('test_collection_rows_count_multi_collections') |
||
359 | collection_list.append(collection_name) |
||
360 | param = {'collection_name': collection_name, |
||
361 | 'dimension': dim, |
||
362 | 'index_file_size': index_file_size, |
||
363 | 'metric_type': MetricType.IP} |
||
364 | connect.create_collection(param) |
||
365 | res = connect.insert(collection_name=collection_name, records=vectors) |
||
366 | connect.flush(collection_list) |
||
367 | for i in range(20): |
||
368 | status, res = connect.count_entities(collection_list[i]) |
||
369 | assert status.OK() |
||
370 | assert res == nq |
||
371 | |||
372 | |||
373 | class TestCollectionCountJAC: |
||
374 | """ |
||
375 | params means different nb, the nb value may trigger merge, or not |
||
376 | """ |
||
377 | |||
378 | @pytest.fixture( |
||
379 | scope="function", |
||
380 | params=[ |
||
381 | 1, |
||
382 | 5000, |
||
383 | 20000, |
||
384 | ], |
||
385 | ) |
||
386 | def insert_nb(self, request): |
||
387 | yield request.param |
||
388 | |||
389 | """ |
||
390 | generate valid create_index params |
||
391 | """ |
||
392 | |||
393 | View Code Duplication | @pytest.fixture( |
|
394 | scope="function", |
||
395 | params=gen_simple_index() |
||
396 | ) |
||
397 | def get_jaccard_index(self, request, connect): |
||
398 | logging.getLogger().info(request.param) |
||
399 | if request.param["index_type"] == IndexType.IVFLAT or request.param["index_type"] == IndexType.FLAT: |
||
400 | return request.param |
||
401 | else: |
||
402 | pytest.skip("Skip index Temporary") |
||
403 | |||
404 | def test_collection_rows_count(self, connect, jac_collection, insert_nb): |
||
405 | ''' |
||
406 | target: test collection rows_count is correct or not |
||
407 | method: create collection and add vectors in it, |
||
408 | assert the value returned by count_entities method is equal to length of vectors |
||
409 | expected: the count is equal to the length of vectors |
||
410 | ''' |
||
411 | nb = insert_nb |
||
412 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
413 | res = connect.insert(collection_name=jac_collection, records=vectors) |
||
414 | connect.flush([jac_collection]) |
||
415 | status, res = connect.count_entities(jac_collection) |
||
416 | assert res == nb |
||
417 | |||
418 | def test_collection_rows_count_after_index_created(self, connect, jac_collection, get_jaccard_index): |
||
419 | ''' |
||
420 | target: test count_entities, after index have been created |
||
421 | method: add vectors in db, and create index, then calling count_entities with correct params |
||
422 | expected: count_entities raise exception |
||
423 | ''' |
||
424 | nb = 100 |
||
425 | index_param = get_jaccard_index["index_param"] |
||
426 | index_type = get_jaccard_index["index_type"] |
||
427 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
428 | res = connect.insert(collection_name=jac_collection, records=vectors) |
||
429 | connect.flush([jac_collection]) |
||
430 | connect.create_index(jac_collection, index_type, index_param) |
||
431 | status, res = connect.count_entities(jac_collection) |
||
432 | assert res == nb |
||
433 | |||
434 | # @pytest.mark.level(2) |
||
435 | # def test_count_without_connection(self, jac_collection, dis_connect): |
||
436 | # ''' |
||
437 | # target: test count_entities, without connection |
||
438 | # method: calling count_entities with correct params, with a disconnected instance |
||
439 | # expected: count_entities raise exception |
||
440 | # ''' |
||
441 | # with pytest.raises(Exception) as e: |
||
442 | # status = dis_connect.count_entities(jac_collection) |
||
443 | |||
444 | def test_collection_rows_count_no_vectors(self, connect, jac_collection): |
||
445 | ''' |
||
446 | target: test collection rows_count is correct or not, if collection is empty |
||
447 | method: create collection and no vectors in it, |
||
448 | assert the value returned by count_entities method is equal to 0 |
||
449 | expected: the count is equal to 0 |
||
450 | ''' |
||
451 | collection_name = gen_unique_str("test_collection") |
||
452 | param = {'collection_name': collection_name, |
||
453 | 'dimension': dim, |
||
454 | 'index_file_size': index_file_size} |
||
455 | connect.create_collection(param) |
||
456 | status, res = connect.count_entities(jac_collection) |
||
457 | assert res == 0 |
||
458 | |||
459 | View Code Duplication | def test_collection_rows_count_multi_collections(self, connect): |
|
460 | ''' |
||
461 | target: test collection rows_count is correct or not with multiple collections of IP |
||
462 | method: create collection and add vectors in it, |
||
463 | assert the value returned by count_entities method is equal to length of vectors |
||
464 | expected: the count is equal to the length of vectors |
||
465 | ''' |
||
466 | nq = 100 |
||
467 | tmp, vectors = gen_binary_vectors(nq, dim) |
||
468 | collection_list = [] |
||
469 | for i in range(20): |
||
470 | collection_name = gen_unique_str('test_collection_rows_count_multi_collections') |
||
471 | collection_list.append(collection_name) |
||
472 | param = {'collection_name': collection_name, |
||
473 | 'dimension': dim, |
||
474 | 'index_file_size': index_file_size, |
||
475 | 'metric_type': MetricType.JACCARD} |
||
476 | connect.create_collection(param) |
||
477 | res = connect.insert(collection_name=collection_name, records=vectors) |
||
478 | connect.flush(collection_list) |
||
479 | for i in range(20): |
||
480 | status, res = connect.count_entities(collection_list[i]) |
||
481 | assert status.OK() |
||
482 | assert res == nq |
||
483 | |||
484 | class TestCollectionCountBinary: |
||
485 | """ |
||
486 | params means different nb, the nb value may trigger merge, or not |
||
487 | """ |
||
488 | |||
489 | @pytest.fixture( |
||
490 | scope="function", |
||
491 | params=[ |
||
492 | 1, |
||
493 | 5000, |
||
494 | 20000, |
||
495 | ], |
||
496 | ) |
||
497 | def insert_nb(self, request): |
||
498 | yield request.param |
||
499 | |||
500 | """ |
||
501 | generate valid create_index params |
||
502 | """ |
||
503 | |||
504 | View Code Duplication | @pytest.fixture( |
|
505 | scope="function", |
||
506 | params=gen_simple_index() |
||
507 | ) |
||
508 | def get_hamming_index(self, request, connect): |
||
509 | logging.getLogger().info(request.param) |
||
510 | if request.param["index_type"] == IndexType.IVFLAT or request.param["index_type"] == IndexType.FLAT: |
||
511 | return request.param |
||
512 | else: |
||
513 | pytest.skip("Skip index Temporary") |
||
514 | |||
515 | @pytest.fixture( |
||
516 | scope="function", |
||
517 | params=gen_simple_index() |
||
518 | ) |
||
519 | def get_substructure_index(self, request, connect): |
||
520 | logging.getLogger().info(request.param) |
||
521 | if request.param["index_type"] == IndexType.FLAT: |
||
522 | return request.param |
||
523 | else: |
||
524 | pytest.skip("Skip index Temporary") |
||
525 | |||
526 | @pytest.fixture( |
||
527 | scope="function", |
||
528 | params=gen_simple_index() |
||
529 | ) |
||
530 | def get_superstructure_index(self, request, connect): |
||
531 | logging.getLogger().info(request.param) |
||
532 | if request.param["index_type"] == IndexType.FLAT: |
||
533 | return request.param |
||
534 | else: |
||
535 | pytest.skip("Skip index Temporary") |
||
536 | |||
537 | def test_collection_rows_count(self, connect, ham_collection, insert_nb): |
||
538 | ''' |
||
539 | target: test collection rows_count is correct or not |
||
540 | method: create collection and add vectors in it, |
||
541 | assert the value returned by count_entities method is equal to length of vectors |
||
542 | expected: the count is equal to the length of vectors |
||
543 | ''' |
||
544 | nb = insert_nb |
||
545 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
546 | res = connect.insert(collection_name=ham_collection, records=vectors) |
||
547 | connect.flush([ham_collection]) |
||
548 | status, res = connect.count_entities(ham_collection) |
||
549 | assert res == nb |
||
550 | |||
551 | def test_collection_rows_count_substructure(self, connect, substructure_collection, insert_nb): |
||
552 | ''' |
||
553 | target: test collection rows_count is correct or not |
||
554 | method: create collection and add vectors in it, |
||
555 | assert the value returned by count_entities method is equal to length of vectors |
||
556 | expected: the count is equal to the length of vectors |
||
557 | ''' |
||
558 | nb = insert_nb |
||
559 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
560 | res = connect.insert(collection_name=substructure_collection, records=vectors) |
||
561 | connect.flush([substructure_collection]) |
||
562 | status, res = connect.count_entities(substructure_collection) |
||
563 | assert res == nb |
||
564 | |||
565 | def test_collection_rows_count_superstructure(self, connect, superstructure_collection, insert_nb): |
||
566 | ''' |
||
567 | target: test collection rows_count is correct or not |
||
568 | method: create collection and add vectors in it, |
||
569 | assert the value returned by count_entities method is equal to length of vectors |
||
570 | expected: the count is equal to the length of vectors |
||
571 | ''' |
||
572 | nb = insert_nb |
||
573 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
574 | res = connect.insert(collection_name=superstructure_collection, records=vectors) |
||
575 | connect.flush([superstructure_collection]) |
||
576 | status, res = connect.count_entities(superstructure_collection) |
||
577 | assert res == nb |
||
578 | |||
579 | def test_collection_rows_count_after_index_created(self, connect, ham_collection, get_hamming_index): |
||
580 | ''' |
||
581 | target: test count_entities, after index have been created |
||
582 | method: add vectors in db, and create index, then calling count_entities with correct params |
||
583 | expected: count_entities raise exception |
||
584 | ''' |
||
585 | nb = 100 |
||
586 | index_type = get_hamming_index["index_type"] |
||
587 | index_param = get_hamming_index["index_param"] |
||
588 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
589 | res = connect.insert(collection_name=ham_collection, records=vectors) |
||
590 | connect.flush([ham_collection]) |
||
591 | connect.create_index(ham_collection, index_type, index_param) |
||
592 | status, res = connect.count_entities(ham_collection) |
||
593 | assert res == nb |
||
594 | |||
595 | def test_collection_rows_count_after_index_created_substructure(self, connect, substructure_collection, get_substructure_index): |
||
596 | ''' |
||
597 | target: test count_entities, after index have been created |
||
598 | method: add vectors in db, and create index, then calling count_entities with correct params |
||
599 | expected: count_entities raise exception |
||
600 | ''' |
||
601 | nb = 100 |
||
602 | index_type = get_substructure_index["index_type"] |
||
603 | index_param = get_substructure_index["index_param"] |
||
604 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
605 | res = connect.insert(collection_name=substructure_collection, records=vectors) |
||
606 | connect.flush([substructure_collection]) |
||
607 | connect.create_index(substructure_collection, index_type, index_param) |
||
608 | status, res = connect.count_entities(substructure_collection) |
||
609 | assert res == nb |
||
610 | |||
611 | def test_collection_rows_count_after_index_created_superstructure(self, connect, superstructure_collection, get_superstructure_index): |
||
612 | ''' |
||
613 | target: test count_entities, after index have been created |
||
614 | method: add vectors in db, and create index, then calling count_entities with correct params |
||
615 | expected: count_entities raise exception |
||
616 | ''' |
||
617 | nb = 100 |
||
618 | index_type = get_superstructure_index["index_type"] |
||
619 | index_param = get_superstructure_index["index_param"] |
||
620 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
621 | res = connect.insert(collection_name=superstructure_collection, records=vectors) |
||
622 | connect.flush([superstructure_collection]) |
||
623 | connect.create_index(superstructure_collection, index_type, index_param) |
||
624 | status, res = connect.count_entities(superstructure_collection) |
||
625 | assert res == nb |
||
626 | |||
627 | # @pytest.mark.level(2) |
||
628 | # def test_count_without_connection(self, ham_collection, dis_connect): |
||
629 | # ''' |
||
630 | # target: test count_entities, without connection |
||
631 | # method: calling count_entities with correct params, with a disconnected instance |
||
632 | # expected: count_entities raise exception |
||
633 | # ''' |
||
634 | # with pytest.raises(Exception) as e: |
||
635 | # status = dis_connect.count_entities(ham_collection) |
||
636 | |||
637 | def test_collection_rows_count_no_vectors(self, connect, ham_collection): |
||
638 | ''' |
||
639 | target: test collection rows_count is correct or not, if collection is empty |
||
640 | method: create collection and no vectors in it, |
||
641 | assert the value returned by count_entities method is equal to 0 |
||
642 | expected: the count is equal to 0 |
||
643 | ''' |
||
644 | collection_name = gen_unique_str("test_collection") |
||
645 | param = {'collection_name': collection_name, |
||
646 | 'dimension': dim, |
||
647 | 'index_file_size': index_file_size} |
||
648 | connect.create_collection(param) |
||
649 | status, res = connect.count_entities(ham_collection) |
||
650 | assert res == 0 |
||
651 | |||
652 | View Code Duplication | def test_collection_rows_count_multi_collections(self, connect): |
|
653 | ''' |
||
654 | target: test collection rows_count is correct or not with multiple collections of IP |
||
655 | method: create collection and add vectors in it, |
||
656 | assert the value returned by count_entities method is equal to length of vectors |
||
657 | expected: the count is equal to the length of vectors |
||
658 | ''' |
||
659 | nq = 100 |
||
660 | tmp, vectors = gen_binary_vectors(nq, dim) |
||
661 | collection_list = [] |
||
662 | for i in range(20): |
||
663 | collection_name = gen_unique_str('test_collection_rows_count_multi_collections') |
||
664 | collection_list.append(collection_name) |
||
665 | param = {'collection_name': collection_name, |
||
666 | 'dimension': dim, |
||
667 | 'index_file_size': index_file_size, |
||
668 | 'metric_type': MetricType.HAMMING} |
||
669 | connect.create_collection(param) |
||
670 | res = connect.insert(collection_name=collection_name, records=vectors) |
||
671 | connect.flush(collection_list) |
||
672 | for i in range(20): |
||
673 | status, res = connect.count_entities(collection_list[i]) |
||
674 | assert status.OK() |
||
675 | assert res == nq |
||
676 | |||
677 | |||
678 | class TestCollectionCountTANIMOTO: |
||
679 | """ |
||
680 | params means different nb, the nb value may trigger merge, or not |
||
681 | """ |
||
682 | |||
683 | @pytest.fixture( |
||
684 | scope="function", |
||
685 | params=[ |
||
686 | 1, |
||
687 | 5000, |
||
688 | 20000, |
||
689 | ], |
||
690 | ) |
||
691 | def insert_nb(self, request): |
||
692 | yield request.param |
||
693 | |||
694 | """ |
||
695 | generate valid create_index params |
||
696 | """ |
||
697 | |||
698 | @pytest.fixture( |
||
699 | scope="function", |
||
700 | params=gen_simple_index() |
||
701 | ) |
||
702 | def get_tanimoto_index(self, request, connect): |
||
703 | logging.getLogger().info(request.param) |
||
704 | if request.param["index_type"] == IndexType.IVFLAT or request.param["index_type"] == IndexType.FLAT: |
||
705 | return request.param |
||
706 | else: |
||
707 | pytest.skip("Skip index Temporary") |
||
708 | |||
709 | def test_collection_rows_count(self, connect, tanimoto_collection, insert_nb): |
||
710 | ''' |
||
711 | target: test collection rows_count is correct or not |
||
712 | method: create collection and add vectors in it, |
||
713 | assert the value returned by count_entities method is equal to length of vectors |
||
714 | expected: the count is equal to the length of vectors |
||
715 | ''' |
||
716 | nb = insert_nb |
||
717 | tmp, vectors = gen_binary_vectors(nb, dim) |
||
718 | res = connect.insert(collection_name=tanimoto_collection, records=vectors) |
||
719 | connect.flush([tanimoto_collection]) |
||
720 | status, res = connect.count_entities(tanimoto_collection) |
||
721 | assert status.OK() |
||
722 | assert res == nb |
||
723 |