diff --git a/tests/python_client/deploy/scripts/action_after_reinstall.py b/tests/python_client/deploy/scripts/action_after_reinstall.py
index df490bb830..0fbcde7edf 100644
--- a/tests/python_client/deploy/scripts/action_after_reinstall.py
+++ b/tests/python_client/deploy/scripts/action_after_reinstall.py
@@ -2,7 +2,7 @@ from pymilvus import connections
 from utils import *
 
 
-def task_1():
+def task_1(data_size):
     """
     task_1:
         before reinstall: create collection and insert data, load and search
@@ -12,10 +12,10 @@ def task_1():
     connections.connect(host="127.0.0.1", port=19530, timeout=60)
     get_collections(prefix)
     load_and_search(prefix)
-    create_collections_and_insert_data(prefix)
+    create_collections_and_insert_data(prefix, data_size)
 
 
-def task_2():
+def task_2(data_zise):
     """
     task_2:
         before reinstall: create collection, insert data and create index, load and search
@@ -25,11 +25,17 @@ def task_2():
     connections.connect(host="127.0.0.1", port=19530, timeout=60)
     get_collections(prefix)
     load_and_search(prefix)
-    create_collections_and_insert_data(prefix)
+    create_collections_and_insert_data(prefix, data_zise)
     create_index(prefix)
     load_and_search(prefix)
 
 
 if __name__ == '__main__':
-    task_1()
-    task_2()
\ No newline at end of file
+    import argparse
+    parser = argparse.ArgumentParser(description='config for deploy test')
+    parser.add_argument('--data_size', type=int, default=3000, help='data size')
+    args = parser.parse_args()
+    data_size = args.data_size
+    print(f"data size: {data_size}")    
+    task_1(data_size)
+    task_2(data_size)
\ No newline at end of file
diff --git a/tests/python_client/deploy/scripts/action_after_upgrade.py b/tests/python_client/deploy/scripts/action_after_upgrade.py
index e664812501..ae21fcd47b 100644
--- a/tests/python_client/deploy/scripts/action_after_upgrade.py
+++ b/tests/python_client/deploy/scripts/action_after_upgrade.py
@@ -2,7 +2,7 @@ from pymilvus import connections
 from utils import *
 
 
-def task_1():
+def task_1(data_size):
     """
     task_1:
         before upgrade: create collection and insert data, load and search
@@ -12,12 +12,12 @@ def task_1():
     connections.connect(host="127.0.0.1", port=19530, timeout=60)
     get_collections(prefix)
     load_and_search(prefix)
-    create_collections_and_insert_data(prefix)
+    create_collections_and_insert_data(prefix, data_size)
     create_index(prefix)
     load_and_search(prefix)
 
 
-def task_2():
+def task_2(data_size):
     """
     task_2:
         before upgrade: create collection, insert data and create index, load and search
@@ -27,11 +27,17 @@ def task_2():
     connections.connect(host="127.0.0.1", port=19530, timeout=60)
     get_collections(prefix)
     load_and_search(prefix)
-    create_collections_and_insert_data(prefix)
+    create_collections_and_insert_data(prefix, data_size)
     create_index(prefix)
     load_and_search(prefix)
 
 
 if __name__ == '__main__':
-    task_1()
-    task_2()
\ No newline at end of file
+    import argparse
+    parser = argparse.ArgumentParser(description='config for deploy test')
+    parser.add_argument('--data_size', type=int, default=3000, help='data size')
+    args = parser.parse_args()
+    data_size = args.data_size
+    print(f"data size: {data_size}")
+    task_1(data_size)
+    task_2(data_size)
\ No newline at end of file
diff --git a/tests/python_client/deploy/scripts/utils.py b/tests/python_client/deploy/scripts/utils.py
index 88f4821344..59715bf0ca 100644
--- a/tests/python_client/deploy/scripts/utils.py
+++ b/tests/python_client/deploy/scripts/utils.py
@@ -64,7 +64,7 @@ def get_collections(prefix):
     return col_list
 
 
-def create_collections_and_insert_data(prefix, count):
+def create_collections_and_insert_data(prefix, count=3000):
     import random
     dim = 128
     nb = count // 10