Add MockCluster, make unit tests reliable (#17032)

Signed-off-by: yah01 <yang.cen@zilliz.com>

internal/querycoord/impl_test.go

@@ -1684,9 +1684,11 @@ func TestGetShardLeaders(t *testing.T) {
 	assert.Equal(t, 0, totalLeaders%3)
 
 	// mock replica all down, without triggering load balance
-	queryCoord.cluster.stopNode(node1.queryNodeID)
-	queryCoord.cluster.stopNode(node2.queryNodeID)
-	queryCoord.cluster.stopNode(node3.queryNodeID)
+	mockCluster := NewMockCluster(queryCoord.cluster)
+	mockCluster.isOnlineHandler = func(nodeID int64) (bool, error) {
+		return false, nil
+	}
+	queryCoord.cluster = mockCluster
 	resp, err = queryCoord.GetShardLeaders(ctx, getShardLeadersReq)
 	assert.NoError(t, err)
 	assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.Status.ErrorCode)

internal/querycoord/mock_cluster.go

@@ -0,0 +1,37 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package querycoord
+
+type MockCluster struct {
+	Cluster
+
+	isOnlineHandler func(nodeID int64) (bool, error)
+}
+
+func NewMockCluster(cluster Cluster) *MockCluster {
+	return &MockCluster{
+		Cluster: cluster,
+	}
+}
+
+func (mock *MockCluster) isOnline(nodeID int64) (bool, error) {
+	if mock.isOnlineHandler != nil {
+		return mock.isOnlineHandler(nodeID)
+	}
+
+	return mock.Cluster.isOnline(nodeID)
+}