From ca4ad826af50104d49c18ec0efe6563a5d38f503 Mon Sep 17 00:00:00 2001
From: zhuwenxing <wenxing.zhu@zilliz.com>
Date: Thu, 12 Oct 2023 22:37:38 -0500
Subject: [PATCH] [test]Add more partitions in br test (#203)

Signed-off-by: zhuwenxing <wenxing.zhu@zilliz.com>
---
 tests/base/client_base.py              |  3 +-
 tests/testcases/test_restore_backup.py | 61 +++++++++++++++++++++++++-
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/tests/base/client_base.py b/tests/base/client_base.py
index 74b4a9e9..278bf59c 100644
--- a/tests/base/client_base.py
+++ b/tests/base/client_base.py
@@ -264,7 +264,8 @@ def prepare_data(self, name=None, db_name="default", nb=ct.default_nb, dim=ct.de
         collection_w = self.init_collection_wrap(name=name, schema=default_schema, active_trace=True)
         # create partitions
         if enable_partition:
-            partition_w = self.init_partition_wrap(collection_wrap=collection_w)
+            for i in range(5):
+                partition_w = self.init_partition_wrap(collection_wrap=collection_w)
         assert collection_w.name == name
         if nb > 0:
             cf.insert_data(collection_w, nb=nb, is_binary=is_binary, auto_id=auto_id, dim=dim)
diff --git a/tests/testcases/test_restore_backup.py b/tests/testcases/test_restore_backup.py
index f5e02c91..5b21dd4c 100644
--- a/tests/testcases/test_restore_backup.py
+++ b/tests/testcases/test_restore_backup.py
@@ -23,7 +23,7 @@ class TestRestoreBackup(TestcaseBase):
 
     @pytest.mark.parametrize("nb", [3000])
     @pytest.mark.parametrize("is_auto_id", [True, False])
-    @pytest.mark.parametrize("enable_partition", [True, False])
+    @pytest.mark.parametrize("enable_partition", [False])
     @pytest.mark.parametrize("is_async", [True, False])
     @pytest.mark.parametrize("collection_need_to_restore", [1, 2, 3])
     @pytest.mark.parametrize("collection_type", ["binary", "float", "all"])
@@ -79,6 +79,65 @@ def test_milvus_restore_back(self, collection_type, collection_need_to_restore,
         for name in restore_collections:
             self.compare_collections(name, name+suffix)
 
+    @pytest.mark.tags(CaseLabel.L1)
+    @pytest.mark.parametrize("nb", [3000])
+    @pytest.mark.parametrize("is_auto_id", [True])
+    @pytest.mark.parametrize("enable_partition", [True])
+    @pytest.mark.parametrize("is_async", [True])
+    @pytest.mark.parametrize("collection_need_to_restore", [3])
+    @pytest.mark.parametrize("collection_type", ["all"])
+    def test_milvus_restore_back_with_multi_partition(self, collection_type, collection_need_to_restore, is_async, is_auto_id, enable_partition, nb):
+        # prepare data
+        names_origin = []
+        back_up_name = cf.gen_unique_str(backup_prefix)
+        if collection_type == "all":
+            for is_binary in [True, False, False]:
+                names_origin.append(cf.gen_unique_str(prefix))
+                self.prepare_data(names_origin[-1], nb=nb, is_binary=is_binary, auto_id=is_auto_id, check_function=False, enable_partition=enable_partition)
+        if collection_type == "float":
+            for is_binary in [False, False, False]:
+                names_origin.append(cf.gen_unique_str(prefix))
+                self.prepare_data(names_origin[-1], nb=nb, is_binary=is_binary, auto_id=is_auto_id, check_function=False, enable_partition=enable_partition)
+        if collection_type == "binary":
+            for is_binary in [True, True, True]:
+                names_origin.append(cf.gen_unique_str(prefix))
+                self.prepare_data(names_origin[-1], nb=nb, is_binary=is_binary, auto_id=is_auto_id, check_function=False, enable_partition=enable_partition)
+        log.info(f"name_origin:{names_origin}, back_up_name: {back_up_name}")
+        for name in names_origin:
+            res, _ = self.utility_wrap.has_collection(name)
+            assert res is True
+        # create backup
+
+        names_need_backup = names_origin
+        payload = {"async": False, "backup_name": back_up_name, "collection_names": names_need_backup}
+        res = client.create_backup(payload)
+        log.info(f"create backup response: {res}")
+        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"]]
+        restore_collections = backup_collections
+        if collection_need_to_restore == "all":
+            payload = {"async": False, "backup_name": back_up_name,
+                       "collection_suffix": suffix}
+        else:
+            restore_collections = names_need_backup[:collection_need_to_restore]
+            payload = {"async": False, "backup_name": back_up_name,
+                       "collection_suffix": suffix, "collection_names": restore_collections}
+        t0 = time.time()
+        res = client.restore_backup(payload)
+        restore_id = res["data"]["id"]
+        log.info(f"restore_backup: {res}")
+        if is_async:
+            res = client.wait_restore_complete(restore_id)
+            assert res is True
+        t1 = time.time()
+        log.info(f"restore {restore_collections} cost time: {t1 - t0}")
+        res, _ = self.utility_wrap.list_collections()
+        for name in restore_collections:
+            assert name + suffix in res
+        for name in restore_collections:
+            self.compare_collections(name, name+suffix)
+    
     @pytest.mark.tags(CaseLabel.L1)
     def test_milvus_restore_back_with_db_support(self):
         # prepare data