conftest.py

import time

# Third party imports
import pytest

# local application imports
from factorys import gen_unique_str, fake, records_factory
from milvus import Milvus, IndexType, MetricType

default_host = "127.0.0.1"
default_grpc_port = 19530
default_http_port = 19121


def pytest_addoption(parser):
    parser.addoption("--ip", action="store", default=default_host)
    parser.addoption("--handler", action="store", default="GRPC")
    parser.addoption("--port", action="store", default=default_grpc_port)


@pytest.fixture(scope="module")
def gip(request):
    ip_ = request.config.getoption("--ip")
    port_ = request.config.getoption("--port")

    return ip_, port_


@pytest.fixture(scope="module")
def ghandler(request):
    handler_ = request.config.getoption("--handler")
    return handler_


@pytest.fixture(scope="module")
def connect(request, ghandler):
    ip = request.config.getoption("--ip")
    handler = request.config.getoption("--handler")
    port_default = default_http_port if handler == "HTTP" else default_grpc_port
    port = request.config.getoption("--port", default=port_default)

    client = Milvus(host=ip, port=port, handler=ghandler)
    # milvus.connect()

    def fin():
        try:
            client.close()
        except:
            pass

    request.addfinalizer(fin)
    return client


@pytest.fixture(scope="module")
def gcon(request, ghandler):
    ip = request.config.getoption("--ip")
    port = request.config.getoption("--port")
    client = Milvus(host=ip, port=port, handler=ghandler)

    def fin():
        try:
            client.close()
            pass
        except Exception as e:
            print(e)
            pass

    request.addfinalizer(fin)
    return client


@pytest.fixture(scope="module")
def args(request):
    ip = request.config.getoption("--ip")
    port = request.config.getoption("--port")
    args = {"ip": ip, "port": port}
    return args


@pytest.fixture(scope="function")
def table(request, connect):
    ori_table_name = getattr(request.module, "table_id", "test")
    table_name = gen_unique_str(ori_table_name)
    dim = getattr(request.module, "dim", "128")
    param = {'collection_name': table_name,
             'dimension': dim,
             'index_type': IndexType.IVFLAT,
             'metric_type': MetricType.L2
             }
    connect.create_collection(param)

    def teardown():
        connect.drop_collection(table_name)

    request.addfinalizer(teardown)

    return table_name


@pytest.fixture(scope="function")
def gcollection(request, gcon):
    table_name = fake.collection_name()
    dim = getattr(request.module, "dim", 128)

    param = {'collection_name': table_name,
             'dimension': dim,
             'index_file_size': 1024,
             'metric_type': MetricType.L2
             }
    gcon.create_collection(param)
    gcon.flush([table_name])

    def teardown():
        status, table_names = gcon.list_collections()
        for name in table_names:
            gcon.drop_collection(name)

    request.addfinalizer(teardown)

    return table_name


@pytest.fixture(scope='function')
def gvector(request, gcon, gcollection):
    dim = getattr(request.module, 'dim', 128)

    records = records_factory(dim, 10000)

    gcon.insert(gcollection, records)
    gcon.flush([gcollection])

    return gcollection