milvus/internal/querycoord/group_balance_test.go
SimFG d7f38a803d
Separate some proto files (#19218)
Signed-off-by: SimFG <bang.fu@zilliz.com>

Signed-off-by: SimFG <bang.fu@zilliz.com>
2022-09-16 16:56:49 +08:00

55 lines
1.4 KiB
Go

package querycoord
import (
"context"
"testing"
"github.com/milvus-io/milvus/api/commonpb"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/stretchr/testify/assert"
)
func TestAddNode(t *testing.T) {
defer removeAllSession()
ctx := context.Background()
coord, err := startQueryCoord(ctx)
assert.NoError(t, err)
defer coord.Stop()
node1, err := startQueryNodeServer(ctx)
assert.NoError(t, err)
defer node1.stop()
node2, err := startQueryNodeServer(ctx)
assert.NoError(t, err)
defer node2.stop()
waitQueryNodeOnline(coord.cluster, node1.queryNodeID)
waitQueryNodeOnline(coord.cluster, node2.queryNodeID)
loadCollectionReq := &querypb.LoadCollectionRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_LoadCollection,
},
CollectionID: defaultCollectionID,
Schema: genDefaultCollectionSchema(false),
ReplicaNumber: 1,
}
status, err := coord.LoadCollection(ctx, loadCollectionReq)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, status.ErrorCode)
waitLoadCollectionDone(ctx, coord, defaultCollectionID)
plans, err := coord.groupBalancer.AddNode(node1.queryNodeID)
assert.NoError(t, err)
assert.Equal(t, 0, len(plans))
plans, err = coord.groupBalancer.AddNode(node2.queryNodeID)
assert.NoError(t, err)
assert.Equal(t, 0, len(plans))
newNodeID := node2.queryNodeID + 1
plans, err = coord.groupBalancer.AddNode(newNodeID)
assert.NoError(t, err)
assert.Equal(t, 1, len(plans))
}