merge tests from 0.9.1 to master (#2432)

* merge tests from 0.9.1
Signed-off-by: Nzw <zw@milvus.io>

* update version
Signed-off-by: Nzw <zw@milvus.io>
Co-authored-by: Nzw <zw@milvus.io>

tests/milvus_python_test/test_collection.py → tests/milvus_python_test/collection/test_collection.py

@@ -1049,6 +1049,7 @@ class TestCollectionLogic(object):
             self.execute(logic_seq, connect)
         else:
             self.execute_with_error(logic_seq, connect)
+        self.tear_down(connect)
 
     def is_right(self, seq):
         if sorted(seq) == seq:
@@ -1089,6 +1090,11 @@ class TestCollectionLogic(object):
                 break
         assert error_flag == True
 
+    def tear_down(self, connect):
+        names = connect.list_collections()[1]
+        for name in names:
+            connect.drop_collection(name)
+
     def gen_params(self):
         collection_name = gen_unique_str("test_collection")
         top_k = 1

tests/milvus_python_test/conftest.py

@@ -9,10 +9,12 @@ from utils import *
 
 index_file_size = 10
 timeout = 1 
+delete_timeout = 60
 
 
 def pytest_addoption(parser):
     parser.addoption("--ip", action="store", default="localhost")
+    parser.addoption("--service", action="store", default="")
     parser.addoption("--port", action="store", default=19530)
     parser.addoption("--http-port", action="store", default=19121)
     parser.addoption("--handler", action="store", default="GRPC")
@@ -35,6 +37,7 @@ def check_server_connection(request):
 @pytest.fixture(scope="module")
 def connect(request):
     ip = request.config.getoption("--ip")
+    service_name = request.config.getoption("--service")
     port = request.config.getoption("--port")
     http_port = request.config.getoption("--http-port")
     handler = request.config.getoption("--handler")
@@ -47,7 +50,7 @@ def connect(request):
         pytest.exit("Milvus server can not connected, exit pytest ...")
     def fin():
         try:
-            # milvus.disconnect()
+            milvus.close()
             pass
         except Exception as e:
             logging.getLogger().info(str(e))
@@ -58,6 +61,7 @@ def connect(request):
 @pytest.fixture(scope="module")
 def dis_connect(request):
     ip = request.config.getoption("--ip")
+    service_name = request.config.getoption("--service")
     port = request.config.getoption("--port")
     http_port = request.config.getoption("--http-port")
     handler = request.config.getoption("--handler")
@@ -71,12 +75,13 @@ def dis_connect(request):
 @pytest.fixture(scope="module")
 def args(request):
     ip = request.config.getoption("--ip")
+    service_name = request.config.getoption("--service")
     port = request.config.getoption("--port")
     http_port = request.config.getoption("--http-port")
     handler = request.config.getoption("--handler")
     if handler == "HTTP":
         port = http_port
-    args = {"ip": ip, "port": port, "handler": handler}
+    args = {"ip": ip, "port": port, "handler": handler, "service_name": service_name}
     return args
 
 
@@ -108,10 +113,10 @@ def collection(request, connect):
         pytest.exit("collection can not be created, exit pytest ...")
 
     def teardown():
-        # status, collection_names = connect.list_collections()
-        # for collection_name in collection_names:
-        #     connect.drop_collection(collection_name)
-        connect.drop_collection(collection_name)
+        status, collection_names = connect.list_collections()
+        for collection_name in collection_names:
+            connect.drop_collection(collection_name, timeout=delete_timeout)
+        # connect.drop_collection(collection_name)
 
     request.addfinalizer(teardown)
 
@@ -135,10 +140,10 @@ def ip_collection(request, connect):
         pytest.exit("collection can not be created, exit pytest ...")
 
     def teardown():
-        # status, collection_names = connect.list_collections()
-        # for collection_name in collection_names:
-        #     connect.drop_collection(collection_name)
-        connect.drop_collection(collection_name)
+        status, collection_names = connect.list_collections()
+        for collection_name in collection_names:
+            connect.drop_collection(collection_name, timeout=delete_timeout)
+        # connect.drop_collection(collection_name)
 
     request.addfinalizer(teardown)
 
@@ -162,10 +167,10 @@ def jac_collection(request, connect):
         pytest.exit("collection can not be created, exit pytest ...")
 
     def teardown():
-        # status, collection_names = connect.list_collections()
-        # for collection_name in collection_names:
-        #     connect.drop_collection(collection_name)
-        connect.drop_collection(collection_name)
+        status, collection_names = connect.list_collections()
+        for collection_name in collection_names:
+            connect.drop_collection(collection_name, timeout=delete_timeout)
+        # connect.drop_collection(collection_name)
 
     request.addfinalizer(teardown)
 
@@ -188,10 +193,10 @@ def ham_collection(request, connect):
         pytest.exit("collection can not be created, exit pytest ...")
 
     def teardown():
-        # status, collection_names = connect.list_collections()
-        # for collection_name in collection_names:
-        #     connect.drop_collection(collection_name)
-        connect.drop_collection(collection_name)
+        status, collection_names = connect.list_collections()
+        for collection_name in collection_names:
+            connect.drop_collection(collection_name, timeout=delete_timeout)
+        # connect.drop_collection(collection_name)
 
     request.addfinalizer(teardown)
 
@@ -214,10 +219,10 @@ def tanimoto_collection(request, connect):
         pytest.exit("collection can not be created, exit pytest ...")
 
     def teardown():
-        # status, collection_names = connect.list_collections()
-        # for collection_name in collection_names:
-        #     connect.drop_collection(collection_name)
-        connect.drop_collection(collection_name)
+        status, collection_names = connect.list_collections()
+        for collection_name in collection_names:
+            connect.drop_collection(collection_name, timeout=delete_timeout)
+        # connect.drop_collection(collection_name)
 
     request.addfinalizer(teardown)
     return collection_name
@@ -239,10 +244,10 @@ def substructure_collection(request, connect):
         pytest.exit("collection can not be created, exit pytest ...")
 
     def teardown():
-        # status, collection_names = connect.list_collections()
-        # for collection_name in collection_names:
-        #     connect.drop_collection(collection_name)
-        connect.drop_collection(collection_name)
+        status, collection_names = connect.list_collections()
+        for collection_name in collection_names:
+            connect.drop_collection(collection_name, timeout=delete_timeout)
+        # connect.drop_collection(collection_name)
 
     request.addfinalizer(teardown)
     return collection_name
@@ -264,10 +269,10 @@ def superstructure_collection(request, connect):
         pytest.exit("collection can not be created, exit pytest ...")
 
     def teardown():
-        # status, collection_names = connect.list_collections()
-        # for collection_name in collection_names:
-        #     connect.drop_collection(collection_name)
-        connect.drop_collection(collection_name)
+        status, collection_names = connect.list_collections()
+        for collection_name in collection_names:
+            connect.drop_collection(collection_name, timeout=delete_timeout)
+        # connect.drop_collection(collection_name)
 
     request.addfinalizer(teardown)
     return collection_name

tests/milvus_python_test/pytest.ini

 [pytest]
 log_format = [%(asctime)s-%(levelname)s-%(name)s]: %(message)s (%(filename)s:%(lineno)s)
+log_date_format = %Y-%m-%d %H:%M:%S
 
 log_cli = true
 log_level = 20
 
-timeout = 600
+timeout = 360
 
 markers = 
     level: test level

tests/milvus_python_test/requirements.txt

@@ -9,3 +9,4 @@ pytest-level==0.1.1
 pytest-xdist==1.23.2
 scikit-learn>=0.19.1
 pymilvus-test>=0.2.0
+kubernetes==10.0.1

tests/milvus_python_test/test_compact.py

@@ -636,7 +636,6 @@ class TestCompactJAC:
                      'index_file_size': index_file_size,
                      'metric_type': MetricType.JACCARD}
             connect.create_collection(param)
