milvus/tests/milvus_python_test/test_add_vectors.py

import time
import random
import pdb
import threading
import logging
from multiprocessing import Pool, Process
import pytest
from milvus import IndexType, MetricType
from utils import *


dim = 128
index_file_size = 10
table_id = "test_add"
ADD_TIMEOUT = 60
nprobe = 1
epsilon = 0.0001
tag = "1970-01-01"


class TestAddBase:
    """
    ******************************************************************
      The following cases are used to test `add_vectors / index / search / delete` mixed function
    ******************************************************************
    """
    @pytest.fixture(
        scope="function",
        params=gen_simple_index_params()
    )
    def get_simple_index_params(self, request, connect):
        if str(connect._cmd("mode")[1]) == "CPU":
            if request.param["index_type"] == IndexType.IVF_SQ8H:
                pytest.skip("sq8h not support in cpu mode")
        if request.param["index_type"] == IndexType.IVF_PQ:
            pytest.skip("Skip PQ Temporary")
        return request.param

    def test_add_vector_create_table(self, connect, table):
        '''
        target: test add vector, then create table again
        method: add vector and create table
        expected: status not ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        param = {'table_name': table,
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        assert not status.OK()

    def test_add_vector_has_table(self, connect, table):
        '''
        target: test add vector, then check table existence
        method: add vector and call HasTable
        expected: table exists, status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        assert assert_has_table(connect, table)

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_delete_table_add_vector(self, connect, table):
        '''
        target: test add vector after table deleted
        method: delete table and add vector
        expected: status not ok
        '''
        status = connect.delete_table(table)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        assert not status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_delete_table_add_vector_another(self, connect, table):
        '''
        target: test add vector to table_1 after table_2 deleted
        method: delete table_2 and add vector to table_1
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        status = connect.delete_table(table)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(param['table_name'], vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_delete_table(self, connect, table):
        '''
        target: test delete table after add vector
        method: add vector and delete table
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        status = connect.delete_table(table)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_delete_another_table(self, connect, table):
        '''
        target: test delete table_1 table after add vector to table_2
        method: add vector and delete table
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_delete_table(self, connect, table):
        '''
        target: test delete table after add vector for a while
        method: add vector, sleep, and delete table
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        time.sleep(1)
        status = connect.delete_table(table)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_delete_another_table(self, connect, table):
        '''
        target: test delete table_1 table after add vector to table_2 for a while
        method: add vector , sleep, and delete table
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        time.sleep(1)
        status = connect.delete_table(param['table_name'])
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_create_index_add_vector(self, connect, table, get_simple_index_params):
        '''
        target: test add vector after build index
        method: build index and add vector
        expected: status ok
        '''
        index_param = get_simple_index_params
        status = connect.create_index(table, index_param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_create_index_add_vector_another(self, connect, table, get_simple_index_params):
        '''
        target: test add vector to table_2 after build index for table_1
        method: build index and add vector
        expected: status ok
        '''
        index_param = get_simple_index_params
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        status = connect.create_index(table, index_param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        connect.delete_table(param['table_name'])
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_create_index(self, connect, table, get_simple_index_params):
        '''
        target: test build index add after vector
        method: add vector and build index
        expected: status ok
        '''
        index_param = get_simple_index_params
        logging.getLogger().info(index_param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        status = connect.create_index(table, index_param)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_create_index_another(self, connect, table, get_simple_index_params):
        '''
        target: test add vector to table_2 after build index for table_1
        method: build index and add vector
        expected: status ok
        '''
        index_param = get_simple_index_params
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        status = connect.create_index(param['table_name'], index_param)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_create_index(self, connect, table, get_simple_index_params):
        '''
        target: test build index add after vector for a while
        method: add vector and build index
        expected: status ok
        '''
        index_param = get_simple_index_params
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        time.sleep(1)
        status = connect.create_index(table, index_param)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_create_index_another(self, connect, table, get_simple_index_params):
        '''
        target: test add vector to table_2 after build index for table_1 for a while
        method: build index and add vector
        expected: status ok
        '''
        index_param = get_simple_index_params
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        time.sleep(1)
        status = connect.create_index(param['table_name'], index_param)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_search_vector_add_vector(self, connect, table):
        '''
        target: test add vector after search table
        method: search table and add vector
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, result = connect.search_vectors(table, 1, nprobe, vector)
        status, ids = connect.add_vectors(table, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_search_vector_add_vector_another(self, connect, table):
        '''
        target: test add vector to table_1 after search table_2
        method: search table and add vector
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, result = connect.search_vectors(table, 1, nprobe, vector)
        status, ids = connect.add_vectors(param['table_name'], vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_search_vector(self, connect, table):
        '''
        target: test search vector after add vector
        method: add vector and search table
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        status, result = connect.search_vectors(table, 1, nprobe, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_search_vector_another(self, connect, table):
        '''
        target: test add vector to table_1 after search table_2
        method: search table and add vector
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        status, result = connect.search_vectors(param['table_name'], 1, nprobe, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_search_vector(self, connect, table):
        '''
        target: test search vector after add vector after a while
        method: add vector, sleep, and search table
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        time.sleep(1)
        status, result = connect.search_vectors(table, 1, nprobe, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_search_vector_another(self, connect, table):
        '''
        target: test add vector to table_1 after search table_2 a while
        method: search table , sleep, and add vector
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(table, vector)
        time.sleep(1)
        status, result = connect.search_vectors(param['table_name'], 1, nprobe, vector)
        assert status.OK()

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

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_ids(self, connect, table):
        '''
        target: test add vectors in table, use customize ids
        method: create table and add vectors in it, check the ids returned and the table length after vectors added
        expected: the length of ids and the table row count
        '''
        nq = 5; top_k = 1; nprobe = 1
        vectors = gen_vectors(nq, dim)
        ids = [i for i in range(nq)]
        status, ids = connect.add_vectors(table, vectors, ids)
        time.sleep(2)
        assert status.OK()
        assert len(ids) == nq
        # check search result
        status, result = connect.search_vectors(table, top_k, nprobe, vectors)
        logging.getLogger().info(result)
        assert len(result) == nq
        for i in range(nq):
            assert result[i][0].id == i

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_twice_ids_no_ids(self, connect, table):
        '''
        target: check the result of add_vectors, with params ids and no ids
        method: test add vectors twice, use customize ids first, and then use no ids
        expected: status not OK 
        '''
        nq = 5; top_k = 1; nprobe = 1
        vectors = gen_vectors(nq, dim)
        ids = [i for i in range(nq)]
        status, ids = connect.add_vectors(table, vectors, ids)
        assert status.OK()
        status, ids = connect.add_vectors(table, vectors)
        logging.getLogger().info(status)
        logging.getLogger().info(ids)
        assert not status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_twice_not_ids_ids(self, connect, table):
        '''
        target: check the result of add_vectors, with params ids and no ids
        method: test add vectors twice, use not ids first, and then use customize ids
        expected: status not OK 
        '''
        nq = 5; top_k = 1; nprobe = 1
        vectors = gen_vectors(nq, dim)
        ids = [i for i in range(nq)]
        status, ids = connect.add_vectors(table, vectors)
        assert status.OK()
        status, ids = connect.add_vectors(table, vectors, ids)
        logging.getLogger().info(status)
        logging.getLogger().info(ids)
        assert not status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_ids_length_not_match(self, connect, table):
        '''
        target: test add vectors in table, use customize ids, len(ids) != len(vectors)
        method: create table and add vectors in it
        expected: raise an exception
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        ids = [i for i in range(1, nq)]
        with pytest.raises(Exception) as e:
            status, ids = connect.add_vectors(table, vectors, ids)

    @pytest.fixture(
        scope="function",
        params=gen_invalid_vector_ids()
    )
    def get_vector_id(self, request):
        yield request.param

    @pytest.mark.level(2)
    def test_add_vectors_ids_invalid(self, connect, table, get_vector_id):
        '''
        target: test add vectors in table, use customize ids, which are not int64
        method: create table and add vectors in it
        expected: raise an exception
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        vector_id = get_vector_id
        ids = [vector_id for i in range(nq)]
        with pytest.raises(Exception) as e:
            status, ids = connect.add_vectors(table, vectors, ids)

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors(self, connect, table):
        '''
        target: test add vectors in table created before
        method: create table and add vectors in it, check the ids returned and the table length after vectors added
        expected: the length of ids and the table row count
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        status, ids = connect.add_vectors(table, vectors)
        assert status.OK()
        assert len(ids) == nq

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_tag(self, connect, table):
        '''
        target: test add vectors in table created before
        method: create table and add vectors in it, with the partition_tag param
        expected: the table row count equals to nq
        '''
        nq = 5
        partition_name = gen_unique_str()
        vectors = gen_vectors(nq, dim)
        status = connect.create_partition(table, partition_name, tag)
        status, ids = connect.add_vectors(table, vectors, partition_tag=tag)
        assert status.OK()
        assert len(ids) == nq

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_tag_A(self, connect, table):
        '''
        target: test add vectors in table created before
        method: create partition and add vectors in it
        expected: the table row count equals to nq
        '''
        nq = 5
        partition_name = gen_unique_str()
        vectors = gen_vectors(nq, dim)
        status = connect.create_partition(table, partition_name, tag)
        status, ids = connect.add_vectors(partition_name, vectors)
        assert status.OK()
        assert len(ids) == nq

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_tag_not_existed(self, connect, table):
        '''
        target: test add vectors in table created before
        method: create table and add vectors in it, with the not existed partition_tag param
        expected: status not ok
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        status, ids = connect.add_vectors(table, vectors, partition_tag=tag)
        assert not status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_tag_not_existed_A(self, connect, table):
        '''
        target: test add vectors in table created before
        method: create partition, add vectors with the not existed partition_tag param
        expected: status not ok
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        new_tag = "new_tag"
        partition_name = gen_unique_str()
        status = connect.create_partition(table, partition_name, tag)
        status, ids = connect.add_vectors(table, vectors, partition_tag=new_tag)
        assert not status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_tag_existed(self, connect, table):
        '''
        target: test add vectors in table created before
        method: create table and add vectors in it repeatly, with the partition_tag param
        expected: the table row count equals to nq
        '''
        nq = 5
        partition_name = gen_unique_str()
        vectors = gen_vectors(nq, dim)
        status = connect.create_partition(table, partition_name, tag)
        status, ids = connect.add_vectors(table, vectors, partition_tag=tag)
        for i in range(5):
            status, ids = connect.add_vectors(table, vectors, partition_tag=tag)
            assert status.OK()
            assert len(ids) == nq

    @pytest.mark.level(2)
    def test_add_vectors_without_connect(self, dis_connect, table):
        '''
        target: test add vectors without connection
        method: create table and add vectors in it, check if added successfully
        expected: raise exception
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        with pytest.raises(Exception) as e:
            status, ids = dis_connect.add_vectors(table, vectors)

    def test_add_table_not_existed(self, connect):
        '''
        target: test add vectors in table, which not existed before
        method: add vectors table not existed, check the status
        expected: status not ok
        '''
        nq = 5
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(gen_unique_str("not_exist_table"), vector)
        assert not status.OK()
        assert not ids

    def test_add_vector_dim_not_matched(self, connect, table):
        '''
        target: test add vector, the vector dimension is not equal to the table dimension
        method: the vector dimension is half of the table dimension, check the status
        expected: status not ok
        '''
        vector = gen_single_vector(int(dim)//2)
        status, ids = connect.add_vectors(table, vector)
        assert not status.OK()

    def test_add_vectors_dim_not_matched(self, connect, table):
        '''
        target: test add vectors, the vector dimension is not equal to the table dimension
        method: the vectors dimension is half of the table dimension, check the status
        expected: status not ok
        '''
        nq = 5
        vectors = gen_vectors(nq, int(dim)//2)
        status, ids = connect.add_vectors(table, vectors)
        assert not status.OK()

    def test_add_vector_query_after_sleep(self, connect, table):
        '''
        target: test add vectors, and search it after sleep
        method: set vector[0][1] as query vectors
        expected: status ok and result length is 1
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        status, ids = connect.add_vectors(table, vectors)
        time.sleep(3)
        status, result = connect.search_vectors(table, 1, nprobe, [vectors[0]])
        assert status.OK()
        assert len(result) == 1

    # TODO: enable
    # @pytest.mark.repeat(10)
    @pytest.mark.timeout(ADD_TIMEOUT)
    def _test_add_vector_with_multiprocessing(self, args):
        '''
        target: test add vectors, with multi processes
        method: 10 processed add vectors concurrently
        expected: status ok and result length is equal to the length off added vectors
        '''
        table = gen_unique_str()
        uri = "tcp://%s:%s" % (args["ip"], args["port"])
        param = {'table_name': table,
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        milvus = get_milvus()
        milvus.connect(uri=uri)
        milvus.create_table(param)
        vector = gen_single_vector(dim)
        process_num = 4
        loop_num = 5
        processes = []
        def add():
            milvus = get_milvus()
            milvus.connect(uri=uri)
            i = 0
            while i < loop_num:
                status, ids = milvus.add_vectors(table, vector)
                i = i + 1
            milvus.disconnect()
        for i in range(process_num):
            p = Process(target=add, args=())
            processes.append(p)
            p.start()
            time.sleep(0.2)
        for p in processes:
            p.join()
        time.sleep(2)
        status, count = milvus.get_table_row_count(table)
        assert count == process_num * loop_num

    def test_add_vector_multi_tables(self, connect):
        '''
        target: test add vectors is correct or not with multiple tables of L2
        method: create 50 tables and add vectors into them in turn
        expected: status ok
        '''
        nq = 100
        vectors = gen_vectors(nq, dim)
        table_list = []
        for i in range(20):
            table_name = gen_unique_str('test_add_vector_multi_tables')
            table_list.append(table_name)
            param = {'table_name': table_name,
                     'dimension': dim,
                     'index_file_size': index_file_size,
                     'metric_type': MetricType.L2}
            connect.create_table(param)
        time.sleep(5)
        for j in range(5):
            for i in range(20):
                status, ids = connect.add_vectors(table_name=table_list[i], records=vectors)
                assert status.OK()

class TestAddIP:
    """
    ******************************************************************
      The following cases are used to test `add_vectors / index / search / delete` mixed function
    ******************************************************************
    """
    @pytest.fixture(
        scope="function",
        params=gen_simple_index_params()
    )
    def get_simple_index_params(self, request, connect):
        if str(connect._cmd("mode")[1]) == "CPU":
            if request.param["index_type"] == IndexType.IVF_SQ8H:
                pytest.skip("sq8h not support in cpu mode")
        if request.param["index_type"] == IndexType.IVF_PQ:
            pytest.skip("Skip PQ Temporary")
        return request.param

    def test_add_vector_create_table(self, connect, ip_table):
        '''
        target: test add vector, then create table again
        method: add vector and create table
        expected: status not ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        param = {'table_name': ip_table,
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        assert not status.OK()

    def test_add_vector_has_table(self, connect, ip_table):
        '''
        target: test add vector, then check table existence
        method: add vector and call HasTable
        expected: table exists, status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        assert assert_has_table(connect, ip_table)

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_delete_table_add_vector(self, connect, ip_table):
        '''
        target: test add vector after table deleted
        method: delete table and add vector
        expected: status not ok
        '''
        status = connect.delete_table(ip_table)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        assert not status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_delete_table_add_vector_another(self, connect, ip_table):
        '''
        target: test add vector to table_1 after table_2 deleted
        method: delete table_2 and add vector to table_1
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        status = connect.delete_table(ip_table)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(param['table_name'], vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_delete_table(self, connect, ip_table):
        '''
        target: test delete table after add vector
        method: add vector and delete table
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        status = connect.delete_table(ip_table)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_delete_another_table(self, connect, ip_table):
        '''
        target: test delete table_1 table after add vector to table_2
        method: add vector and delete table
        expected: status ok
        '''
        param = {'table_name': 'test_add_vector_delete_another_table',
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        status = connect.delete_table(param['table_name'])
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_delete_table(self, connect, ip_table):
        '''
        target: test delete table after add vector for a while
        method: add vector, sleep, and delete table
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        time.sleep(1)
        status = connect.delete_table(ip_table)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_delete_another_table(self, connect, ip_table):
        '''
        target: test delete table_1 table after add vector to table_2 for a while
        method: add vector , sleep, and delete table
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        time.sleep(1)
        status = connect.delete_table(param['table_name'])
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_create_index_add_vector(self, connect, ip_table, get_simple_index_params):
        '''
        target: test add vector after build index
        method: build index and add vector
        expected: status ok
        '''
        index_param = get_simple_index_params
        status = connect.create_index(ip_table, index_param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_create_index_add_vector_another(self, connect, ip_table, get_simple_index_params):
        '''
        target: test add vector to table_2 after build index for table_1
        method: build index and add vector
        expected: status ok
        '''
        index_param = get_simple_index_params
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        status = connect.create_index(ip_table, index_param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_create_index(self, connect, ip_table, get_simple_index_params):
        '''
        target: test build index add after vector
        method: add vector and build index
        expected: status ok
        '''
        index_param = get_simple_index_params
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        status, mode = connect._cmd("mode")
        assert status.OK()
        status = connect.create_index(ip_table, index_param)
        if str(mode) == "GPU" and (index_param["index_type"] == IndexType.IVF_PQ):
            assert not status.OK()
        else:
            assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_create_index_another(self, connect, ip_table, get_simple_index_params):
        '''
        target: test add vector to table_2 after build index for table_1
        method: build index and add vector
        expected: status ok
        '''
        index_param = get_simple_index_params
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        status = connect.create_index(param['table_name'], index_param)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_create_index(self, connect, ip_table, get_simple_index_params):
        '''
        target: test build index add after vector for a while
        method: add vector and build index
        expected: status ok
        '''
        index_param = get_simple_index_params
        if index_param["index_type"] == IndexType.IVF_PQ:
            pytest.skip("Skip some PQ cases")
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        time.sleep(1)
        status = connect.create_index(ip_table, index_param)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_create_index_another(self, connect, ip_table, get_simple_index_params):
        '''
        target: test add vector to table_2 after build index for table_1 for a while
        method: build index and add vector
        expected: status ok
        '''
        index_param = get_simple_index_params
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        time.sleep(1)
        status = connect.create_index(param['table_name'], index_param)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_search_vector_add_vector(self, connect, ip_table):
        '''
        target: test add vector after search table
        method: search table and add vector
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, result = connect.search_vectors(ip_table, 1, nprobe, vector)
        status, ids = connect.add_vectors(ip_table, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_search_vector_add_vector_another(self, connect, ip_table):
        '''
        target: test add vector to table_1 after search table_2
        method: search table and add vector
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, result = connect.search_vectors(ip_table, 1, nprobe, vector)
        status, ids = connect.add_vectors(param['table_name'], vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_search_vector(self, connect, ip_table):
        '''
        target: test search vector after add vector
        method: add vector and search table
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        status, result = connect.search_vectors(ip_table, 1, nprobe, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_search_vector_another(self, connect, ip_table):
        '''
        target: test add vector to table_1 after search table_2
        method: search table and add vector
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        status, result = connect.search_vectors(param['table_name'], 1, nprobe, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_search_vector(self, connect, ip_table):
        '''
        target: test search vector after add vector after a while
        method: add vector, sleep, and search table
        expected: status ok
        '''
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        time.sleep(1)
        status, result = connect.search_vectors(ip_table, 1, nprobe, vector)
        assert status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vector_sleep_search_vector_another(self, connect, ip_table):
        '''
        target: test add vector to table_1 after search table_2 a while
        method: search table , sleep, and add vector
        expected: status ok
        '''
        param = {'table_name': gen_unique_str(),
                 'dimension': dim,
                 'index_file_size': index_file_size,
                 'metric_type': MetricType.L2}
        status = connect.create_table(param)
        vector = gen_single_vector(dim)
        status, ids = connect.add_vectors(ip_table, vector)
        time.sleep(1)
        status, result = connect.search_vectors(param['table_name'], 1, nprobe, vector)
        assert status.OK()

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

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_ids(self, connect, ip_table):
        '''
        target: test add vectors in table, use customize ids
        method: create table and add vectors in it, check the ids returned and the table length after vectors added
        expected: the length of ids and the table row count
        '''
        nq = 5; top_k = 1; nprobe = 1
        vectors = gen_vectors(nq, dim)
        ids = [i for i in range(nq)]
        status, ids = connect.add_vectors(ip_table, vectors, ids)
        time.sleep(2)
        assert status.OK()
        assert len(ids) == nq
        # check search result
        status, result = connect.search_vectors(ip_table, top_k, nprobe, vectors)
        logging.getLogger().info(result)
        assert len(result) == nq
        for i in range(nq):
            assert result[i][0].id == i

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_twice_ids_no_ids(self, connect, ip_table):
        '''
        target: check the result of add_vectors, with params ids and no ids
        method: test add vectors twice, use customize ids first, and then use no ids
        expected: status not OK 
        '''
        nq = 5; top_k = 1; nprobe = 1
        vectors = gen_vectors(nq, dim)
        ids = [i for i in range(nq)]
        status, ids = connect.add_vectors(ip_table, vectors, ids)
        assert status.OK()
        status, ids = connect.add_vectors(ip_table, vectors)
        logging.getLogger().info(status)
        logging.getLogger().info(ids)
        assert not status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_twice_not_ids_ids(self, connect, ip_table):
        '''
        target: check the result of add_vectors, with params ids and no ids
        method: test add vectors twice, use not ids first, and then use customize ids
        expected: status not OK 
        '''
        nq = 5; top_k = 1; nprobe = 1
        vectors = gen_vectors(nq, dim)
        ids = [i for i in range(nq)]
        status, ids = connect.add_vectors(ip_table, vectors)
        assert status.OK()
        status, ids = connect.add_vectors(ip_table, vectors, ids)
        logging.getLogger().info(status)
        logging.getLogger().info(ids)
        assert not status.OK()

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors_ids_length_not_match(self, connect, ip_table):
        '''
        target: test add vectors in table, use customize ids, len(ids) != len(vectors)
        method: create table and add vectors in it
        expected: raise an exception
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        ids = [i for i in range(1, nq)]
        with pytest.raises(Exception) as e:
            status, ids = connect.add_vectors(ip_table, vectors, ids)

    @pytest.fixture(
        scope="function",
        params=gen_invalid_vector_ids()
    )
    def get_vector_id(self, request):
        yield request.param

    @pytest.mark.level(2)
    def test_add_vectors_ids_invalid(self, connect, ip_table, get_vector_id):
        '''
        target: test add vectors in table, use customize ids, which are not int64
        method: create table and add vectors in it
        expected: raise an exception
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        vector_id = get_vector_id
        ids = [vector_id for i in range(nq)]
        with pytest.raises(Exception) as e:
            status, ids = connect.add_vectors(ip_table, vectors, ids)

    @pytest.mark.timeout(ADD_TIMEOUT)
    def test_add_vectors(self, connect, ip_table):
        '''
        target: test add vectors in table created before
        method: create table and add vectors in it, check the ids returned and the table length after vectors added
        expected: the length of ids and the table row count
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        status, ids = connect.add_vectors(ip_table, vectors)
        assert status.OK()
        assert len(ids) == nq

    @pytest.mark.level(2)
    def test_add_vectors_without_connect(self, dis_connect, ip_table):
        '''
        target: test add vectors without connection
        method: create table and add vectors in it, check if added successfully
        expected: raise exception
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        with pytest.raises(Exception) as e:
            status, ids = dis_connect.add_vectors(ip_table, vectors)

    def test_add_vector_dim_not_matched(self, connect, ip_table):
        '''
        target: test add vector, the vector dimension is not equal to the table dimension
        method: the vector dimension is half of the table dimension, check the status
        expected: status not ok
        '''
        vector = gen_single_vector(int(dim)//2)
        status, ids = connect.add_vectors(ip_table, vector)
        assert not status.OK()

    def test_add_vectors_dim_not_matched(self, connect, ip_table):
        '''
        target: test add vectors, the vector dimension is not equal to the table dimension
        method: the vectors dimension is half of the table dimension, check the status
        expected: status not ok
        '''
        nq = 5
        vectors = gen_vectors(nq, int(dim)//2)
        status, ids = connect.add_vectors(ip_table, vectors)
        assert not status.OK()

    def test_add_vector_query_after_sleep(self, connect, ip_table):
        '''
        target: test add vectors, and search it after sleep
        method: set vector[0][1] as query vectors
        expected: status ok and result length is 1
        '''
        nq = 5
        vectors = gen_vectors(nq, dim)
        status, ids = connect.add_vectors(ip_table, vectors)
        time.sleep(3)
        status, result = connect.search_vectors(ip_table, 1, nprobe, [vectors[0]])
        assert status.OK()
        assert len(result) == 1

    def test_add_vector_multi_tables(self, connect):
        '''
        target: test add vectors is correct or not with multiple tables of IP
        method: create 50 tables and add vectors into them in turn
        expected: status ok
        '''
        nq = 100
        vectors = gen_vectors(nq, dim)
        table_list = []
        for i in range(20):
            table_name = gen_unique_str('test_add_vector_multi_tables')
            table_list.append(table_name)
            param = {'table_name': table_name,
                     'dimension': dim,
                     'index_file_size': index_file_size,
                     'metric_type': MetricType.IP}
            connect.create_table(param)
        time.sleep(2)
        for j in range(10):
            for i in range(20):
                status, ids = connect.add_vectors(table_name=table_list[i], records=vectors)
                assert status.OK()

class TestAddAdvance:
    @pytest.fixture(
        scope="function",
        params=[
            1,
            10,
            100,
            1000,
            pytest.param(5000 - 1, marks=pytest.mark.xfail),
            pytest.param(5000, marks=pytest.mark.xfail),
            pytest.param(5000 + 1, marks=pytest.mark.xfail),
        ],
    )
    def insert_count(self, request):
        yield request.param

    def test_insert_much(self, connect, table, insert_count):
        '''
        target: test add vectors with different length of vectors
        method: set different vectors as add method params
        expected: length of ids is equal to the length of vectors
        '''
        nb = insert_count
        insert_vec_list = gen_vectors(nb, dim)
        status, ids = connect.add_vectors(table, insert_vec_list)
        assert len(ids) == nb
        assert status.OK()

    def test_insert_much_ip(self, connect, ip_table, insert_count):
        '''
        target: test add vectors with different length of vectors
        method: set different vectors as add method params
        expected: length of ids is equal to the length of vectors
        '''
        nb = insert_count
        insert_vec_list = gen_vectors(nb, dim)
        status, ids = connect.add_vectors(ip_table, insert_vec_list)
        assert len(ids) == nb
        assert status.OK()

    def test_insert_much_jaccard(self, connect, jac_table, insert_count):
        '''
        target: test add vectors with different length of vectors
        method: set different vectors as add method params
        expected: length of ids is equal to the length of vectors
        '''
        nb = insert_count
        tmp, insert_vec_list = gen_binary_vectors(nb, dim)
        status, ids = connect.add_vectors(jac_table, insert_vec_list)
        assert len(ids) == nb
        assert status.OK()

    def test_insert_much_hamming(self, connect, ham_table, insert_count):
        '''
        target: test add vectors with different length of vectors
        method: set different vectors as add method params
        expected: length of ids is equal to the length of vectors
        '''
        nb = insert_count
        tmp, insert_vec_list = gen_binary_vectors(nb, dim)
        status, ids = connect.add_vectors(ham_table, insert_vec_list)
        assert len(ids) == nb
        assert status.OK()

    def test_insert_much_tanimoto(self, connect, tanimoto_table, insert_count):
        '''
        target: test add vectors with different length of vectors
        method: set different vectors as add method params
        expected: length of ids is equal to the length of vectors
        '''
        nb = insert_count
        tmp, insert_vec_list = gen_binary_vectors(nb, dim)
        status, ids = connect.add_vectors(tanimoto_table, insert_vec_list)
        assert len(ids) == nb
        assert status.OK()


class TestNameInvalid(object):
    """
    Test adding vectors with invalid table names
    """
    @pytest.fixture(
        scope="function",
        params=gen_invalid_table_names()
    )
    def get_table_name(self, request):
        yield request.param

    @pytest.fixture(
        scope="function",
        params=gen_invalid_table_names()
    )
    def get_tag_name(self, request):
        yield request.param

    @pytest.mark.level(2)
    def test_add_vectors_with_invalid_table_name(self, connect, get_table_name):
        table_name = get_table_name
        vectors = gen_vectors(1, dim)
        status, result = connect.add_vectors(table_name, vectors)
        assert not status.OK()

    @pytest.mark.level(2)
    def test_add_vectors_with_invalid_tag_name(self, connect, get_table_name, get_tag_name):
        table_name = get_table_name
        tag_name = get_tag_name
        vectors = gen_vectors(1, dim)
        status, result = connect.add_vectors(table_name, vectors, partition_tag=tag_name)
        assert not status.OK()


class TestAddTableVectorsInvalid(object):
    single_vector = gen_single_vector(dim)
    vectors = gen_vectors(2, dim)

    """
    Test adding vectors with invalid vectors
    """
    @pytest.fixture(
        scope="function",
        params=gen_invalid_vectors()
    )
    def gen_vector(self, request):
        yield request.param

    @pytest.mark.level(2)
    def test_add_vector_with_invalid_vectors(self, connect, table, gen_vector):
        tmp_single_vector = copy.deepcopy(self.single_vector)
        tmp_single_vector[0][1] = gen_vector
        with pytest.raises(Exception) as e:
            status, result = connect.add_vectors(table, tmp_single_vector)

    @pytest.mark.level(2)
    def test_add_vectors_with_invalid_vectors(self, connect, table, gen_vector):
        tmp_vectors = copy.deepcopy(self.vectors)
        tmp_vectors[1][1] = gen_vector
        with pytest.raises(Exception) as e:
            status, result = connect.add_vectors(table, tmp_vectors)

    @pytest.mark.level(2)
    def test_add_vectors_with_invalid_vectors_jaccard(self, connect, jac_table, gen_vector):
        tmp_vectors = copy.deepcopy(self.vectors)
        tmp_vectors[1][1] = gen_vector
        with pytest.raises(Exception) as e:
            status, result = connect.add_vectors(jac_table, tmp_vectors)

    @pytest.mark.level(2)
    def test_add_vectors_with_invalid_vectors_hamming(self, connect, ham_table, gen_vector):
        tmp_vectors = copy.deepcopy(self.vectors)
        tmp_vectors[1][1] = gen_vector
        with pytest.raises(Exception) as e:
            status, result = connect.add_vectors(ham_table, tmp_vectors)