diff --git a/internal/datacoord/server_test.go b/internal/datacoord/server_test.go
index cb9070cd9e..5d8a4537b5 100644
--- a/internal/datacoord/server_test.go
+++ b/internal/datacoord/server_test.go
@@ -4756,10 +4756,10 @@ func TestDataNodeTtChannel(t *testing.T) {
 func TestUpdateAutoBalanceConfigLoop(t *testing.T) {
 	Params.Save(Params.DataCoordCfg.CheckAutoBalanceConfigInterval.Key, "1")
 	defer Params.Reset(Params.DataCoordCfg.CheckAutoBalanceConfigInterval.Key)
-	Params.Save(Params.DataCoordCfg.AutoBalance.Key, "false")
-	defer Params.Reset(Params.DataCoordCfg.AutoBalance.Key)
 
 	t.Run("test old node exist", func(t *testing.T) {
+		Params.Save(Params.DataCoordCfg.AutoBalance.Key, "false")
+		defer Params.Reset(Params.DataCoordCfg.AutoBalance.Key)
 		oldSessions := make(map[string]*sessionutil.Session)
 		oldSessions["s1"] = sessionutil.NewSession(context.Background())
 
@@ -4778,6 +4778,8 @@ func TestUpdateAutoBalanceConfigLoop(t *testing.T) {
 	})
 
 	t.Run("test all old node down", func(t *testing.T) {
+		Params.Save(Params.DataCoordCfg.AutoBalance.Key, "false")
+		defer Params.Reset(Params.DataCoordCfg.AutoBalance.Key)
 		server := &Server{}
 		mockSession := sessionutil.NewMockSession(t)
 		mockSession.EXPECT().GetSessionsWithVersionRange(mock.Anything, mock.Anything).Return(nil, 0, nil).Maybe()
diff --git a/internal/querycoordv2/server_test.go b/internal/querycoordv2/server_test.go
index 109b325d5a..80262832fd 100644
--- a/internal/querycoordv2/server_test.go
+++ b/internal/querycoordv2/server_test.go
@@ -362,10 +362,10 @@ func (suite *ServerSuite) TestUpdateAutoBalanceConfigLoop() {
 
 	Params.Save(Params.QueryCoordCfg.CheckAutoBalanceConfigInterval.Key, "1")
 	defer Params.Reset(Params.QueryCoordCfg.CheckAutoBalanceConfigInterval.Key)
-	Params.Save(Params.QueryCoordCfg.AutoBalance.Key, "false")
-	defer Params.Reset(Params.QueryCoordCfg.AutoBalance.Key)
 
 	suite.Run("test old node exist", func() {
+		Params.Save(Params.QueryCoordCfg.AutoBalance.Key, "false")
+		defer Params.Reset(Params.QueryCoordCfg.AutoBalance.Key)
 		server := &Server{}
 		mockSession := sessionutil.NewMockSession(suite.T())
 		server.session = mockSession
@@ -384,6 +384,8 @@ func (suite *ServerSuite) TestUpdateAutoBalanceConfigLoop() {
 	})
 
 	suite.Run("all old node down", func() {
+		Params.Save(Params.QueryCoordCfg.AutoBalance.Key, "false")
+		defer Params.Reset(Params.QueryCoordCfg.AutoBalance.Key)
 		server := &Server{}
 		mockSession := sessionutil.NewMockSession(suite.T())
 		server.session = mockSession