def test_connect_uri_null(self, args): ''' target: test connect with null uri method: uri set null expected: connected is True ''' milvus = Milvus() uri_value = "" if self.local_ip(args): milvus.connect(uri=uri_value, timeout=1) assert milvus.connected() else: with pytest.raises(Exception) as e: milvus.connect(uri=uri_value, timeout=1) assert not milvus.connected()
def test_connect_correct_ip_port(self, args): ''' target: test connect with corrent ip and port value method: set correct ip and port expected: connected is True ''' milvus = Milvus() milvus.connect(host=args["ip"], port=args["port"]) assert milvus.connected()
def _test_connect_ip_localhost(self, args): ''' target: test connect with ip value: localhost method: set host localhost expected: connected is True ''' milvus = Milvus() milvus.connect(host='localhost', port=args["port"]) assert milvus.connected()
def test_connect_connected(self, args): ''' target: test connect and disconnect with corrent ip and port value, assert connected value method: set correct ip and port expected: connected is False ''' milvus = Milvus() milvus.connect(host=args["ip"], port=args["port"]) milvus.disconnect() assert not milvus.connected()
def test_connect_uri(self, args): ''' target: test connect with correct uri method: uri format and value are both correct expected: connected is True ''' milvus = Milvus() uri_value = "tcp://%s:%s" % (args["ip"], args["port"]) milvus.connect(uri=uri_value) assert milvus.connected()
def test_connect_param_priority_no_port(self, args): ''' target: both host_ip_port / uri are both given, if port is null, use the uri params method: port set "", check if wrong uri connection is ok expected: connect raise an exception and connected is false ''' milvus = Milvus() uri_value = "tcp://%s:19540" % args["ip"] milvus.connect(host=args["ip"], port="", uri=uri_value) assert milvus.connected()
def test_connect_param_priority_both_hostip_uri(self, args): ''' target: both host_ip_port / uri are both given, and not null, use the uri params method: check if wrong uri connection is ok expected: connect raise an exception and connected is false ''' milvus = Milvus() uri_value = "tcp://%s:%s" % (args["ip"], args["port"]) with pytest.raises(Exception) as e: milvus.connect(host=args["ip"], port=19540, uri=uri_value, timeout=1) assert not milvus.connected()
def test_connect_param_priority_uri(self, args): ''' target: both host_ip_port / uri are both given, if host is null, use the uri params method: host set "", check if correct uri connection is ok expected: connected is False ''' milvus = Milvus() uri_value = "tcp://%s:%s" % (args["ip"], args["port"]) with pytest.raises(Exception) as e: milvus.connect(host="", port=args["port"], uri=uri_value, timeout=1) assert not milvus.connected()
def test_connect_wrong_ip_null(self, args): ''' target: test connect with wrong ip value method: set host null expected: not use default ip, connected is False ''' milvus = Milvus() ip = "" with pytest.raises(Exception) as e: milvus.connect(host=ip, port=args["port"], timeout=1) assert not milvus.connected()
def test_connect_with_invalid_uri(self, get_invalid_uri): ''' target: test uri connect with invalid uri value method: set port in gen_invalid_uris expected: connected is False ''' milvus = Milvus() uri_value = get_invalid_uri with pytest.raises(Exception) as e: milvus.connect(uri=uri_value, timeout=1) assert not milvus.connected()
def test_connect_with_invalid_port(self, args, get_invalid_port): ''' target: test ip:port connect with invalid port value method: set port in gen_invalid_ports expected: connected is False ''' milvus = Milvus() port = get_invalid_port with pytest.raises(Exception) as e: milvus.connect(host=args["ip"], port=port, timeout=1) assert not milvus.connected()
def test_connect_wrong_uri_wrong_ip_null(self, args): ''' target: test uri connect with ip value wouldn't connected method: set uri ip null expected: connected is True ''' milvus = Milvus() uri_value = "tcp://:%s" % args["port"] with pytest.raises(Exception) as e: milvus.connect(uri=uri_value, timeout=1) assert not milvus.connected()
def test_connect_repeatedly(self, args): ''' target: test connect repeatedly method: connect again expected: status.code is 0, and status.message shows have connected already ''' milvus = Milvus() uri_value = "tcp://%s:%s" % (args["ip"], args["port"]) milvus.connect(uri=uri_value) milvus.connect(uri=uri_value) assert milvus.connected()
def test_connect_disconnect_repeatedly_once(self, args): ''' target: test connect and disconnect repeatedly method: disconnect, and then connect, assert connect status expected: status.code is 0 ''' milvus = Milvus() uri_value = "tcp://%s:%s" % (args["ip"], args["port"]) milvus.connect(uri=uri_value) milvus.disconnect() milvus.connect(uri=uri_value) assert milvus.connected()
def test_connect_disconnect_repeatedly_times(self, args): ''' target: test connect and disconnect for 10 times repeatedly method: disconnect, and then connect, assert connect status expected: status.code is 0 ''' times = 10 milvus = Milvus() uri_value = "tcp://%s:%s" % (args["ip"], args["port"]) milvus.connect(uri=uri_value) for i in range(times): milvus.disconnect() milvus.connect(uri=uri_value) assert milvus.connected()
import copy import time import random import mock import sys sys.path.append('.') from milvus import ParamError, Milvus import pytest client = Milvus() client._handler.connected = mock.Mock(return_value=True) client.connected = mock.Mock(return_value=True) def test_create_table_param(gcon): _PARAM = { "table_name": "table_name_{}".format(str(random.randint(0, 10000))), "dimension": 128 } table_param = copy.deepcopy(_PARAM) status = gcon.create_table(table_param) assert status.OK() time.sleep(1) gcon.drop_table(table_param["table_name"]) table_param["table_name"] = 12343
def test_connect_with_invalid_ip(self, args, get_invalid_ip): milvus = Milvus() ip = get_invalid_ip with pytest.raises(Exception) as e: milvus.connect(host=ip, port=args["port"], timeout=1) assert not milvus.connected()
def test_true_connect(self, gip): cnn = Milvus() cnn.connect(*gip) assert cnn.status.OK assert cnn.connected()
class Connection: def __init__(self, name, uri, max_retry=1, error_handlers=None, **kwargs): self.name = name self.uri = uri self.max_retry = max_retry self.retried = 0 self.conn = Milvus() self.error_handlers = [] if not error_handlers else error_handlers self.on_retry_func = kwargs.get('on_retry_func', None) # define search hook self.conn.set_hook(search_in_file=Searchook()) # self._connect() def __str__(self): return 'Connection:name=\"{}\";uri=\"{}\"'.format(self.name, self.uri) def _connect(self, metadata=None): try: self.conn.connect(uri=self.uri) except Exception as e: if not self.error_handlers: raise exceptions.ConnectionConnectError(message=str(e), metadata=metadata) for handler in self.error_handlers: handler(e, metadata=metadata) @property def can_retry(self): return self.retried < self.max_retry @property def connected(self): return self.conn.connected() def on_retry(self): if self.on_retry_func: self.on_retry_func(self) else: self.retried > 1 and logger.warning('{} is retrying {}'.format( self, self.retried)) def on_connect(self, metadata=None): while not self.connected and self.can_retry: self.retried += 1 self.on_retry() self._connect(metadata=metadata) if not self.can_retry and not self.connected: raise exceptions.ConnectionConnectError( message='Max retry {} reached!'.format(self.max_retry, metadata=metadata)) self.retried = 0 def connect(self, func, exception_handler=None): @wraps(func) def inner(*args, **kwargs): self.on_connect() try: return func(*args, **kwargs) except Exception as e: if exception_handler: exception_handler(e) else: raise e return inner
class MilvusClient(object): def __init__(self, table_name=None, ip=None, port=None, timeout=60): self._milvus = Milvus() self._table_name = table_name try: i = 1 start_time = time.time() if not ip: self._milvus.connect(host=SERVER_HOST_DEFAULT, port=SERVER_PORT_DEFAULT) else: # retry connect for remote server while time.time() < start_time + timeout: try: self._milvus.connect(host=ip, port=port) if self._milvus.connected() is True: logger.debug( "Try connect times: %d, %s" % (i, round(time.time() - start_time, 2))) break except Exception as e: logger.debug("Milvus connect failed") i = i + 1 except Exception as e: raise e def __str__(self): return 'Milvus table %s' % self._table_name def check_status(self, status): if not status.OK(): logger.error(status.message) # raise Exception("Status not ok") def create_table(self, table_name, dimension, index_file_size, metric_type): if not self._table_name: self._table_name = table_name if metric_type == "l2": metric_type = MetricType.L2 elif metric_type == "ip": metric_type = MetricType.IP elif metric_type == "jaccard": metric_type = MetricType.JACCARD elif metric_type == "hamming": metric_type = MetricType.HAMMING else: logger.error("Not supported metric_type: %s" % metric_type) create_param = { 'table_name': table_name, 'dimension': dimension, 'index_file_size': index_file_size, "metric_type": metric_type } status = self._milvus.create_table(create_param) self.check_status(status) @time_wrapper def insert(self, X, ids=None): status, result = self._milvus.add_vectors(self._table_name, X, ids) self.check_status(status) return status, result @time_wrapper def create_index(self, index_type, nlist): index_params = { "index_type": INDEX_MAP[index_type], "nlist": nlist, } logger.info("Building index start, table_name: %s, index_params: %s" % (self._table_name, json.dumps(index_params))) status = self._milvus.create_index(self._table_name, index=index_params) self.check_status(status) def describe_index(self): status, result = self._milvus.describe_index(self._table_name) index_type = None for k, v in INDEX_MAP.items(): if result._index_type == v: index_type = k break nlist = result._nlist res = {"index_type": index_type, "nlist": nlist} return res def drop_index(self): logger.info("Drop index: %s" % self._table_name) return self._milvus.drop_index(self._table_name) @time_wrapper def query(self, X, top_k, nprobe): status, result = self._milvus.search_vectors(self._table_name, top_k, nprobe, X) self.check_status(status) return result def count(self): return self._milvus.get_table_row_count(self._table_name)[1] def delete(self, timeout=60): logger.info("Start delete table: %s" % self._table_name) self._milvus.delete_table(self._table_name) i = 0 while i < timeout: if self.count(): time.sleep(1) i = i + 1 continue else: break if i >= timeout: logger.error("Delete table timeout") def describe(self): return self._milvus.describe_table(self._table_name) def show_tables(self): return self._milvus.show_tables() def exists_table(self, table_name=None): if table_name is None: table_name = self._table_name status, res = self._milvus.has_table(table_name) self.check_status(status) return res @time_wrapper def preload_table(self): return self._milvus.preload_table(self._table_name, timeout=3000) def get_server_version(self): status, res = self._milvus.server_version() return res def get_server_mode(self): return self.cmd("mode") def get_server_commit(self): return self.cmd("build_commit_id") def get_server_config(self): return json.loads(self.cmd("get_config *")) def get_mem_info(self): result = json.loads(self.cmd("get_system_info")) result_human = { # unit: Gb "memory_used": round(int(result["memory_used"]) / (1024 * 1024 * 1024), 2) } return result_human def cmd(self, command): status, res = self._milvus._cmd(command) logger.info("Server command: %s, result: %s" % (command, res)) self.check_status(status) return res