-        time.sleep(6)
         for i in range(num_collections):
             status, ids = connect.insert(collection_name=collection_list[i], records=vectors)
             assert status.OK()
@@ -646,6 +645,8 @@ class TestCompactJAC:
             assert status.OK()
             status = connect.compact(collection_list[i])
             assert status.OK()
+            status = connect.drop_collection(collection_list[i])
+            assert status.OK()
 
     @pytest.mark.timeout(COMPACT_TIMEOUT)
     def test_add_vector_after_compact(self, connect, jac_collection):

tests/milvus_python_test/test_partition.py

@@ -12,9 +12,9 @@ from utils import *
 dim = 128
 index_file_size = 10
 collection_id = "test_partition"
-ADD_TIMEOUT = 60
 nprobe = 1
 tag = "1970-01-01"
+TIMEOUT = 120
 
 
 class TestCreateBase:
@@ -33,7 +33,9 @@ class TestCreateBase:
         status = connect.create_partition(collection, tag)
         assert status.OK()
 
-    def _test_create_partition_limit(self, connect, collection):
+    @pytest.mark.level(3)
+    @pytest.mark.timeout(TIMEOUT)
+    def test_create_partition_limit(self, connect, collection):
         '''
         target: test create partitions, check status returned
         method: call function: create_partition for 4097 times

tests/milvus_python_test/test_restart.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
+collection_id = "test_partition_restart"
+nprobe = 1
+tag = "1970-01-01"
+
+
+class TestRestartBase:
+
+    """
+    ******************************************************************
+      The following cases are used to test `create_partition` function 
+    ******************************************************************
+    """
+    @pytest.fixture(scope="function", autouse=True)
+    def skip_check(self, connect, args):
+        if args["service_name"].find("shards") != -1:
+            reason = "Skip restart cases in shards mode"
+            logging.getLogger().info(reason)
+            pytest.skip(reason)
+
+
+    @pytest.mark.level(2)
+    def test_create_partition_insert_restart(self, connect, collection, args):
+        '''
+        target: return the same row count after server restart
+        method: call function: create partition, then insert, restart server and assert row count
+        expected: status ok, and row count keep the same
+        '''
+        status = connect.create_partition(collection, tag)
+        assert status.OK()
+        nq = 1000
+        vectors = gen_vectors(nq, dim)
+        ids = [i for i in range(nq)]
+        status, ids = connect.insert(collection, vectors, ids, partition_tag=tag)
+        assert status.OK()
+        status = connect.flush([collection])
+        assert status.OK()
+        status, res = connect.count_entities(collection)
+        logging.getLogger().info(res)
+        assert res == nq
+
+        # restart server
+        if restart_server(args["service_name"]):
+            logging.getLogger().info("Restart success") 
+        else:
+            logging.getLogger().info("Restart failed")
+        # assert row count again
+
+        # debug
+        new_connect = get_milvus(args["ip"], args["port"], handler=args["handler"]) 
+        status, res = new_connect.count_entities(collection)
+        logging.getLogger().info(status)
+        logging.getLogger().info(res)
+        assert status.OK()
+        assert res == nq

tests/milvus_python_test/utils.py

-# STL imports
+import os
+import sys
 import random
 import string
 import struct
-import sys
 import logging
 import time, datetime
 import copy
@@ -625,3 +625,55 @@ def assert_equal_vector(v1, v2):
         assert False
     for i in range(len(v1)):
         assert abs(v1[i] - v2[i]) < epsilon
+
+
+def restart_server(helm_release_name):
+    res = True
+    timeout = 120
+    from kubernetes import client, config
+    client.rest.logger.setLevel(logging.WARNING)
+
+    namespace = "milvus"
+    # service_name = "%s.%s.svc.cluster.local" % (helm_release_name, namespace)
+    config.load_kube_config()
+    v1 = client.CoreV1Api()
+    pod_name = None
+    # config_map_names = v1.list_namespaced_config_map(namespace, pretty='true')
+    # body = {"replicas": 0}
+    pods = v1.list_namespaced_pod(namespace)
+    for i in pods.items:
+        if i.metadata.name.find(helm_release_name) != -1 and i.metadata.name.find("mysql") == -1:
+            pod_name = i.metadata.name
+            break
+            # v1.patch_namespaced_config_map(config_map_name, namespace, body, pretty='true')
+    # status_res = v1.read_namespaced_service_status(helm_release_name, namespace, pretty='true')
+    # print(status_res)
+    if pod_name is not None:
+        try:
+            v1.delete_namespaced_pod(pod_name, namespace)
+        except Exception as e:
+            logging.error(str(e))
+            logging.error("Exception when calling CoreV1Api->delete_namespaced_pod")
+            res = False
+            return res
+        time.sleep(5)
+        # check if restart successfully
+        pods = v1.list_namespaced_pod(namespace)
+        for i in pods.items:
+            pod_name_tmp = i.metadata.name
+            if pod_name_tmp.find(helm_release_name) != -1:
+                logging.debug(pod_name_tmp)
+                start_time = time.time()
+                while time.time() - start_time > timeout:
+                    status_res = v1.read_namespaced_pod_status(pod_name_tmp, namespace, pretty='true')
+                    if status_res.status.phase == "Running":
+                        break
+                    time.sleep(1)
+                if time.time() - start_time > timeout:
+                    logging.error("Restart pod: %s timeout" % pod_name_tmp)
+                    res = False
+                    return res
+    else:
+        logging.error("Pod: %s not found" % helm_release_name)
+        res = False
+    return res