add new feature test (#134)

Signed-off-by: zhuwenxing <[email protected]>

Author: zhuwenxing

tests/base/client_base.py

@@ -1,5 +1,5 @@
 import sys
-from pymilvus import DefaultConfig, DataType
+from pymilvus import DefaultConfig, DataType, db
 
 sys.path.append("..")
 from base.connections_wrapper import ApiConnectionsWrapper
@@ -248,12 +248,13 @@ def collection_insert_multi_segments_one_shard(self, collection_prefix, num_of_s
             assert collection_w.num_entities == nb_of_segment * (i + 1)
         return collection_w
 
-    def prepare_data(self, name=None, nb=ct.default_nb, dim=ct.default_dim, is_binary=False, auto_id=False,
+    def prepare_data(self, name=None, db_name="default", nb=ct.default_nb, dim=ct.default_dim, is_binary=False, auto_id=False,
                      primary_field=ct.default_int64_field_name, is_flushed=True, check_function=False):
         """
         prepare data for test case
         """
         self._connect()
+        db.using_database(db_name)
         prefix = "backup_e2e_"
         name = cf.gen_unique_str(prefix) if name is None else name
         default_schema = cf.gen_default_collection_schema(auto_id=auto_id, dim=dim, primary_field=primary_field)
@@ -333,7 +334,9 @@ def is_binary_by_schema(self, schema):
             if field.dtype == DataType.FLOAT_VECTOR:
                 return False
 
-    def compare_collections(self, src_name, dist_name):
+    def compare_collections(self, src_name, dist_name, output_fields=None):
+        if output_fields is None:
+            output_fields = [ct.default_int64_field_name, ct.default_json_field_name]
         collection_src, _ = self.collection_wrap.init_collection(name=src_name)
         collection_dist, _ = self.collection_wrap.init_collection(name=dist_name)
         assert collection_src.num_entities == collection_dist.num_entities, \
@@ -349,10 +352,12 @@ def compare_collections(self, src_name, dist_name):
             else:
                 coll.create_index(ct.default_float_vec_field_name, ct.default_index, index_name=cf.gen_unique_str())
             coll.load()
-        src_res = collection_src.query(expr=f'{ct.default_int64_field_name} > 0',
-                                       output_fields=[ct.default_int64_field_name, ct.default_json_field_name])
-        dist_res = collection_dist.query(expr=f'{ct.default_int64_field_name} > 0',
-                                         output_fields=[ct.default_int64_field_name, ct.default_json_field_name])
+        src_res = collection_src.query(expr=f'{ct.default_int64_field_name} >= 0',
+                                       output_fields=output_fields)
+        log.info(f"src res: {len(src_res)}")
+        dist_res = collection_dist.query(expr=f'{ct.default_int64_field_name} >= 0',
+                                         output_fields=output_fields)
+        log.info(f"dist res: {len(dist_res)}")
         assert len(dist_res) == len(src_res)
 
     def check_collection_binary(self, name):

tests/common/common_func.py

@@ -160,9 +160,9 @@ def gen_collection_schema_all_datatype(description=ct.default_desc,
     return schema
 
 
-def gen_collection_schema(fields, primary_field=None, description=ct.default_desc, auto_id=False):
+def gen_collection_schema(fields, primary_field=None, description=ct.default_desc, auto_id=False, **kwargs):
     schema, _ = ApiCollectionSchemaWrapper().init_collection_schema(fields=fields, primary_field=primary_field,
-                                                                    description=description, auto_id=auto_id)
+                                                                    description=description, auto_id=auto_id, **kwargs)
     return schema
 
 
@@ -319,7 +319,6 @@ def gen_default_binary_dataframe_data(nb=ct.default_nb, dim=ct.default_dim, star
     })
     return df, binary_raw_values
 
-
 def gen_default_list_data(nb=ct.default_nb, dim=ct.default_dim):
     int_values = [i for i in range(nb)]
     float_values = [np.float32(i) for i in range(nb)]

tests/requirements.txt

@@ -17,7 +17,7 @@ pytest-print==0.2.1
 pytest-level==0.1.1
 pytest-xdist==2.5.0
 pytest-loguru==0.2.0
-pymilvus==2.2.9.dev14
+pymilvus==2.2.9.dev18
 pytest-rerunfailures==9.1.1
 git+https://github.com/Projectplace/pytest-tags
 ndg-httpsclient

tests/testcases/test_restore_backup.py

@@ -1,6 +1,7 @@
 import time
 import pytest
-
+import numpy as np
+from pymilvus import db
 from base.client_base import TestcaseBase
 from common import common_func as cf
 from common import common_type as ct
@@ -75,3 +76,123 @@ def test_milvus_restore_back(self, collection_type, collection_need_to_restore,
         for name in restore_collections:
             self.compare_collections(name, name+suffix)
 
+    def test_milvus_restore_back_with_db_support(self):
+        # prepare data
+        self._connect()
+        names_origin = []
+        for i in range(2):
+            db_name = cf.gen_unique_str("db")
+            db.create_database(db_name)
+            db.using_database(db_name)
+            collection_name = cf.gen_unique_str(prefix)
+            self.prepare_data(name=collection_name, db_name=db_name, nb=3000, is_binary=False, auto_id=True, check_function=False)
+            assert collection_name in self.utility_wrap.list_collections()[0]
+            names_origin.append(f"{db_name}.{collection_name}")
+        log.info(f"name_origin:{names_origin}")
+        # create backup
+        back_up_name = cf.gen_unique_str(backup_prefix)
+        payload = {
+            "async": False,
+            "backup_name": back_up_name,
+            "collection_names": names_origin,
+        }
+        log.info(f"payload: {payload}")
+        res = client.create_backup(payload)
+        log.info(f"create backup response: {res}")
+        res = client.list_backup()
+        log.info(f"list_backup {res}")
+        if "data" in res:
+            all_backup = [r["name"] for r in res["data"]]
+        else:
+            all_backup = []
+        assert back_up_name in all_backup
+        payload = {"async": False, "backup_name": back_up_name,
+                   # "collection_names": names_origin,
+                   "collection_suffix": suffix}
+        log.info(f"restore payload: {payload}")
+        res = client.restore_backup(payload)
+        log.info(f"restore_backup: {res}")
+        for name in names_origin:
+            db_name = name.split(".")[0]
+            collection_name = name.split(".")[1]
+            db.using_database(db_name)
+
+            res, _ = self.utility_wrap.list_collections()
+            log.info(f"collection list in db {db_name}: {res}")
+            assert collection_name + suffix in res
+            self.compare_collections(collection_name, collection_name + suffix)
+
+    @pytest.mark.parametrize("include_partition_key", [True, False])
+    @pytest.mark.parametrize("include_dynamic", [True, False])
+    @pytest.mark.parametrize("include_json", [True, False])
+    def test_milvus_restore_back_with_new_feature_support(self, include_json, include_dynamic, include_partition_key):
+        self._connect()
+        name_origin = cf.gen_unique_str(prefix)
+        back_up_name = cf.gen_unique_str(backup_prefix)
+        if include_json:
+            fields = [cf.gen_int64_field(name="int64", is_primary=True),
+                      cf.gen_int64_field(name="key"),
+                      cf.gen_json_field(name="json"),
+                      cf.gen_float_vec_field(name="float_vector", dim=128),
+                      ]
+        else:
+            fields = [cf.gen_int64_field(name="int64", is_primary=True),
+                      cf.gen_int64_field(name="key"),
+                      cf.gen_float_vec_field(name="float_vector", dim=128),
+                      ]
+        if include_partition_key:
+            partition_key = "key"
+            default_schema = cf.gen_collection_schema(fields,
+                                                      enable_dynamic_field=include_dynamic,
+                                                      partition_key_field=partition_key)
+        else:
+            default_schema = cf.gen_collection_schema(fields,
+                                                      enable_dynamic_field=include_dynamic)
+
+        collection_w = self.init_collection_wrap(name=name_origin, schema=default_schema, active_trace=True)
+        nb = 3000
+        if include_json:
+            data = [
+                [i for i in range(nb)],
+                [i % 3 for i in range(nb)],
+                [{f"key_{str(i)}": i} for i in range(nb)],
+                [[np.float32(i) for i in range(128)] for _ in range(nb)],
+            ]
+        else:
+            data = [
+                [i for i in range(nb)],
+                [i % 3 for i in range(nb)],
+                [[np.float32(i) for i in range(128)] for _ in range(nb)],
+            ]
+        collection_w.insert(data=data)
+
+        res = client.create_backup({"async": False, "backup_name": back_up_name, "collection_names": [name_origin]})
+        log.info(f"create_backup {res}")
+        res = client.list_backup()
+        log.info(f"list_backup {res}")
+        if "data" in res:
+            all_backup = [r["name"] for r in res["data"]]
+        else:
+            all_backup = []
+        assert back_up_name in all_backup
+        backup = client.get_backup(back_up_name)
+        assert backup["data"]["name"] == back_up_name
+        backup_collections = [backup["collection_name"]for backup in backup["data"]["collection_backups"]]
+        assert name_origin in backup_collections
+        res = client.restore_backup({"async": False, "backup_name": back_up_name, "collection_names": [name_origin],
+                                     "collection_suffix": suffix})
+        log.info(f"restore_backup: {res}")
+        res, _ = self.utility_wrap.list_collections()
+        assert name_origin + suffix in res
+        output_fields = None
+        if not include_json:
+            output_fields = [ct.default_int64_field_name]
+        self.compare_collections(name_origin, name_origin + suffix, output_fields=output_fields)
+        res = client.delete_backup(back_up_name)
+        res = client.list_backup()
+        if "data" in res:
+            all_backup = [r["name"] for r in res["data"]]
+        else:
+            all_backup = []
+        assert back_up_name not in all_backup
+