From add51f6c24fe1ca4b7590906b59dfbaa46658324 Mon Sep 17 00:00:00 2001
From: jingkl <34296482+jingkl@users.noreply.github.com>
Date: Mon, 20 Dec 2021 19:31:44 +0800
Subject: [PATCH] Modify the testcase with search (#13779)

Signed-off-by: jingkl <jingjing.jia@zilliz.com>
---
 tests/python_client/base/client_base.py       | 47 +------------------
 .../python_client/testcases/test_search_20.py |  4 +-
 tests/python_client/testcases/test_utility.py |  4 +-
 3 files changed, 6 insertions(+), 49 deletions(-)

diff --git a/tests/python_client/base/client_base.py b/tests/python_client/base/client_base.py
index 5be54d2239..b99b8c100f 100644
--- a/tests/python_client/base/client_base.py
+++ b/tests/python_client/base/client_base.py
@@ -133,50 +133,7 @@ class TestcaseBase(Base):
                                       check_task=check_task, check_items=check_items,
                                       **kwargs)
         return partition_wrap
-
-    def init_collection_without_load(self, prefix, insert_data=False, nb=ct.default_nb,
-                                partition_num=0, is_binary=False, is_all_data_type=False,
-                                auto_id=False, dim=ct.default_dim, is_index=False):
-        """
-        target: create specified collections
-        method: 1. create collections (binary/non-binary, default/all data type, auto_id or not)
-                2. create partitions if specified
-                3. insert specified (binary/non-binary, default/all data type) data
-                   into each partition if any
-        expected: return collection and raw data, insert ids
-        """
-        log.info("Test case of search interface: initialize before test case")
-        self._connect()
-        collection_name = cf.gen_unique_str(prefix)
-        vectors = []
-        binary_raw_vectors = []
-        insert_ids = []
-        time_stamp = 0
-        # 1 create collection
-        default_schema = cf.gen_default_collection_schema(auto_id=auto_id, dim=dim)
-        if is_binary:
-            default_schema = cf.gen_default_binary_collection_schema(auto_id=auto_id, dim=dim)
-        if is_all_data_type:
-            default_schema = cf.gen_collection_schema_all_datatype(auto_id=auto_id, dim=dim)
-        log.info("init_collection_general: collection creation")
-        collection_w = self.init_collection_wrap(name=collection_name,
-                                                 schema=default_schema)
-        # 2 add extra partitions if specified (default is 1 partition named "_default")
-        if partition_num > 0:
-            cf.gen_partitions(collection_w, partition_num)
-            
-        # 3 insert data if specified
-        if insert_data:
-            collection_w, vectors, binary_raw_vectors, insert_ids, time_stamp = \
-                cf.insert_data(collection_w, nb, is_binary, is_all_data_type, auto_id=auto_id, dim=dim)
-            assert collection_w.is_empty is False
-            assert collection_w.num_entities == nb
-            log.info("insert_data: inserted data into collection %s (num_entities: %s)"
-                     % (collection_w.name, nb))
-
-        return collection_w, vectors, binary_raw_vectors, insert_ids, time_stamp
-       
-
+ 
     def init_collection_general(self, prefix, insert_data=False, nb=ct.default_nb,
                                 partition_num=0, is_binary=False, is_all_data_type=False,
                                 auto_id=False, dim=ct.default_dim, is_index=False):
@@ -237,7 +194,7 @@ class TestcaseBase(Base):
         df_default = cf.gen_default_dataframe_data(nb=half, start=half)
         collection_w.insert(df_default)
         conn.flush([collection_w.name])
-        collection_w.load()
+        collection_w.load(partition_names=[partition_w.name, "_default"])
         return collection_w, partition_w,df_partition, df_default
    
 
diff --git a/tests/python_client/testcases/test_search_20.py b/tests/python_client/testcases/test_search_20.py
index 58afac61db..e9ecf27872 100644
--- a/tests/python_client/testcases/test_search_20.py
+++ b/tests/python_client/testcases/test_search_20.py
@@ -486,7 +486,7 @@ class TestCollectionSearchInvalid(TestcaseBase):
         """
         # 1. initialize with data
         partition_num = 1
-        collection_w = self.init_collection_without_load(prefix, True, 10, partition_num)[0]
+        collection_w = self.init_collection_general(prefix, True, 10, partition_num, is_index=True)[0]
         par = collection_w.partitions
         par_name = par[partition_num].name
         # 2. release partition
@@ -1157,7 +1157,7 @@ class TestCollectionSearch(TestcaseBase):
     #     expected: search successfully
     #     """
     #     # 1. initialize without data
-    #     collection_w, _, _, insert_ids, time_stamp = self.init_collection_without_load(prefix, True, nb,
+    #     collection_w, _, _, insert_ids, time_stamp = self.init_collection_general(prefix, True, nb,
     #                                                                               1, auto_id=auto_id,
     #                                                                               dim=dim)[0:5]
     #     # 2. release collection
diff --git a/tests/python_client/testcases/test_utility.py b/tests/python_client/testcases/test_utility.py
index b51b6595f5..2802898db0 100644
--- a/tests/python_client/testcases/test_utility.py
+++ b/tests/python_client/testcases/test_utility.py
@@ -832,8 +832,8 @@ class TestUtilityBase(TestcaseBase):
         """
         half = ct.default_nb
         # insert entities into two partitions, collection flush and load
-        collection_w, _, partition_e, _ = self.insert_entities_into_two_partitions(half)
-        partition_e.release()
+        collection_w, partition_w, _, _ = self.insert_entities_into_two_partitions_in_half(half)
+        partition_w.release()
         res = self.utility_wrap.loading_progress(collection_w.name)[0]
         assert res[num_total_entities] == half * 2
         assert res[num_loaded_entities] == half