From 7e286222a67371f22570c1494fcd17ff33a87537 Mon Sep 17 00:00:00 2001
From: zhuwenxing <wenxing.zhu@zilliz.com>
Date: Mon, 6 Dec 2021 21:12:01 +0800
Subject: [PATCH] [skip ci]Update load balance testcase (#12818)

Signed-off-by: zhuwenxing <wenxing.zhu@zilliz.com>
---
 .../loadbalance/test_auto_load_balance.py     | 46 +++++++++++++------
 1 file changed, 31 insertions(+), 15 deletions(-)

diff --git a/tests/python_client/loadbalance/test_auto_load_balance.py b/tests/python_client/loadbalance/test_auto_load_balance.py
index 5f35f2c16a..09ea205305 100644
--- a/tests/python_client/loadbalance/test_auto_load_balance.py
+++ b/tests/python_client/loadbalance/test_auto_load_balance.py
@@ -1,11 +1,12 @@
 from time import sleep
-from pymilvus import connections
+from pymilvus import connections, list_collections, utility
 from chaos.checker import (CreateChecker, InsertFlushChecker,
                            SearchChecker, QueryChecker, IndexChecker, Op, assert_statistic)
 from common.milvus_sys import MilvusSys
 from utils.util_log import test_log as log
 from chaos import chaos_commons as cc
 from common.common_type import CaseLabel
+from common import common_func as cf
 from chaos import constants
 from customize.milvus_operator import MilvusOperator
 
@@ -15,9 +16,10 @@ namespace = "chaos-testing"
 
 
 def install_milvus(release_name):
-    cus_configs = {'spec.components.image': 'milvusdb/milvus-dev:master-latest',
+    cus_configs = {'spec.components.image': 'milvusdb/milvus-dev:master-20211206-b20a238',
                    'metadata.namespace': namespace,
-                   'metadata.name': release_name
+                   'metadata.name': release_name,
+                   'spec.components.proxy.serviceType': 'LoadBalancer'
                    }
     milvus_op = MilvusOperator()
     log.info(f"install milvus with configs: {cus_configs}")
@@ -35,14 +37,12 @@ def install_milvus(release_name):
 
 
 def scale_up_milvus(release_name):
-    cus_configs = {'spec.components.image': 'milvusdb/milvus-dev:master-latest',
-                   'metadata.namespace': namespace,
-                   'metadata.name': release_name,
+    cus_configs = {
                    'spec.components.queryNode.replicas': 2
-                   }
+                  }
     milvus_op = MilvusOperator()
     log.info(f"scale up milvus with configs: {cus_configs}")
-    milvus_op.upgrade(release_name, cus_configs)
+    milvus_op.upgrade(release_name, cus_configs, namespace=namespace)
     healthy = milvus_op.wait_for_healthy(release_name, namespace, timeout=1200)
     log.info(f"milvus healthy: {healthy}")
     if healthy:
@@ -58,9 +58,9 @@ def scale_up_milvus(release_name):
 class TestAutoLoadBalance(object):
 
 
-    # def teardown_method(self):
-    #     milvus_op = MilvusOperator()
-    #     milvus_op.uninstall(self.release_name, namespace)
+    def teardown_method(self):
+        milvus_op = MilvusOperator()
+        milvus_op.uninstall(self.release_name, namespace)
 
     def test_auto_load_balance(self):
         """
@@ -82,20 +82,36 @@ class TestAutoLoadBalance(object):
         }
         cc.start_monitor_threads(self.health_checkers)
         # wait  
-        sleep(constants.WAIT_PER_OP * 4)
-
+        sleep(constants.WAIT_PER_OP * 10)
+        all_collections = list_collections()
+        for c in all_collections:
+            seg_info = utility.get_query_segment_info(c)
+            seg_distribution = cf.get_segment_distribution(seg_info)
+            for k in seg_distribution.keys():
+                log.info(f"collection {c}'s segment distribution in node {k} is {seg_distribution[k]['sealed']}")
+                log.info(f"collection {c}'s growing segment distribution in node {k} is {seg_distribution[k]['growing']}")
         # first assert
+        log.info("first assert")
         assert_statistic(self.health_checkers)
 
         # scale up
+        log.info("scale up milvus")
         scale_up_milvus(self.release_name)
         # reset counting
         cc.reset_counting(self.health_checkers)
-        sleep(constants.WAIT_PER_OP * 4)
+        sleep(constants.WAIT_PER_OP * 10)
+        all_collections = list_collections()
+        for c in all_collections:
+            seg_info = utility.get_query_segment_info(c)
+            seg_distribution = cf.get_segment_distribution(seg_info)
+            for k in seg_distribution.keys():
+                log.info(f"collection {c}'s sealed segment distribution in node {k} is {seg_distribution[k]['sealed']}")
+                log.info(f"collection {c}'s growing segment distribution in node {k} is {seg_distribution[k]['growing']}")
         # second assert
+        log.info("second assert")
         assert_statistic(self.health_checkers)
 
         # TODO assert segment distribution
-        
+
         # assert all expectations
         assert_expectations()