milvus-io--milvus/internal/proxy/impl_test.go

// 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 proxy

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"

	"github.com/apache/pulsar-client-go/pulsar"
	"github.com/bytedance/mockey"
	"github.com/cockroachdb/errors"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/otel/trace"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/internal/allocator"
	grpcmixcoordclient "github.com/milvus-io/milvus/internal/distributed/mixcoord/client"
	"github.com/milvus-io/milvus/internal/distributed/streaming"
	mhttp "github.com/milvus-io/milvus/internal/http"
	"github.com/milvus-io/milvus/internal/mocks"
	"github.com/milvus-io/milvus/internal/mocks/distributed/mock_streaming"
	"github.com/milvus-io/milvus/internal/proxy/shardclient"
	"github.com/milvus-io/milvus/internal/util/dependency"
	"github.com/milvus-io/milvus/internal/util/segcore"
	"github.com/milvus-io/milvus/internal/util/sessionutil"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
	"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/proxypb"
	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
	"github.com/milvus-io/milvus/pkg/v3/proto/rootcoordpb"
	"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
	pulsar2 "github.com/milvus-io/milvus/pkg/v3/streaming/walimpls/impls/pulsar"
	"github.com/milvus-io/milvus/pkg/v3/util"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v3/util/ratelimitutil"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

func TestProxy_InvalidateCollectionMetaCache_remove_stream(t *testing.T) {
	paramtable.Init()
	cache := globalMetaCache
	globalMetaCache = nil
	defer func() { globalMetaCache = cache }()

	chMgr := NewMockChannelsMgr(t)
	chMgr.EXPECT().removeDMLStream(mock.Anything).Return()

	node := &Proxy{chMgr: chMgr}
	_ = node.initRateCollector()
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	ctx := context.Background()
	req := &proxypb.InvalidateCollMetaCacheRequest{
		Base: &commonpb.MsgBase{MsgType: commonpb.MsgType_DropCollection},
	}

	status, err := node.InvalidateCollectionMetaCache(ctx, req)
	assert.NoError(t, err)
	assert.Equal(t, commonpb.ErrorCode_Success, status.GetErrorCode())
}

func TestProxy_CheckHealth(t *testing.T) {
	t.Run("not healthy", func(t *testing.T) {
		node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
		node.simpleLimiter = NewSimpleLimiter(0, 0)
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		ctx := context.Background()
		resp, err := node.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
		assert.NoError(t, err)
		assert.Equal(t, false, resp.IsHealthy)
		assert.Equal(t, 1, len(resp.Reasons))
	})

	t.Run("proxy health check is ok", func(t *testing.T) {
		mixc := &mocks.MockMixCoordClient{}
		mixc.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(&milvuspb.CheckHealthResponse{IsHealthy: true}, nil)
		node := &Proxy{
			mixCoord: mixc,
			session:  &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}},
		}
		node.simpleLimiter = NewSimpleLimiter(0, 0)
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()
		resp, err := node.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
		assert.NoError(t, err)
		assert.Equal(t, true, resp.IsHealthy)
		assert.Empty(t, resp.Reasons)
	})

	t.Run("proxy health check is fail", func(t *testing.T) {
		checkHealthFunc1 := func(ctx context.Context,
			req *milvuspb.CheckHealthRequest,
			opts ...grpc.CallOption,
		) (*milvuspb.CheckHealthResponse, error) {
			return &milvuspb.CheckHealthResponse{
				IsHealthy: false,
				Reasons:   []string{"unHealth"},
			}, nil
		}
		qc := &mocks.MockQueryCoordClient{}
		qc.EXPECT().CheckHealth(mock.Anything, mock.Anything).Return(nil, errors.New("test"))
		node := &Proxy{
			session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}},
			mixCoord: NewMixCoordMock(func(mock *MixCoordMock) {
				mock.checkHealthFunc = checkHealthFunc1
			}),
		}
		node.simpleLimiter = NewSimpleLimiter(0, 0)
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()
		resp, err := node.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
		assert.NoError(t, err)
		assert.Equal(t, false, resp.IsHealthy)
		assert.Equal(t, 1, len(resp.Reasons))
	})

	t.Run("check quota state", func(t *testing.T) {
		node := &Proxy{
			mixCoord: NewMixCoordMock(),
		}
		node.simpleLimiter = NewSimpleLimiter(0, 0)
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		resp, err := node.CheckHealth(context.Background(), &milvuspb.CheckHealthRequest{})
		assert.NoError(t, err)
		assert.Equal(t, true, resp.IsHealthy)
		assert.Equal(t, 0, len(resp.GetQuotaStates()))
		assert.Equal(t, 0, len(resp.GetReasons()))

		states := []milvuspb.QuotaState{milvuspb.QuotaState_DenyToWrite, milvuspb.QuotaState_DenyToRead}
		codes := []commonpb.ErrorCode{commonpb.ErrorCode_MemoryQuotaExhausted, commonpb.ErrorCode_ForceDeny}
		err = node.simpleLimiter.SetRates(&proxypb.LimiterNode{
			Limiter: &proxypb.Limiter{},
			// db level
			Children: map[int64]*proxypb.LimiterNode{
				1: {
					Limiter: &proxypb.Limiter{},
					// collection level
					Children: map[int64]*proxypb.LimiterNode{
						100: {
							Limiter: &proxypb.Limiter{
								States: states,
								Codes:  codes,
							},
							Children: make(map[int64]*proxypb.LimiterNode),
						},
					},
				},
			},
		})
		assert.NoError(t, err)

		resp, err = node.CheckHealth(context.Background(), &milvuspb.CheckHealthRequest{})
		assert.NoError(t, err)
		assert.Equal(t, true, resp.IsHealthy)
	})
}

func TestProxyRoleDescriptionValidation(t *testing.T) {
	paramtable.Init()
	node := &Proxy{mixCoord: NewMixCoordMock()}
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	ctx := context.Background()

	paramtable.Get().Save(Params.ProxyCfg.MaxRoleDescriptionLength.Key, "4")
	defer paramtable.Get().Reset(Params.ProxyCfg.MaxRoleDescriptionLength.Key)

	status, err := node.CreateRole(ctx, &milvuspb.CreateRoleRequest{
		Entity: &milvuspb.RoleEntity{
			Name:        "role_desc",
			Description: "12345",
		},
	})
	require.NoError(t, err)
	assert.NotEqual(t, commonpb.ErrorCode_Success, status.GetErrorCode())

	status, err = node.AlterRole(ctx, &milvuspb.AlterRoleRequest{
		RoleName:    "role_desc",
		Description: "12345",
	})
	require.NoError(t, err)
	assert.NotEqual(t, commonpb.ErrorCode_Success, status.GetErrorCode())

	status, err = node.AlterRole(ctx, &milvuspb.AlterRoleRequest{
		RoleName:    util.RoleAdmin,
		Description: "ok",
	})
	require.NoError(t, err)
	assert.NotEqual(t, commonpb.ErrorCode_Success, status.GetErrorCode())

	status, err = node.RestoreRBAC(ctx, &milvuspb.RestoreRBACMetaRequest{
		RBACMeta: &milvuspb.RBACMeta{
			Roles: []*milvuspb.RoleEntity{
				{
					Name:        "role_desc",
					Description: "12345",
				},
			},
		},
	})
	require.NoError(t, err)
	assert.NotEqual(t, commonpb.ErrorCode_Success, status.GetErrorCode())
}

func TestProxyRoleDescriptionForwarding(t *testing.T) {
	paramtable.Init()
	mixCoord := mocks.NewMockMixCoordClient(t)
	node := &Proxy{mixCoord: mixCoord}
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	ctx := context.Background()

	mixCoord.EXPECT().CreateRole(mock.Anything, mock.MatchedBy(func(req *milvuspb.CreateRoleRequest) bool {
		return req.GetBase().GetMsgType() == commonpb.MsgType_CreateRole &&
			req.GetEntity().GetName() == "role_desc_forward" &&
			req.GetEntity().GetDescription() == "role description"
	})).Return(merr.Success(), nil).Once()
	status, err := node.CreateRole(ctx, &milvuspb.CreateRoleRequest{
		Entity: &milvuspb.RoleEntity{
			Name:        "role_desc_forward",
			Description: "role description",
		},
	})
	require.NoError(t, err)
	require.True(t, merr.Ok(status))

	mixCoord.EXPECT().AlterRole(mock.Anything, mock.MatchedBy(func(req *milvuspb.AlterRoleRequest) bool {
		return req.GetBase().GetMsgType() == commonpb.MsgType_AlterRole &&
			req.GetRoleName() == "role_desc_forward" &&
			req.GetDescription() == "updated role description"
	})).Return(merr.Success(), nil).Once()
	status, err = node.AlterRole(ctx, &milvuspb.AlterRoleRequest{
		RoleName:    "role_desc_forward",
		Description: "updated role description",
	})
	require.NoError(t, err)
	require.True(t, merr.Ok(status))
}

func TestProxyRenameCollection(t *testing.T) {
	t.Run("not healthy", func(t *testing.T) {
		node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		ctx := context.Background()
		resp, err := node.RenameCollection(ctx, &milvuspb.RenameCollectionRequest{})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp), merr.ErrServiceNotReady)
	})

	t.Run("rename with illegal new collection name", func(t *testing.T) {
		node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()
		resp, err := node.RenameCollection(ctx, &milvuspb.RenameCollectionRequest{NewName: "$#^%#&#$*!)#@!"})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp), merr.ErrParameterInvalid)
	})

	t.Run("rename fail", func(t *testing.T) {
		rc := mocks.NewMockMixCoordClient(t)
		rc.On("RenameCollection", mock.Anything, mock.Anything).
			Return(nil, errors.New("fail"))
		node := &Proxy{
			session:  &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}},
			mixCoord: rc,
		}
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()

		resp, err := node.RenameCollection(ctx, &milvuspb.RenameCollectionRequest{NewName: "new"})
		assert.Error(t, err)
		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
	})

	t.Run("rename ok", func(t *testing.T) {
		rc := mocks.NewMockMixCoordClient(t)
		rc.On("RenameCollection", mock.Anything, mock.Anything).
			Return(merr.Success(), nil)
		node := &Proxy{
			session:  &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}},
			mixCoord: rc,
		}
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()

		resp, err := node.RenameCollection(ctx, &milvuspb.RenameCollectionRequest{NewName: "new"})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
	})
}

func TestProxyFunctionEdit(t *testing.T) {
	mockey.PatchConvey("TestProxy_AddFunction", t, func() {
		m1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
			return nil
		}).Build()
		m2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()
		defer m1.UnPatch()
		defer m2.UnPatch()
		node := createTestProxy()
		defer node.sched.Close()

		_, err := node.AddCollectionFunction(context.Background(), &milvuspb.AddCollectionFunctionRequest{
			Base: &commonpb.MsgBase{
				MsgType: commonpb.MsgType_AddCollectionFunction,
			},
			CollectionName: "test_collection",
			CollectionID:   1,
			FunctionSchema: &schemapb.FunctionSchema{
				Name:             "test_function",
				Type:             schemapb.FunctionType_TextEmbedding,
				InputFieldNames:  []string{},
				OutputFieldNames: []string{},
				Params:           []*commonpb.KeyValuePair{},
			},
		})
		assert.NoError(t, err)
	})

	mockey.PatchConvey("TestProxy_DropFunction", t, func() {
		m1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
			return nil
		}).Build()
		m2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()
		defer m1.UnPatch()
		defer m2.UnPatch()
		node := createTestProxy()
		defer node.sched.Close()

		_, err := node.DropCollectionFunction(context.Background(), &milvuspb.DropCollectionFunctionRequest{
			Base: &commonpb.MsgBase{
				MsgType: commonpb.MsgType_AddCollectionFunction,
			},
			CollectionName: "test_collection",
			CollectionID:   1,
			FunctionName:   "test",
		})
		assert.NoError(t, err)
	})

	mockey.PatchConvey("TestProxy_AlterFunction", t, func() {
		m1 := mockey.Mock((*ddTaskQueue).Enqueue).To(func(t task) error {
			return nil
		}).Build()
		m2 := mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()
		defer m1.UnPatch()
		defer m2.UnPatch()
		node := createTestProxy()
		defer node.sched.Close()

		_, err := node.AlterCollectionFunction(context.Background(), &milvuspb.AlterCollectionFunctionRequest{
			Base: &commonpb.MsgBase{
				MsgType: commonpb.MsgType_AddCollectionFunction,
			},
			CollectionName: "test_collection",
			CollectionID:   1,
			FunctionName:   "test",
			FunctionSchema: &schemapb.FunctionSchema{
				Name:             "test_function",
				Type:             schemapb.FunctionType_TextEmbedding,
				InputFieldNames:  []string{},
				OutputFieldNames: []string{},
				Params:           []*commonpb.KeyValuePair{},
			},
		})
		assert.NoError(t, err)
	})
}

func TestProxy_ResourceGroup(t *testing.T) {
	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.simpleLimiter = NewSimpleLimiter(0, 0)
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	qc := mocks.NewMockMixCoordClient(t)
	node.SetMixCoordClient(qc)
	qc.EXPECT().ListPolicy(ctx, mock.Anything).Return(&internalpb.ListPolicyResponse{
		Status: merr.Success(),
	}, nil).Maybe()
	qc.EXPECT().ShowLoadCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe()

	tsoAllocatorIns := newMockTsoAllocator()
	node.sched, err = newTaskScheduler(node.ctx, tsoAllocatorIns)
	assert.NoError(t, err)
	node.sched.Start()
	defer node.sched.Close()

	// mgr := newShardClientMgr()
	InitMetaCache(ctx, qc)

	successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}

	t.Run("create resource group", func(t *testing.T) {
		qc.EXPECT().CreateResourceGroup(mock.Anything, mock.Anything).Return(successStatus, nil)
		resp, err := node.CreateResourceGroup(ctx, &milvuspb.CreateResourceGroupRequest{
			ResourceGroup: "rg",
		})
		assert.NoError(t, err)
		assert.Equal(t, resp.ErrorCode, commonpb.ErrorCode_Success)
	})

	t.Run("drop resource group", func(t *testing.T) {
		qc.EXPECT().DropResourceGroup(mock.Anything, mock.Anything).Return(successStatus, nil)
		resp, err := node.DropResourceGroup(ctx, &milvuspb.DropResourceGroupRequest{
			ResourceGroup: "rg",
		})
		assert.NoError(t, err)
		assert.Equal(t, resp.ErrorCode, commonpb.ErrorCode_Success)
	})

	t.Run("transfer node", func(t *testing.T) {
		qc.EXPECT().TransferNode(mock.Anything, mock.Anything).Return(successStatus, nil)
		resp, err := node.TransferNode(ctx, &milvuspb.TransferNodeRequest{
			SourceResourceGroup: "rg1",
			TargetResourceGroup: "rg2",
			NumNode:             1,
		})
		assert.NoError(t, err)
		assert.Equal(t, resp.ErrorCode, commonpb.ErrorCode_Success)
	})

	t.Run("list resource group", func(t *testing.T) {
		qc.EXPECT().ListResourceGroups(mock.Anything, mock.Anything).Return(&milvuspb.ListResourceGroupsResponse{Status: successStatus}, nil)
		resp, err := node.ListResourceGroups(ctx, &milvuspb.ListResourceGroupsRequest{})
		assert.NoError(t, err)
		assert.True(t, merr.Ok(resp.GetStatus()))
	})

	t.Run("describe resource group", func(t *testing.T) {
		qc.EXPECT().DescribeResourceGroup(mock.Anything, mock.Anything).Return(&querypb.DescribeResourceGroupResponse{
			Status: successStatus,
			ResourceGroup: &querypb.ResourceGroupInfo{
				Name:             "rg",
				Capacity:         1,
				NumAvailableNode: 1,
				NumLoadedReplica: nil,
				NumOutgoingNode:  nil,
				NumIncomingNode:  nil,
			},
		}, nil)
		resp, err := node.DescribeResourceGroup(ctx, &milvuspb.DescribeResourceGroupRequest{
			ResourceGroup: "rg",
		})
		assert.NoError(t, err)
		assert.True(t, merr.Ok(resp.GetStatus()))
	})
}

func TestProxy_InvalidResourceGroupName(t *testing.T) {
	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.simpleLimiter = NewSimpleLimiter(0, 0)
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	qc := mocks.NewMockMixCoordClient(t)
	node.SetMixCoordClient(qc)
	qc.EXPECT().DropResourceGroup(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
	qc.EXPECT().ListPolicy(ctx, mock.Anything).Return(&internalpb.ListPolicyResponse{
		Status: merr.Success(),
	}, nil).Maybe()
	tsoAllocatorIns := newMockTsoAllocator()
	node.sched, err = newTaskScheduler(node.ctx, tsoAllocatorIns)
	assert.NoError(t, err)
	node.sched.Start()
	defer node.sched.Close()

	// mgr := newShardClientMgr()
	InitMetaCache(ctx, qc)

	t.Run("create resource group", func(t *testing.T) {
		resp, err := node.CreateResourceGroup(ctx, &milvuspb.CreateResourceGroupRequest{
			ResourceGroup: "...",
		})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp), merr.ErrParameterInvalid)
	})

	t.Run("drop resource group", func(t *testing.T) {
		resp, err := node.DropResourceGroup(ctx, &milvuspb.DropResourceGroupRequest{
			ResourceGroup: "...",
		})
		assert.NoError(t, err)
		assert.Equal(t, resp.ErrorCode, commonpb.ErrorCode_Success)
	})

	t.Run("transfer node", func(t *testing.T) {
		resp, err := node.TransferNode(ctx, &milvuspb.TransferNodeRequest{
			SourceResourceGroup: "...",
			TargetResourceGroup: "!!!",
			NumNode:             1,
		})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp), merr.ErrParameterInvalid)
	})

	t.Run("transfer replica", func(t *testing.T) {
		resp, err := node.TransferReplica(ctx, &milvuspb.TransferReplicaRequest{
			SourceResourceGroup: "...",
			TargetResourceGroup: "!!!",
			NumReplica:          1,
			CollectionName:      "collection1",
		})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp), merr.ErrParameterInvalid)
	})
}

// createTestProxy creates a test proxy instance with all necessary setup
func createTestProxy() *Proxy {
	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, _ := NewProxy(ctx, factory)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}

	node.sched, _ = newTaskScheduler(ctx, node.tsoAllocator)
	node.sched.Start()

	return node
}

func TestProxy_FlushAll_Success(t *testing.T) {
	mockey.PatchConvey("TestProxy_FlushAll_Success", t, func() {
		// Mock global meta cache methods
		globalMetaCache = &MetaCache{}
		mockey.Mock(globalMetaCache.GetCollectionID).To(func(ctx context.Context, dbName, collectionName string) (UniqueID, error) {
			return UniqueID(0), nil
		}).Build()
		mockey.Mock(globalMetaCache.RemoveDatabase).Return().Build()

		// Mock paramtable initialization
		mockey.Mock(paramtable.Init).Return().Build()
		mockey.Mock((*paramtable.ComponentParam).Save).Return().Build()

		successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
		mockey.Mock((*grpcmixcoordclient.Client).ListDatabases).To(func(ctx context.Context, req *milvuspb.ListDatabasesRequest, opts ...grpc.CallOption) (*milvuspb.ListDatabasesResponse, error) {
			return &milvuspb.ListDatabasesResponse{Status: successStatus, DbNames: []string{"default"}}, nil
		}).Build()
		// Mock grpc mix coord client FlushAll method for default database
		mockey.Mock((*grpcmixcoordclient.Client).ShowCollections).To(func(ctx context.Context, req *milvuspb.ShowCollectionsRequest, opts ...grpc.CallOption) (*milvuspb.ShowCollectionsResponse, error) {
			return &milvuspb.ShowCollectionsResponse{Status: successStatus}, nil
		}).Build()

		// Act: Execute test
		node := createTestProxy()
		defer node.sched.Close()

		messageID := pulsar2.NewPulsarID(pulsar.EarliestMessageID())
		msg := message.NewFlushAllMessageBuilderV2().
			WithVChannel("test-vchannel").
			WithHeader(&message.FlushAllMessageHeader{}).
			WithBody(&message.FlushAllMessageBody{}).
			MustBuildMutable().WithTimeTick(1000).
			WithLastConfirmed(messageID)
		milvusMsg := message.ImmutableMessageToMilvusMessage(commonpb.WALName_Pulsar.String(), msg.IntoImmutableMessage(messageID))

		mixcoord := &grpcmixcoordclient.Client{}
		node.mixCoord = mixcoord
		mockey.Mock((*grpcmixcoordclient.Client).FlushAll).To(func(ctx context.Context, req *datapb.FlushAllRequest, opts ...grpc.CallOption) (*datapb.FlushAllResponse, error) {
			return &datapb.FlushAllResponse{
				Status: successStatus,
				FlushAllMsgs: map[string]*commonpb.ImmutableMessage{
					"channel1": milvusMsg,
				},
			}, nil
		}).Build()

		resp, err := node.FlushAll(context.Background(), &milvuspb.FlushAllRequest{})

		// Assert: Verify results
		assert.NoError(t, err)
		assert.True(t, merr.Ok(resp.GetStatus()))
	})
}

func TestProxy_FlushAll_ServerAbnormal(t *testing.T) {
	mockey.PatchConvey("TestProxy_FlushAll_ServerAbnormal", t, func() {
		// Mock global meta cache methods
		globalMetaCache = &MetaCache{}
		mockey.Mock(globalMetaCache.GetCollectionID).To(func(ctx context.Context, dbName, collectionName string) (UniqueID, error) {
			return UniqueID(0), nil
		}).Build()
		mockey.Mock(globalMetaCache.RemoveDatabase).Return().Build()

		// Mock paramtable initialization
		mockey.Mock(paramtable.Init).Return().Build()
		mockey.Mock((*paramtable.ComponentParam).Save).Return().Build()

		// Act: Execute test
		node := createTestProxy()
		defer node.sched.Close()

		mixcoord := &grpcmixcoordclient.Client{}
		node.mixCoord = mixcoord

		// Set node state to abnormal
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		resp, err := node.FlushAll(context.Background(), &milvuspb.FlushAllRequest{})

		// Assert: Verify results
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrServiceNotReady)
	})
}

func TestProxy_GetFlushAllState(t *testing.T) {
	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}
	node.mixCoord = mocks.NewMockMixCoordClient(t)

	// set expectations
	successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
	node.mixCoord.(*mocks.MockMixCoordClient).EXPECT().GetFlushAllState(mock.Anything, mock.Anything).
		Return(&milvuspb.GetFlushAllStateResponse{Status: successStatus}, nil).Maybe()

	t.Run("GetFlushAllState success", func(t *testing.T) {
		resp, err := node.GetFlushAllState(ctx, &milvuspb.GetFlushAllStateRequest{})
		assert.NoError(t, err)
		assert.True(t, merr.Ok(resp.GetStatus()))
	})

	t.Run("GetFlushAllState failed, server is abnormal", func(t *testing.T) {
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		resp, err := node.GetFlushAllState(ctx, &milvuspb.GetFlushAllStateRequest{})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrServiceNotReady)
		node.UpdateStateCode(commonpb.StateCode_Healthy)
	})

	t.Run("DataCoord GetFlushAllState failed", func(t *testing.T) {
		node.mixCoord.(*mocks.MockMixCoordClient).ExpectedCalls = nil
		node.mixCoord.(*mocks.MockMixCoordClient).EXPECT().GetFlushAllState(mock.Anything, mock.Anything).
			Return(&milvuspb.GetFlushAllStateResponse{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    "mock err",
				},
			}, nil)
		resp, err := node.GetFlushAllState(ctx, &milvuspb.GetFlushAllStateRequest{})
		assert.NoError(t, err)
		assert.Equal(t, resp.GetStatus().GetErrorCode(), commonpb.ErrorCode_UnexpectedError)
	})
}

func TestProxy_GetFlushState(t *testing.T) {
	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}
	node.mixCoord = mocks.NewMockMixCoordClient(t)

	// set expectations
	successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
	node.mixCoord.(*mocks.MockMixCoordClient).EXPECT().GetFlushState(mock.Anything, mock.Anything, mock.Anything).
		Return(&milvuspb.GetFlushStateResponse{Status: successStatus}, nil).Maybe()

	t.Run("GetFlushState success", func(t *testing.T) {
		resp, err := node.GetFlushState(ctx, &milvuspb.GetFlushStateRequest{})
		assert.NoError(t, err)
		assert.Equal(t, resp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
	})

	t.Run("GetFlushState failed, server is abnormal", func(t *testing.T) {
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		resp, err := node.GetFlushState(ctx, &milvuspb.GetFlushStateRequest{})
		assert.NoError(t, err)
		assert.Equal(t, resp.GetStatus().GetErrorCode(), commonpb.ErrorCode_NotReadyServe)
		node.UpdateStateCode(commonpb.StateCode_Healthy)
	})

	t.Run("GetFlushState with collection name", func(t *testing.T) {
		resp, err := node.GetFlushState(ctx, &milvuspb.GetFlushStateRequest{
			CollectionName: "*",
		})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrParameterInvalid)

		cacheBak := globalMetaCache
		defer func() { globalMetaCache = cacheBak }()
		cache := NewMockCache(t)
		cache.On("GetCollectionID",
			mock.Anything, // context.Context
			mock.AnythingOfType("string"),
			mock.AnythingOfType("string"),
		).Return(UniqueID(0), nil).Maybe()
		globalMetaCache = cache

		resp, err = node.GetFlushState(ctx, &milvuspb.GetFlushStateRequest{
			CollectionName: "collection1",
		})
		assert.NoError(t, err)
		assert.Equal(t, resp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
	})

	t.Run("DataCoord GetFlushState failed", func(t *testing.T) {
		node.mixCoord.(*mocks.MockMixCoordClient).ExpectedCalls = nil
		node.mixCoord.(*mocks.MockMixCoordClient).EXPECT().GetFlushState(mock.Anything, mock.Anything, mock.Anything).
			Return(&milvuspb.GetFlushStateResponse{
				Status: &commonpb.Status{
					ErrorCode: commonpb.ErrorCode_UnexpectedError,
					Reason:    "mock err",
				},
			}, nil)
		resp, err := node.GetFlushState(ctx, &milvuspb.GetFlushStateRequest{})
		assert.NoError(t, err)
		assert.Equal(t, resp.GetStatus().GetErrorCode(), commonpb.ErrorCode_UnexpectedError)
	})

	t.Run("GetFlushState return error", func(t *testing.T) {
		node.mixCoord.(*mocks.MockMixCoordClient).ExpectedCalls = nil
		node.mixCoord.(*mocks.MockMixCoordClient).EXPECT().GetFlushState(mock.Anything, mock.Anything, mock.Anything).
			Return(nil, errors.New("fake error"))
		resp, err := node.GetFlushState(ctx, &milvuspb.GetFlushStateRequest{})
		assert.NoError(t, err)
		assert.Equal(t, resp.GetStatus().GetErrorCode(), commonpb.ErrorCode_UnexpectedError)
	})
}

func TestProxy_GetReplicas(t *testing.T) {
	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}
	mockQC := mocks.NewMockMixCoordClient(t)
	node.mixCoord = mockQC

	// set expectations
	successStatus := &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}
	t.Run("success", func(t *testing.T) {
		mockQC.EXPECT().GetReplicas(mock.Anything, mock.AnythingOfType("*milvuspb.GetReplicasRequest")).Return(&milvuspb.GetReplicasResponse{Status: successStatus}, nil)
		resp, err := node.GetReplicas(ctx, &milvuspb.GetReplicasRequest{
			CollectionID: 1000,
		})
		assert.NoError(t, err)
		assert.True(t, merr.Ok(resp.GetStatus()))
	})

	t.Run("proxy_not_healthy", func(t *testing.T) {
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		resp, err := node.GetReplicas(ctx, &milvuspb.GetReplicasRequest{
			CollectionID: 1000,
		})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrServiceNotReady)
		node.UpdateStateCode(commonpb.StateCode_Healthy)
	})

	t.Run("QueryCoordClient_returnsError", func(t *testing.T) {
		mockQC.ExpectedCalls = nil
		mockQC.EXPECT().GetReplicas(mock.Anything, mock.AnythingOfType("*milvuspb.GetReplicasRequest")).Return(nil, errors.New("mocked"))

		resp, err := node.GetReplicas(ctx, &milvuspb.GetReplicasRequest{
			CollectionID: 1000,
		})
		assert.NoError(t, err)
		assert.Equal(t, resp.GetStatus().GetErrorCode(), commonpb.ErrorCode_UnexpectedError)
	})
}

func TestProxy_Connect(t *testing.T) {
	t.Run("proxy unhealthy", func(t *testing.T) {
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)

		resp, err := node.Connect(context.TODO(), nil)
		assert.NoError(t, err)
		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("failed to list database", func(t *testing.T) {
		r := mocks.NewMockMixCoordClient(t)
		r.On("ListDatabases",
			mock.Anything,
			mock.Anything,
		).Return(nil, errors.New("error mock ListDatabases"))

		node := &Proxy{mixCoord: r}
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		resp, err := node.Connect(context.TODO(), nil)
		assert.NoError(t, err)
		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("list database error", func(t *testing.T) {
		r := mocks.NewMockMixCoordClient(t)
		r.On("ListDatabases",
			mock.Anything,
			mock.Anything,
		).Return(&milvuspb.ListDatabasesResponse{
			Status: merr.Status(merr.WrapErrServiceNotReady(paramtable.GetRole(), paramtable.GetNodeID(), "initialization")),
		}, nil)

		node := &Proxy{mixCoord: r}
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		resp, err := node.Connect(context.TODO(), nil)
		assert.NoError(t, err)
		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("database not found", func(t *testing.T) {
		md := metadata.New(map[string]string{
			"dbName": "20230525",
		})
		ctx := metadata.NewIncomingContext(context.TODO(), md)

		r := mocks.NewMockMixCoordClient(t)
		r.On("ListDatabases",
			mock.Anything,
			mock.Anything,
		).Return(&milvuspb.ListDatabasesResponse{
			Status:  merr.Success(),
			DbNames: []string{},
		}, nil)

		node := &Proxy{mixCoord: r}
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		resp, err := node.Connect(ctx, nil)
		assert.NoError(t, err)
		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("failed to allocate ts", func(t *testing.T) {
		md := metadata.New(map[string]string{
			"dbName": "20230525",
		})
		ctx := metadata.NewIncomingContext(context.TODO(), md)

		r := mocks.NewMockMixCoordClient(t)
		r.On("ListDatabases",
			mock.Anything,
			mock.Anything,
		).Return(&milvuspb.ListDatabasesResponse{
			Status:  merr.Success(),
			DbNames: []string{"20230525"},
		}, nil)

		m := newMockTimestampAllocator(t)
		m.On("AllocTimestamp",
			mock.Anything,
			mock.Anything,
		).Return(nil, errors.New("error mock AllocateTimestamp"))
		alloc, _ := newTimestampAllocator(m, 199)
		node := Proxy{
			tsoAllocator: alloc,
			mixCoord:     r,
		}
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		resp, err := node.Connect(ctx, nil)
		assert.NoError(t, err)
		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("normal case", func(t *testing.T) {
		md := metadata.New(map[string]string{
			"dbName": "20230525",
		})
		ctx := metadata.NewIncomingContext(context.TODO(), md)

		r := mocks.NewMockMixCoordClient(t)
		r.On("ListDatabases",
			mock.Anything,
			mock.Anything,
		).Return(&milvuspb.ListDatabasesResponse{
			Status:  merr.Success(),
			DbNames: []string{"20230525"},
		}, nil)

		m := newMockTimestampAllocator(t)
		m.On("AllocTimestamp",
			mock.Anything,
			mock.Anything,
		).Return(&rootcoordpb.AllocTimestampResponse{
			Status:    merr.Success(),
			Timestamp: 20230518,
			Count:     1,
		}, nil)
		alloc, _ := newTimestampAllocator(m, 199)
		node := Proxy{
			tsoAllocator: alloc,
			mixCoord:     r,
		}
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		resp, err := node.Connect(ctx, &milvuspb.ConnectRequest{
			ClientInfo: &commonpb.ClientInfo{},
		})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})
}

func TestProxy_ListClientInfos(t *testing.T) {
	t.Run("proxy unhealthy", func(t *testing.T) {
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)

		resp, err := node.ListClientInfos(context.TODO(), nil)
		assert.NoError(t, err)
		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("normal case", func(t *testing.T) {
		node := Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		resp, err := node.ListClientInfos(context.TODO(), nil)
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})
}

func TestProxyCreateDatabase(t *testing.T) {
	paramtable.Init()

	t.Run("not healthy", func(t *testing.T) {
		node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		ctx := context.Background()
		resp, err := node.CreateDatabase(ctx, &milvuspb.CreateDatabaseRequest{})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp), merr.ErrServiceNotReady)
	})

	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}
	node.simpleLimiter = NewSimpleLimiter(0, 0)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
	node.sched.ddQueue.setMaxTaskNum(10)
	assert.NoError(t, err)
	err = node.sched.Start()
	assert.NoError(t, err)
	defer node.sched.Close()

	t.Run("create database fail", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("CreateDatabase", mock.Anything, mock.Anything).
			Return(nil, errors.New("fail"))
		node.mixCoord = mixc
		ctx := context.Background()
		resp, err := node.CreateDatabase(ctx, &milvuspb.CreateDatabaseRequest{DbName: "db"})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
	})

	t.Run("create database ok", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("CreateDatabase", mock.Anything, mock.Anything).
			Return(merr.Success(), nil)
		node.mixCoord = mixc
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()

		resp, err := node.CreateDatabase(ctx, &milvuspb.CreateDatabaseRequest{DbName: "db"})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
	})
}

func TestProxyDropDatabase(t *testing.T) {
	paramtable.Init()

	t.Run("not healthy", func(t *testing.T) {
		node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		ctx := context.Background()
		resp, err := node.DropDatabase(ctx, &milvuspb.DropDatabaseRequest{})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp), merr.ErrServiceNotReady)
	})

	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	node.initRateCollector()
	assert.NoError(t, err)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}
	node.simpleLimiter = NewSimpleLimiter(0, 0)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
	node.sched.ddQueue.setMaxTaskNum(10)
	assert.NoError(t, err)
	err = node.sched.Start()
	assert.NoError(t, err)
	defer node.sched.Close()

	t.Run("drop database fail", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("DropDatabase", mock.Anything, mock.Anything).
			Return(nil, errors.New("fail"))
		node.mixCoord = mixc
		ctx := context.Background()
		resp, err := node.DropDatabase(ctx, &milvuspb.DropDatabaseRequest{DbName: "db"})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
	})

	t.Run("drop database ok", func(t *testing.T) {
		mix := mocks.NewMockMixCoordClient(t)
		mix.EXPECT().DropDatabase(mock.Anything, mock.Anything).Return(merr.Success(), nil)
		mix.EXPECT().ListPolicy(mock.Anything, mock.Anything).Return(&internalpb.ListPolicyResponse{
			Status: merr.Success(),
		}, nil).Maybe()
		mix.EXPECT().ShowLoadCollections(mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{}, nil).Maybe()
		node.mixCoord = mix
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		cacheBak := globalMetaCache
		defer func() { globalMetaCache = cacheBak }()
		cache := NewMockCache(t)
		cache.EXPECT().RemoveDatabase(mock.Anything, mock.AnythingOfType("string")).Return()
		globalMetaCache = cache

		ctx := context.Background()

		resp, err := node.DropDatabase(ctx, &milvuspb.DropDatabaseRequest{DbName: "db"})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
	})
}

func TestProxyListDatabase(t *testing.T) {
	paramtable.Init()

	t.Run("not healthy", func(t *testing.T) {
		node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		ctx := context.Background()
		resp, err := node.ListDatabases(ctx, &milvuspb.ListDatabasesRequest{})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrServiceNotReady)
	})

	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}
	node.simpleLimiter = NewSimpleLimiter(0, 0)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
	node.sched.ddQueue.setMaxTaskNum(10)
	assert.NoError(t, err)
	err = node.sched.Start()
	assert.NoError(t, err)
	defer node.sched.Close()

	t.Run("list database fail", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("ListDatabases", mock.Anything, mock.Anything).
			Return(nil, errors.New("fail"))
		node.mixCoord = mixc
		ctx := context.Background()
		resp, err := node.ListDatabases(ctx, &milvuspb.ListDatabasesRequest{})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetStatus().GetErrorCode())
	})

	t.Run("list database ok", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("ListDatabases", mock.Anything, mock.Anything).
			Return(&milvuspb.ListDatabasesResponse{
				Status: merr.Success(),
			}, nil)
		node.mixCoord = mixc
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()

		resp, err := node.ListDatabases(ctx, &milvuspb.ListDatabasesRequest{})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})
}

func TestProxyAlterDatabase(t *testing.T) {
	paramtable.Init()

	t.Run("not healthy", func(t *testing.T) {
		node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		ctx := context.Background()
		resp, err := node.AlterDatabase(ctx, &milvuspb.AlterDatabaseRequest{})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp), merr.ErrServiceNotReady)
	})

	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}
	node.simpleLimiter = NewSimpleLimiter(0, 0)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
	node.sched.ddQueue.setMaxTaskNum(10)
	assert.NoError(t, err)
	err = node.sched.Start()
	assert.NoError(t, err)
	defer node.sched.Close()

	t.Run("alter database fail", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("AlterDatabase", mock.Anything, mock.Anything).Return(nil, errors.New("fail"))
		node.mixCoord = mixc
		ctx := context.Background()
		resp, err := node.AlterDatabase(ctx, &milvuspb.AlterDatabaseRequest{})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
	})

	t.Run("alter database ok", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("AlterDatabase", mock.Anything, mock.Anything).
			Return(merr.Success(), nil)
		node.mixCoord = mixc
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()

		resp, err := node.AlterDatabase(ctx, &milvuspb.AlterDatabaseRequest{})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
	})
}

func TestProxyDescribeDatabase(t *testing.T) {
	paramtable.Init()

	t.Run("not healthy", func(t *testing.T) {
		node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		ctx := context.Background()
		resp, err := node.DescribeDatabase(ctx, &milvuspb.DescribeDatabaseRequest{})
		assert.NoError(t, err)
		assert.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrServiceNotReady)
	})

	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()

	node, err := NewProxy(ctx, factory)
	assert.NoError(t, err)
	node.tsoAllocator = &timestampAllocator{
		tso: newMockTimestampAllocatorInterface(),
	}
	node.simpleLimiter = NewSimpleLimiter(0, 0)
	node.UpdateStateCode(commonpb.StateCode_Healthy)
	node.sched, err = newTaskScheduler(ctx, node.tsoAllocator)
	node.sched.ddQueue.setMaxTaskNum(10)
	assert.NoError(t, err)
	err = node.sched.Start()
	assert.NoError(t, err)
	defer node.sched.Close()

	t.Run("describe database fail", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("DescribeDatabase", mock.Anything, mock.Anything).Return(nil, errors.New("fail"))
		node.mixCoord = mixc
		ctx := context.Background()
		resp, err := node.DescribeDatabase(ctx, &milvuspb.DescribeDatabaseRequest{})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetStatus().GetErrorCode())
	})

	t.Run("describe database ok", func(t *testing.T) {
		mixc := mocks.NewMockMixCoordClient(t)
		mixc.On("DescribeDatabase", mock.Anything, mock.Anything).Return(&rootcoordpb.DescribeDatabaseResponse{Status: merr.Success()}, nil)
		node.mixCoord = mixc
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		ctx := context.Background()

		resp, err := node.DescribeDatabase(ctx, &milvuspb.DescribeDatabaseRequest{})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})
}

func TestProxyDescribeCollection(t *testing.T) {
	paramtable.Init()
	node := &Proxy{session: &sessionutil.Session{SessionRaw: sessionutil.SessionRaw{ServerID: 1}}}
	ctx := context.Background()
	mixCoord := mocks.NewMockMixCoordClient(t)
	mixCoord.On("DescribeCollection", mock.Anything, mock.MatchedBy(func(req *milvuspb.DescribeCollectionRequest) bool {
		return req.DbName == "test_1" && req.CollectionName == "test_collection"
	})).Return(&milvuspb.DescribeCollectionResponse{
		Status:       merr.Success(),
		CollectionID: 1,
		Schema: &schemapb.CollectionSchema{
			Name: "test_collection",
			Fields: []*schemapb.FieldSchema{
				{Name: "pk", DataType: schemapb.DataType_Int64},
			},
		},
	}, nil).Maybe()
	mixCoord.On("ShowPartitions", mock.Anything, mock.Anything).Return(&milvuspb.ShowPartitionsResponse{
		Status:               merr.Success(),
		PartitionNames:       []string{"default"},
		CreatedTimestamps:    []uint64{1},
		CreatedUtcTimestamps: []uint64{1},
		PartitionIDs:         []int64{1},
	}, nil).Maybe()
	mixCoord.On("ShowLoadCollections", mock.Anything, mock.Anything).Return(&querypb.ShowCollectionsResponse{
		Status: merr.Success(),
	}, nil).Maybe()
	mixCoord.On("DescribeCollection", mock.Anything, mock.Anything).Return(nil, merr.ErrCollectionNotFound).Maybe()
	var err error
	globalMetaCache, err = NewMetaCache(mixCoord)
	assert.NoError(t, err)

	t.Run("not healthy", func(t *testing.T) {
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		defer node.UpdateStateCode(commonpb.StateCode_Healthy)
		resp, err := node.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{})
		assert.NoError(t, err)
		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("collection not exists", func(t *testing.T) {
		resp, err := node.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
			DbName:         "test_1",
			CollectionName: "test_not_exists",
		})
		assert.NoError(t, err)
		assert.Contains(t, resp.GetStatus().GetReason(), "can't find collection[database=test_1][collection=test_not_exists]")
		assert.Equal(t, commonpb.ErrorCode_CollectionNotExists, resp.GetStatus().GetErrorCode())
	})

	t.Run("collection id not exists", func(t *testing.T) {
		resp, err := node.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
			DbName:       "test_1",
			CollectionID: 1000,
		})
		assert.NoError(t, err)
		assert.Contains(t, resp.GetStatus().GetReason(), "can't find collection[database=test_1][collection=]")
		assert.Equal(t, commonpb.ErrorCode_CollectionNotExists, resp.GetStatus().GetErrorCode())
	})

	t.Run("db not exists", func(t *testing.T) {
		resp, err := node.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
			DbName:         "db_not_exists",
			CollectionName: "test_collection",
		})
		assert.NoError(t, err)
		assert.Contains(t, resp.GetStatus().GetReason(), "can't find collection[database=db_not_exists][collection=test_collection]")
		assert.Equal(t, commonpb.ErrorCode_CollectionNotExists, resp.GetStatus().GetErrorCode())
	})

	t.Run("describe collection ok", func(t *testing.T) {
		resp, err := node.DescribeCollection(ctx, &milvuspb.DescribeCollectionRequest{
			DbName:         "test_1",
			CollectionName: "test_collection",
		})
		assert.NoError(t, err)
		assert.Empty(t, resp.GetStatus().GetReason())
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("batch describe collection ok", func(t *testing.T) {
		resp, err := node.BatchDescribeCollection(ctx, &milvuspb.BatchDescribeCollectionRequest{
			DbName:         "test_1",
			CollectionName: []string{"test_collection"},
		})
		assert.NoError(t, err)
		assert.Empty(t, resp.GetStatus().GetReason())
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
		assert.Equal(t, 1, len(resp.GetResponses()))
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetResponses()[0].GetStatus().GetErrorCode())
	})
}

func TestProxy_AllocTimestamp(t *testing.T) {
	t.Run("proxy unhealthy", func(t *testing.T) {
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)

		resp, err := node.AllocTimestamp(context.TODO(), nil)
		assert.NoError(t, err)
		assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("success", func(t *testing.T) {
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		node.tsoAllocator = &timestampAllocator{
			tso: newMockTimestampAllocatorInterface(),
		}

		resp, err := node.AllocTimestamp(context.TODO(), nil)
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})

	t.Run("failed", func(t *testing.T) {
		m := newMockTimestampAllocator(t)
		m.On("AllocTimestamp",
			mock.Anything,
			mock.Anything,
		).Return(&rootcoordpb.AllocTimestampResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    "failed",
			},
			Timestamp: 20230518,
			Count:     1,
		}, nil)

		alloc, _ := newTimestampAllocator(m, 199)
		node := Proxy{
			tsoAllocator: alloc,
		}
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		resp, err := node.AllocTimestamp(context.TODO(), nil)
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetStatus().GetErrorCode())
	})
}

func TestProxy_Delete(t *testing.T) {
	collectionName := "test_delete"
	collectionID := int64(111)
	partitionName := "default"
	partitionID := int64(222)
	channels := []string{"test_vchannel"}
	dbName := "test_1"
	collSchema := &schemapb.CollectionSchema{
		Name:        collectionName,
		Description: "",
		AutoID:      false,
		Fields: []*schemapb.FieldSchema{
			{
				FieldID:      common.StartOfUserFieldID,
				Name:         "pk",
				IsPrimaryKey: true,
				DataType:     schemapb.DataType_Int64,
			},
			{
				FieldID:      common.StartOfUserFieldID + 1,
				Name:         "non_pk",
				IsPrimaryKey: false,
				DataType:     schemapb.DataType_Int64,
			},
		},
	}
	schema := newSchemaInfo(collSchema)
	basicInfo := &collectionInfo{
		collID: collectionID,
	}
	paramtable.Init()

	t.Run("delete run failed", func(t *testing.T) {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		chMgr := NewMockChannelsMgr(t)

		req := &milvuspb.DeleteRequest{
			CollectionName: collectionName,
			DbName:         dbName,
			PartitionName:  partitionName,
			Expr:           "pk in [1, 2, 3]",
		}
		cache := NewMockCache(t)
		cache.EXPECT().GetDatabaseInfo(mock.Anything, mock.Anything).Return(&databaseInfo{dbID: 0}, nil)
		cache.On("GetCollectionID",
			mock.Anything, // context.Context
			mock.AnythingOfType("string"),
			mock.AnythingOfType("string"),
		).Return(collectionID, nil)
		cache.On("GetCollectionSchema",
			mock.Anything, // context.Context
			mock.AnythingOfType("string"),
			mock.AnythingOfType("string"),
		).Return(schema, nil)
		cache.On("GetPartitionID",
			mock.Anything, // context.Context
			mock.AnythingOfType("string"),
			mock.AnythingOfType("string"),
			mock.AnythingOfType("string"),
		).Return(partitionID, nil)
		cache.On("GetCollectionInfo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(basicInfo, nil)
		chMgr.On("getVChannels", mock.Anything).Return(channels, nil)
		chMgr.On("getChannels", mock.Anything).Return(nil, errors.New("mock error"))
		globalMetaCache = cache
		rc := mocks.NewMockRootCoordClient(t)
		tsoAllocator := &mockTsoAllocator{}
		idAllocator, err := allocator.NewIDAllocator(ctx, rc, 0)
		assert.NoError(t, err)

		queue, err := newTaskScheduler(ctx, tsoAllocator)
		assert.NoError(t, err)

		node := &Proxy{chMgr: chMgr, rowIDAllocator: idAllocator, sched: queue}
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		resp, err := node.Delete(ctx, req)
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp.GetStatus()))
	})
}

func TestProxy_ImportV2(t *testing.T) {
	ctx := context.Background()
	mockErr := errors.New("mock error")

	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()

	t.Run("ImportV2", func(t *testing.T) {
		// server is not healthy
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		rsp, err := node.ImportV2(ctx, nil)
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		factory := dependency.NewDefaultFactory(true)
		node, err = NewProxy(ctx, factory)
		assert.NoError(t, err)
		node.UpdateStateCode(commonpb.StateCode_Healthy)
		node.tsoAllocator = &timestampAllocator{
			tso: newMockTimestampAllocatorInterface(),
		}
		scheduler, err := newTaskScheduler(ctx, node.tsoAllocator)
		assert.NoError(t, err)
		node.sched = scheduler
		err = node.sched.Start()
		assert.NoError(t, err)
		chMgr := NewMockChannelsMgr(t)
		node.chMgr = chMgr

		// no such collection
		// checkExternalCollectionBlockedForWrite skips error when GetCollectionSchema fails,
		// so task will be enqueued and PreExecute will call GetCollectionID which returns error.
		mc := NewMockCache(t)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(nil, mockErr).Once()
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, mockErr)
		globalMetaCache = mc
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{CollectionName: "aaa"})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// get schema failed in PreExecute
		mc = NewMockCache(t)
		// checkExternalCollectionBlockedForWrite skips error, task enqueued.
		// PreExecute calls GetCollectionID (succeeds), then GetCollectionSchema (fails).
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(nil, mockErr).Once()
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(nil, mockErr).Once()
		globalMetaCache = mc
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{CollectionName: "aaa"})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// schema has no fields
		mc = NewMockCache(t)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(&schemaInfo{
			CollectionSchema: &schemapb.CollectionSchema{},
		}, nil).Once()
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(&schemaInfo{
			CollectionSchema: &schemapb.CollectionSchema{},
		}, nil).Once()
		globalMetaCache = mc
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{CollectionName: "aaa"})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// get channel failed
		mc = NewMockCache(t)
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(&schemaInfo{
			CollectionSchema: &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
				{IsPartitionKey: true},
			}},
		}, nil)
		globalMetaCache = mc
		chMgr.EXPECT().getVChannels(mock.Anything).Return(nil, mockErr).Once()
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{CollectionName: "aaa"})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// set partition name and with partition key
		chMgr = NewMockChannelsMgr(t)
		chMgr.EXPECT().getVChannels(mock.Anything).Return([]string{"ch0"}, nil)
		node.chMgr = chMgr
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{CollectionName: "aaa", PartitionName: "bbb"})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// get partitions failed
		mc = NewMockCache(t)
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(&schemaInfo{
			CollectionSchema: &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
				{IsPartitionKey: true},
			}},
		}, nil)
		mc.EXPECT().GetPartitions(mock.Anything, mock.Anything, mock.Anything).Return(nil, mockErr)
		globalMetaCache = mc
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{CollectionName: "aaa"})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// get partitionID failed
		mc = NewMockCache(t)
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(&schemaInfo{
			CollectionSchema: &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{{FieldID: 1}}},
		}, nil)
		mc.EXPECT().GetPartitionID(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(0, mockErr)
		globalMetaCache = mc
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{CollectionName: "aaa", PartitionName: "bbb"})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// no file
		mc = NewMockCache(t)
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(&schemaInfo{
			CollectionSchema: &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{{FieldID: 1}}},
		}, nil)
		mc.EXPECT().GetPartitionID(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		globalMetaCache = mc
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{CollectionName: "aaa", PartitionName: "bbb"})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// illegal file type
		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{
			CollectionName: "aaa",
			PartitionName:  "bbb",
			Files: []*internalpb.ImportFile{{
				Id:    1,
				Paths: []string{"a.cpp"},
			}},
		})
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())

		// normal case
		mc = NewMockCache(t)
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		mc.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, mock.Anything).Return(&schemaInfo{
			CollectionSchema: &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{{FieldID: 1}}},
		}, nil)
		mc.EXPECT().GetPartitionID(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(int64(1), nil)
		mc.EXPECT().GetDatabaseInfo(mock.Anything, mock.Anything).Return(&databaseInfo{
			dbID: 1,
		}, nil)
		globalMetaCache = mc

		mixCoord := mocks.NewMockMixCoordClient(t)
		mixCoord.EXPECT().ImportV2(mock.Anything, mock.Anything).Return(&internalpb.ImportResponse{
			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
			JobID:  "123456789",
		}, nil)
		node.mixCoord = mixCoord

		rsp, err = node.ImportV2(ctx, &internalpb.ImportRequest{
			CollectionName: "aaa",
			Files: []*internalpb.ImportFile{{
				Id:    1,
				Paths: []string{"a.json"},
			}},
		})
		assert.NoError(t, err)
		assert.Equal(t, int32(0), rsp.GetStatus().GetCode())
		assert.Equal(t, "123456789", rsp.GetJobID())
	})

	t.Run("GetImportProgress", func(t *testing.T) {
		// server is not healthy
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		rsp, err := node.GetImportProgress(ctx, nil)
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		// normal case
		mixCoord := mocks.NewMockMixCoordClient(t)
		mixCoord.EXPECT().GetImportProgress(mock.Anything, mock.Anything).Return(nil, nil)
		node.mixCoord = mixCoord
		rsp, err = node.GetImportProgress(ctx, &internalpb.GetImportProgressRequest{})
		assert.NoError(t, err)
		assert.Equal(t, int32(0), rsp.GetStatus().GetCode())
	})

	t.Run("ListImports", func(t *testing.T) {
		// server is not healthy
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)
		rsp, err := node.ListImports(ctx, nil)
		assert.NoError(t, err)
		assert.NotEqual(t, int32(0), rsp.GetStatus().GetCode())
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		// normal case
		mc := NewMockCache(t)
		mc.EXPECT().GetCollectionID(mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
		globalMetaCache = mc
		mixCoord := mocks.NewMockMixCoordClient(t)
		mixCoord.EXPECT().ListImports(mock.Anything, mock.Anything).Return(nil, nil)
		node.mixCoord = mixCoord
		rsp, err = node.ListImports(ctx, &internalpb.ListImportsRequest{
			CollectionName: "col",
		})
		assert.NoError(t, err)
		assert.Equal(t, int32(0), rsp.GetStatus().GetCode())
	})
}

func TestGetCollectionRateSubLabel(t *testing.T) {
	d := "db1"
	collectionName := "test1"

	t.Run("normal", func(t *testing.T) {
		subLabel := GetCollectionRateSubLabel(&milvuspb.QueryRequest{
			DbName:         d,
			CollectionName: collectionName,
		})
		assert.Equal(t, ratelimitutil.GetCollectionSubLabel(d, collectionName), subLabel)
	})

	t.Run("fail", func(t *testing.T) {
		{
			subLabel := GetCollectionRateSubLabel(&milvuspb.QueryRequest{
				DbName:         "",
				CollectionName: collectionName,
			})
			assert.Equal(t, "", subLabel)
		}
		{
			subLabel := GetCollectionRateSubLabel(&milvuspb.QueryRequest{
				DbName:         d,
				CollectionName: "",
			})
			assert.Equal(t, "", subLabel)
		}
	})
}

func TestProxy_InvalidateShardLeaderCache(t *testing.T) {
	t.Run("proxy unhealthy", func(t *testing.T) {
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Abnormal)

		resp, err := node.InvalidateShardLeaderCache(context.TODO(), nil)
		assert.NoError(t, err)
		assert.False(t, merr.Ok(resp))
	})

	t.Run("success", func(t *testing.T) {
		node := &Proxy{}
		node.UpdateStateCode(commonpb.StateCode_Healthy)

		cacheBak := globalMetaCache
		defer func() { globalMetaCache = cacheBak }()
		// set expectations
		mockShardClientMgr := shardclient.NewMockShardClientManager(t)
		mockShardClientMgr.EXPECT().InvalidateShardLeaderCache(mock.Anything).Return()
		node.shardMgr = mockShardClientMgr

		resp, err := node.InvalidateShardLeaderCache(context.TODO(), &proxypb.InvalidateShardLeaderCacheRequest{})
		assert.NoError(t, err)
		assert.True(t, merr.Ok(resp))
	})
}

func TestRegisterRestRouter(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()

	mixc := mocks.NewMockMixCoordClient(t)
	mixc.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(nil, errors.New("error"))

	proxy := &Proxy{
		mixCoord: mixc,
	}
	proxy.RegisterRestRouter(router)

	tests := []struct {
		path       string
		statusCode int
	}{
		{path: mhttp.QCTargetPath, statusCode: http.StatusInternalServerError},
		{path: mhttp.QCDistPath, statusCode: http.StatusInternalServerError},
		{path: mhttp.QCAllTasksPath, statusCode: http.StatusInternalServerError},
		{path: mhttp.DNSyncTasksPath, statusCode: http.StatusInternalServerError},
		{path: mhttp.DCCompactionTasksPath, statusCode: http.StatusInternalServerError},
		{path: mhttp.DCImportTasksPath, statusCode: http.StatusInternalServerError},
		{path: mhttp.DCBuildIndexTasksPath, statusCode: http.StatusInternalServerError},
		{path: mhttp.DNSyncTasksPath, statusCode: http.StatusInternalServerError},
	}

	for _, tt := range tests {
		t.Run(tt.path, func(t *testing.T) {
			req, _ := http.NewRequest("GET", tt.path, nil)
			w := httptest.NewRecorder()
			router.ServeHTTP(w, req)
			assert.Equal(t, tt.statusCode, w.Code)
		})
	}
}

func TestRunAnalyzer(t *testing.T) {
	paramtable.Init()
	ctx := context.Background()

	cache := globalMetaCache
	globalMetaCache = nil
	defer func() { globalMetaCache = cache }()

	p := &Proxy{}

	tsoAllocatorIns := newMockTsoAllocator()
	sched, err := newTaskScheduler(ctx, tsoAllocatorIns)
	require.NoError(t, err)
	sched.Start()
	defer sched.Close()

	p.sched = sched

	t.Run("run analyzer err with node not healthy", func(t *testing.T) {
		p.UpdateStateCode(commonpb.StateCode_Abnormal)
		resp, err := p.RunAnalyzer(context.Background(), &milvuspb.RunAnalyzerRequest{
			Placeholder: [][]byte{[]byte("test doc")},
		})
		require.NoError(t, err)
		require.Error(t, merr.Error(resp.GetStatus()))
	})

	p.UpdateStateCode(commonpb.StateCode_Healthy)
	t.Run("run analyzer with mixcoord success", func(t *testing.T) {
		mockMixcoord := mocks.NewMockMixCoordClient(t)
		p.mixCoord = mockMixcoord
		mockMixcoord.EXPECT().RunAnalyzer(mock.Anything, mock.Anything, mock.Anything).Return(&milvuspb.RunAnalyzerResponse{Status: merr.Status(nil)}, nil)

		resp, err := p.RunAnalyzer(context.Background(), &milvuspb.RunAnalyzerRequest{
			Placeholder: [][]byte{[]byte("test doc")},
		})

		require.NoError(t, err)
		require.NoError(t, merr.Error(resp.GetStatus()))
	})

	t.Run("run analyzer with mixcoord failed", func(t *testing.T) {
		mockMixcoord := mocks.NewMockMixCoordClient(t)
		p.mixCoord = mockMixcoord
		mockMixcoord.EXPECT().RunAnalyzer(mock.Anything, mock.Anything, mock.Anything).Return(&milvuspb.RunAnalyzerResponse{Status: merr.Status(fmt.Errorf("mock error"))}, nil)

		resp, err := p.RunAnalyzer(context.Background(), &milvuspb.RunAnalyzerRequest{
			Placeholder: [][]byte{[]byte("test doc")},
		})
		require.NoError(t, err)
		require.Error(t, merr.Error(resp.GetStatus()))
	})

	t.Run("run analyzer from loaded collection field", func(t *testing.T) {
		mockCache := NewMockCache(t)
		globalMetaCache = mockCache

		fieldMap := &typeutil.ConcurrentMap[string, int64]{}
		fieldMap.Insert("test_text", 100)
		mockCache.EXPECT().GetCollectionID(mock.Anything, mock.Anything, "test_collection").Return(1, nil)
		mockCache.EXPECT().GetCollectionSchema(mock.Anything, mock.Anything, "test_collection").Return(&schemaInfo{
			CollectionSchema: &schemapb.CollectionSchema{
				Fields: []*schemapb.FieldSchema{{
					FieldID: 100,
					Name:    "test_text",
				}},
			},
			fieldMap: fieldMap,
		}, nil)

		lb := shardclient.NewMockLBPolicy(t)
		lb.EXPECT().ExecuteOneChannel(mock.Anything, mock.Anything).Return(nil)
		p.lbPolicy = lb

		resp, err := p.RunAnalyzer(context.Background(), &milvuspb.RunAnalyzerRequest{
			Placeholder:    [][]byte{[]byte("test doc")},
			CollectionName: "test_collection",
			FieldName:      "test_text",
		})

		require.NoError(t, err)
		require.NoError(t, merr.Error(resp.GetStatus()))
	})
}

func Test_GetSegmentsInfo(t *testing.T) {
	t.Run("normal case", func(t *testing.T) {
		mockMixCoord := mocks.NewMockMixCoordClient(t)
		mockMixCoord.EXPECT().GetSegmentInfo(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, request *datapb.GetSegmentInfoRequest, opts ...grpc.CallOption) (*datapb.GetSegmentInfoResponse, error) {
			segmentInfos := make([]*datapb.SegmentInfo, 0)
			for _, segID := range request.SegmentIDs {
				segmentInfos = append(segmentInfos, &datapb.SegmentInfo{
					ID:            segID,
					CollectionID:  1,
					PartitionID:   2,
					InsertChannel: "ch-1",
					NumOfRows:     1024,
					State:         commonpb.SegmentState_Flushed,
					MaxRowNum:     65535,
					Binlogs: []*datapb.FieldBinlog{
						{
							FieldID: 0,
							Binlogs: []*datapb.Binlog{
								{
									LogID: 1,
								},
								{
									LogID: 5,
								},
							},
						},
						{
							FieldID: 1,
							Binlogs: []*datapb.Binlog{
								{
									LogID: 2,
								},
								{
									LogID: 6,
								},
							},
						},
						{
							FieldID: 100,
							Binlogs: []*datapb.Binlog{
								{
									LogID: 3,
								},
								{
									LogID: 7,
								},
							},
						},
						{
							FieldID: 101,
							Binlogs: []*datapb.Binlog{
								{
									LogID: 4,
								},
								{
									LogID: 8,
								},
							},
						},
					},
					Statslogs: nil,
					Deltalogs: nil,
					Level:     datapb.SegmentLevel_L1,
					IsSorted:  true,
				})
			}

			return &datapb.GetSegmentInfoResponse{
				Status: merr.Success(),
				Infos:  segmentInfos,
			}, nil
		})

		ctx := context.Background()
		p := &Proxy{
			ctx:      ctx,
			mixCoord: mockMixCoord,
		}
		p.UpdateStateCode(commonpb.StateCode_Healthy)

		resp, err := p.GetSegmentsInfo(ctx, &internalpb.GetSegmentsInfoRequest{
			CollectionID: 1,
			SegmentIDs:   []int64{4, 5, 6},
		})
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.GetErrorCode())
		assert.Equal(t, 3, len(resp.GetSegmentInfos()))
		assert.Equal(t, 4, len(resp.GetSegmentInfos()[0].GetInsertLogs()))
		assert.Equal(t, int64(1), resp.GetSegmentInfos()[0].GetCollectionID())
		assert.Equal(t, int64(2), resp.GetSegmentInfos()[0].GetPartitionID())
		assert.Equal(t, "ch-1", resp.GetSegmentInfos()[0].GetVChannel())
		assert.ElementsMatch(t, []int64{1, 5}, resp.GetSegmentInfos()[0].GetInsertLogs()[0].GetLogIDs())
		assert.ElementsMatch(t, []int64{2, 6}, resp.GetSegmentInfos()[0].GetInsertLogs()[1].GetLogIDs())
		assert.ElementsMatch(t, []int64{3, 7}, resp.GetSegmentInfos()[0].GetInsertLogs()[2].GetLogIDs())
		assert.ElementsMatch(t, []int64{4, 8}, resp.GetSegmentInfos()[0].GetInsertLogs()[3].GetLogIDs())
	})
}

func TestProxy_AddFileResource(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Healthy)
		proxy.mixCoord = NewMixCoordMock()

		req := &milvuspb.AddFileResourceRequest{
			Base: &commonpb.MsgBase{},
			Name: "test_resource",
			Path: "/path/to/resource",
		}

		resp, err := proxy.AddFileResource(context.Background(), req)
		assert.NoError(t, err)
		assert.NoError(t, merr.Error(resp))
	})

	t.Run("proxy not healthy", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Abnormal)

		req := &milvuspb.AddFileResourceRequest{
			Name: "test_resource",
			Path: "/path/to/resource",
		}

		resp, err := proxy.AddFileResource(context.Background(), req)
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp))
	})

	t.Run("mixCoord error", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Healthy)

		mockMixCoord := mocks.NewMockMixCoordClient(t)
		mockMixCoord.EXPECT().AddFileResource(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("mock error"))
		proxy.mixCoord = mockMixCoord

		req := &milvuspb.AddFileResourceRequest{
			Name: "test_resource",
			Path: "/path/to/resource",
		}

		resp, err := proxy.AddFileResource(context.Background(), req)
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp))
	})
}

func TestProxy_RemoveFileResource(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Healthy)
		proxy.mixCoord = NewMixCoordMock()

		req := &milvuspb.RemoveFileResourceRequest{
			Base: &commonpb.MsgBase{},
			Name: "test_resource",
		}

		resp, err := proxy.RemoveFileResource(context.Background(), req)
		assert.NoError(t, err)
		assert.NoError(t, merr.Error(resp))
	})

	t.Run("proxy not healthy", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Abnormal)
		req := &milvuspb.RemoveFileResourceRequest{
			Name: "test_resource",
		}

		resp, err := proxy.RemoveFileResource(context.Background(), req)
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp))
	})

	t.Run("mixCoord error", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Healthy)

		mockMixCoord := mocks.NewMockMixCoordClient(t)
		mockMixCoord.EXPECT().RemoveFileResource(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("mock error"))
		proxy.mixCoord = mockMixCoord

		req := &milvuspb.RemoveFileResourceRequest{
			Name: "test_resource",
		}

		resp, err := proxy.RemoveFileResource(context.Background(), req)
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp))
	})
}

func TestProxy_ListFileResources(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Healthy)
		proxy.mixCoord = NewMixCoordMock()

		req := &milvuspb.ListFileResourcesRequest{
			Base: &commonpb.MsgBase{},
		}

		resp, err := proxy.ListFileResources(context.Background(), req)
		assert.NoError(t, err)
		assert.NoError(t, merr.Error(resp.GetStatus()))
	})

	t.Run("proxy not healthy", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Abnormal)

		req := &milvuspb.ListFileResourcesRequest{}
		resp, err := proxy.ListFileResources(context.Background(), req)
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp.GetStatus()))
	})

	t.Run("mixCoord error", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Healthy)

		mockMixCoord := mocks.NewMockMixCoordClient(t)
		mockMixCoord.EXPECT().ListFileResources(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("mock error"))
		proxy.mixCoord = mockMixCoord

		req := &milvuspb.ListFileResourcesRequest{}
		resp, err := proxy.ListFileResources(context.Background(), req)
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp.GetStatus()))
	})
}

func TestProxy_ComputePhraseMatchSlop(t *testing.T) {
	t.Run("proxy not healthy", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Abnormal)

		req := &milvuspb.ComputePhraseMatchSlopRequest{
			AnalyzerParams: `{"tokenizer": "standard"}`,
			QueryText:      "hello world",
			DataTexts:      []string{"hello world", "world hello"},
		}

		resp, err := proxy.ComputePhraseMatchSlop(context.Background(), req)
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp.GetStatus()))
	})

	t.Run("success_with_analyzer_params", func(t *testing.T) {
		// Create test proxy with mock mixCoord
		node := createTestProxy()
		defer node.sched.Close()

		// Mock mixCoord.ComputePhraseMatchSlop
		mockMixCoord := NewMixCoordMock()
		node.mixCoord = mockMixCoord

		req := &milvuspb.ComputePhraseMatchSlopRequest{
			AnalyzerParams: `{"tokenizer": "standard"}`,
			QueryText:      "hello world",
			DataTexts:      []string{"hello world", "world hello", "foo bar"},
		}

		resp, err := node.ComputePhraseMatchSlop(context.Background(), req)
		assert.NoError(t, err)
		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
	})
}

func TestHandleIfSearchByPK_BM25Detection(t *testing.T) {
	// Test that getBM25FunctionOfAnnsField correctly identifies BM25 functions
	t.Run("detect BM25 function by output field ID", func(t *testing.T) {
		functions := []*schemapb.FunctionSchema{
			{
				Name:             "bm25_func",
				Type:             schemapb.FunctionType_BM25,
				InputFieldNames:  []string{"text_field"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"sparse_field"},
				OutputFieldIds:   []int64{101},
			},
		}

		// Should find BM25 function when querying with output field ID
		fn, found := getBM25FunctionOfAnnsField(101, functions)
		assert.True(t, found)
		assert.Equal(t, "bm25_func", fn.Name)
		assert.Equal(t, schemapb.FunctionType_BM25, fn.Type)
		assert.Equal(t, "text_field", fn.InputFieldNames[0])

		// Should not find BM25 function when querying with input field ID
		_, found = getBM25FunctionOfAnnsField(100, functions)
		assert.False(t, found)

		// Should not find BM25 function when querying with non-existent field ID
		_, found = getBM25FunctionOfAnnsField(999, functions)
		assert.False(t, found)
	})

	t.Run("non-BM25 function should not be detected", func(t *testing.T) {
		functions := []*schemapb.FunctionSchema{
			{
				Name:             "embedding_func",
				Type:             schemapb.FunctionType_TextEmbedding,
				InputFieldNames:  []string{"text_field"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"vector_field"},
				OutputFieldIds:   []int64{101},
			},
		}

		// Should not find BM25 function for non-BM25 function type
		_, found := getBM25FunctionOfAnnsField(101, functions)
		assert.False(t, found)
	})

	t.Run("multiple functions with one BM25", func(t *testing.T) {
		functions := []*schemapb.FunctionSchema{
			{
				Name:             "embedding_func",
				Type:             schemapb.FunctionType_TextEmbedding,
				InputFieldNames:  []string{"text_field"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"vector_field"},
				OutputFieldIds:   []int64{102},
			},
			{
				Name:             "bm25_func",
				Type:             schemapb.FunctionType_BM25,
				InputFieldNames:  []string{"text_field"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"sparse_field"},
				OutputFieldIds:   []int64{101},
			},
		}

		// Should find BM25 function
		fn, found := getBM25FunctionOfAnnsField(101, functions)
		assert.True(t, found)
		assert.Equal(t, "bm25_func", fn.Name)

		// Should not find BM25 for embedding function output
		_, found = getBM25FunctionOfAnnsField(102, functions)
		assert.False(t, found)
	})
}

func TestHandleIfSearchByPK_PreservesNamespaceInInternalQuery(t *testing.T) {
	mockey.PatchConvey("TestHandleIfSearchByPK_PreservesNamespaceInInternalQuery", t, func() {
		paramtable.Init()
		originalCache := globalMetaCache
		defer func() { globalMetaCache = originalCache }()

		namespace := "tenant_a"
		schema := &schemapb.CollectionSchema{
			Name:            "test_collection",
			EnableNamespace: true,
			Fields: []*schemapb.FieldSchema{
				{FieldID: 100, Name: "id", IsPrimaryKey: true, DataType: schemapb.DataType_Int64},
				{
					FieldID:    101,
					Name:       "vec",
					DataType:   schemapb.DataType_FloatVector,
					TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "2"}},
				},
			},
		}

		cache := NewMockCache(t)
		cache.EXPECT().
			GetCollectionInfo(mock.Anything, "default", "test_collection", int64(0)).
			Return(&collectionInfo{schema: newSchemaInfo(schema)}, nil)
		globalMetaCache = cache

		var capturedNamespace *string
		mockey.Mock((*Proxy).query).To(func(_ *Proxy, _ context.Context, qt *queryTask, _ trace.Span) (*milvuspb.QueryResults, segcore.StorageCost, error) {
			capturedNamespace = qt.request.Namespace
			return &milvuspb.QueryResults{
				Status: merr.Success(),
				FieldsData: []*schemapb.FieldData{
					{
						FieldName: "id",
						FieldId:   100,
						Type:      schemapb.DataType_Int64,
						Field: &schemapb.FieldData_Scalars{
							Scalars: &schemapb.ScalarField{
								Data: &schemapb.ScalarField_LongData{
									LongData: &schemapb.LongArray{Data: []int64{1}},
								},
							},
						},
					},
					{
						FieldName: "vec",
						FieldId:   101,
						Type:      schemapb.DataType_FloatVector,
						Field: &schemapb.FieldData_Vectors{
							Vectors: &schemapb.VectorField{
								Dim: 2,
								Data: &schemapb.VectorField_FloatVector{
									FloatVector: &schemapb.FloatArray{Data: []float32{0.1, 0.2}},
								},
							},
						},
					},
				},
			}, segcore.StorageCost{}, nil
		}).Build()

		node := &Proxy{}
		req := &milvuspb.SearchRequest{
			DbName:         "default",
			CollectionName: "test_collection",
			Namespace:      &namespace,
			SearchInput: &milvuspb.SearchRequest_Ids{
				Ids: &schemapb.IDs{
					IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{1}}},
				},
			},
			SearchParams: []*commonpb.KeyValuePair{{Key: AnnsFieldKey, Value: "vec"}},
		}

		_, err := node.handleIfSearchByPK(context.Background(), req)
		assert.NoError(t, err)
		require.NotNil(t, capturedNamespace)
		assert.Equal(t, namespace, *capturedNamespace)
	})
}

func TestProxy_ManualCompaction_ExternalCollection(t *testing.T) {
	// Save and restore globalMetaCache
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	// Create external collection schema
	externalSchema := &schemapb.CollectionSchema{
		Name:           "external_col",
		ExternalSource: "s3://bucket/path",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionID).Return(int64(1), nil).Build()
	m2 := mockey.Mock((*MetaCache).GetCollectionInfo).Return(&collectionInfo{
		schema: newSchemaInfo(externalSchema),
	}, nil).Build()
	defer m1.UnPatch()
	defer m2.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &milvuspb.ManualCompactionRequest{
		DbName:         "default",
		CollectionName: "external_col",
	}

	resp, err := proxy.ManualCompaction(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp.GetStatus()))
}

func TestProxy_Insert_ExternalCollection(t *testing.T) {
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(externalSchema), nil).Build()
	defer m1.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &milvuspb.InsertRequest{
		DbName:         "default",
		CollectionName: "external_col",
	}

	resp, err := proxy.Insert(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp.GetStatus()))
	assert.Contains(t, resp.GetStatus().GetReason(), "insert operation is not supported for external collection")
}

func TestProxy_Delete_ExternalCollection(t *testing.T) {
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(externalSchema), nil).Build()
	defer m1.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &milvuspb.DeleteRequest{
		DbName:         "default",
		CollectionName: "external_col",
		Expr:           "id in [1, 2, 3]",
	}

	resp, err := proxy.Delete(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp.GetStatus()))
	assert.Contains(t, resp.GetStatus().GetReason(), "delete operation is not supported for external collection")
}

func TestProxy_Upsert_ExternalCollection(t *testing.T) {
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(externalSchema), nil).Build()
	defer m1.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &milvuspb.UpsertRequest{
		DbName:         "default",
		CollectionName: "external_col",
	}

	resp, err := proxy.Upsert(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp.GetStatus()))
	assert.Contains(t, resp.GetStatus().GetReason(), "upsert operation is not supported for external collection")
}

func TestProxy_Flush_ExternalCollection(t *testing.T) {
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(externalSchema), nil).Build()
	defer m1.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &milvuspb.FlushRequest{
		DbName:          "default",
		CollectionNames: []string{"external_col"},
	}

	resp, err := proxy.Flush(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp.GetStatus()))
	assert.Contains(t, resp.GetStatus().GetReason(), "flush operation is not supported for external collection")
}

func TestProxy_CreatePartition_ExternalCollection(t *testing.T) {
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(externalSchema), nil).Build()
	defer m1.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &milvuspb.CreatePartitionRequest{
		DbName:         "default",
		CollectionName: "external_col",
		PartitionName:  "partition1",
	}

	resp, err := proxy.CreatePartition(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp))
	assert.Contains(t, resp.GetReason(), "create partition operation is not supported for external collection")
}

func TestProxy_DropPartition_ExternalCollection(t *testing.T) {
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(externalSchema), nil).Build()
	defer m1.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &milvuspb.DropPartitionRequest{
		DbName:         "default",
		CollectionName: "external_col",
		PartitionName:  "partition1",
	}

	resp, err := proxy.DropPartition(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp))
	assert.Contains(t, resp.GetReason(), "drop partition operation is not supported for external collection")
}

func TestProxy_ImportV2_ExternalCollection(t *testing.T) {
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(externalSchema), nil).Build()
	defer m1.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &internalpb.ImportRequest{
		DbName:         "default",
		CollectionName: "external_col",
	}

	resp, err := proxy.ImportV2(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp.GetStatus()))
	assert.Contains(t, resp.GetStatus().GetReason(), "import operation is not supported for external collection")
}

func TestProxy_AddCollectionField_ExternalCollection(t *testing.T) {
	node := createTestProxy()
	defer node.sched.Close()

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}
	mockDescribe := mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
		Status: merr.Success(),
		Schema: externalSchema,
	}, nil).Build()
	defer mockDescribe.UnPatch()

	mockEnqueue := mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, queued task) error {
		_ = queued.OnEnqueue()
		addTask := queued.(*addCollectionFieldTask)
		assert.Equal(t, externalSchema, addTask.oldSchema)
		addTask.result = merr.Success()
		return nil
	}).Build()
	defer mockEnqueue.UnPatch()

	mockWait := mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()
	defer mockWait.UnPatch()

	resp, err := node.AddCollectionField(context.Background(), &milvuspb.AddCollectionFieldRequest{
		DbName:         "default",
		CollectionName: "external_col",
		Schema:         []byte{1},
	})
	assert.NoError(t, err)
	assert.True(t, merr.Ok(resp))
}

func TestProxy_AddCollectionField_TextValidation(t *testing.T) {
	baseSchema := &schemapb.CollectionSchema{
		Name: "test_coll",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
			{
				FieldID:    101,
				Name:       "vec",
				DataType:   schemapb.DataType_FloatVector,
				TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "128"}},
			},
		},
	}

	for _, tc := range []struct {
		name        string
		storageV3   string
		field       *schemapb.FieldSchema
		wantErr     bool
		errContains string
	}{
		{
			name:      "text nullable allowed when storage v3 enabled",
			storageV3: "true",
			field: &schemapb.FieldSchema{
				Name:     "text_field",
				DataType: schemapb.DataType_Text,
				Nullable: true,
			},
		},
		{
			name:      "text requires storage v3",
			storageV3: "false",
			field: &schemapb.FieldSchema{
				Name:     "text_field",
				DataType: schemapb.DataType_Text,
				Nullable: true,
			},
			wantErr:     true,
			errContains: "TEXT field requires StorageV3",
		},
		{
			name:      "text must be nullable",
			storageV3: "true",
			field: &schemapb.FieldSchema{
				Name:     "text_not_nullable",
				DataType: schemapb.DataType_Text,
				Nullable: false,
			},
			wantErr:     true,
			errContains: "added field must be nullable",
		},
		{
			name:      "text default value rejected",
			storageV3: "true",
			field: &schemapb.FieldSchema{
				Name:     "text_default",
				DataType: schemapb.DataType_Text,
				Nullable: true,
				DefaultValue: &schemapb.ValueField{
					Data: &schemapb.ValueField_StringData{StringData: "default text"},
				},
			},
			wantErr:     true,
			errContains: "default value is not supported when adding TEXT field",
		},
	} {
		t.Run(tc.name, func(t *testing.T) {
			paramtable.Get().Save(paramtable.Get().CommonCfg.UseLoonFFI.Key, tc.storageV3)
			t.Cleanup(func() {
				paramtable.Get().Reset(paramtable.Get().CommonCfg.UseLoonFFI.Key)
			})

			node := createTestProxy()
			defer node.sched.Close()

			mockDescribe := mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: proto.Clone(baseSchema).(*schemapb.CollectionSchema),
			}, nil).Build()
			defer mockDescribe.UnPatch()

			mockEnqueue := mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, queued task) error {
				require.NoError(t, queued.OnEnqueue())
				addTask := queued.(*addCollectionFieldTask)
				err := addTask.PreExecute(context.Background())
				if err != nil {
					addTask.result = merr.Status(err)
				} else {
					addTask.result = merr.Success()
				}
				addTask.Notify(err)
				return nil
			}).Build()
			defer mockEnqueue.UnPatch()

			fieldBytes, err := proto.Marshal(tc.field)
			require.NoError(t, err)
			resp, err := node.AddCollectionField(context.Background(), &milvuspb.AddCollectionFieldRequest{
				DbName:         "default",
				CollectionName: "test_coll",
				Schema:         fieldBytes,
			})
			require.NoError(t, err)
			if tc.wantErr {
				require.Error(t, merr.Error(resp))
				require.Contains(t, resp.GetReason(), tc.errContains)
				return
			}
			require.True(t, merr.Ok(resp), resp.GetReason())
		})
	}
}

func TestProxy_AddCollectionField_DoesNotBlockOnSchemaVersion(t *testing.T) {
	t.Run("does not query schema version stats before enqueue", func(t *testing.T) {
		mockey.PatchConvey("AddCollectionField skips schema version stats gate", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: &schemapb.CollectionSchema{Name: "test_coll"},
			}, nil).Build()
			mockey.Mock((*Proxy).GetCollectionStatistics).To(
				func(*Proxy, context.Context, *milvuspb.GetCollectionStatisticsRequest) (*milvuspb.GetCollectionStatisticsResponse, error) {
					require.FailNow(t, "AddCollectionField should not query collection statistics")
					return nil, errors.New("unexpected GetCollectionStatistics call")
				}).Build()
			mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
				_ = t.OnEnqueue()
				addTask := t.(*addCollectionFieldTask)
				addTask.result = merr.Success()
				return nil
			}).Build()
			mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()

			resp, err := node.AddCollectionField(context.Background(), &milvuspb.AddCollectionFieldRequest{
				DbName:         "default",
				CollectionName: "test_coll",
			})
			assert.NoError(t, err)
			assert.True(t, merr.Ok(resp))
		})
	})

	t.Run("allows overlapping requests for same collection", func(t *testing.T) {
		mockey.PatchConvey("AddCollectionField has no proxy-local in-flight gate", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: &schemapb.CollectionSchema{Name: "test_coll"},
			}, nil).Build()
			mockey.Mock((*Proxy).GetCollectionStatistics).Return(&milvuspb.GetCollectionStatisticsResponse{
				Status: merr.Success(),
			}, nil).Build()
			mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
				_ = t.OnEnqueue()
				addTask := t.(*addCollectionFieldTask)
				addTask.result = merr.Success()
				return nil
			}).Build()

			firstWaitStarted := make(chan struct{})
			releaseFirst := make(chan struct{})
			var waitCalls atomic.Int32
			mockey.Mock((*TaskCondition).WaitToFinish).To(func(*TaskCondition) error {
				if waitCalls.Add(1) == 1 {
					close(firstWaitStarted)
					<-releaseFirst
				}
				return nil
			}).Build()

			firstDone := make(chan *commonpb.Status, 1)
			go func() {
				resp, err := node.AddCollectionField(context.Background(), &milvuspb.AddCollectionFieldRequest{
					DbName:         "default",
					CollectionName: "test_coll",
				})
				require.NoError(t, err)
				firstDone <- resp
			}()
			<-firstWaitStarted

			resp, err := node.AddCollectionField(context.Background(), &milvuspb.AddCollectionFieldRequest{
				DbName:         "default",
				CollectionName: "test_coll",
			})
			assert.NoError(t, err)
			assert.True(t, merr.Ok(resp))
			assert.Equal(t, int32(2), waitCalls.Load())

			close(releaseFirst)
			assert.True(t, merr.Ok(<-firstDone))
		})
	})
}

func TestProxy_AlterCollectionField_ExternalCollection(t *testing.T) {
	cache := globalMetaCache
	defer func() { globalMetaCache = cache }()
	globalMetaCache = &MetaCache{}

	externalSchema := &schemapb.CollectionSchema{
		Name: "external_col",
		Fields: []*schemapb.FieldSchema{
			{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
		},
	}

	m1 := mockey.Mock((*MetaCache).GetCollectionSchema).Return(newSchemaInfo(externalSchema), nil).Build()
	defer m1.UnPatch()

	proxy := &Proxy{}
	proxy.UpdateStateCode(commonpb.StateCode_Healthy)

	req := &milvuspb.AlterCollectionFieldRequest{
		DbName:         "default",
		CollectionName: "external_col",
		FieldName:      "id",
	}

	resp, err := proxy.AlterCollectionField(context.Background(), req)
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp))
	assert.Contains(t, resp.GetReason(), "alter field operation is not supported for external collection")
}

func TestProxy_GetReplicateConfiguration_Success(t *testing.T) {
	expectedConfig := &commonpb.ReplicateConfiguration{
		Clusters: []*commonpb.MilvusCluster{
			{
				ClusterId: "primary",
				ConnectionParam: &commonpb.ConnectionParam{
					Uri: "http://primary:19530",
				},
				Pchannels: []string{"channel1"},
			},
		},
		CrossClusterTopology: []*commonpb.CrossClusterTopology{
			{SourceClusterId: "primary", TargetClusterId: "secondary"},
		},
	}

	replicateService := mock_streaming.NewMockReplicateService(t)
	replicateService.EXPECT().GetReplicateConfiguration(mock.Anything).Return(expectedConfig, nil)
	mockWAL := mock_streaming.NewMockWALAccesser(t)
	mockWAL.EXPECT().Replicate().Return(replicateService)
	streaming.SetWALForTest(mockWAL)
	defer streaming.SetWALForTest(nil)

	node := &Proxy{}
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	resp, err := node.GetReplicateConfiguration(context.Background(), &milvuspb.GetReplicateConfigurationRequest{})
	assert.NoError(t, err)
	assert.True(t, merr.Ok(resp.GetStatus()))
	assert.NotNil(t, resp.GetConfiguration())
	assert.Equal(t, "primary", resp.GetConfiguration().GetClusters()[0].GetClusterId())
	assert.Len(t, resp.GetConfiguration().GetCrossClusterTopology(), 1)
}

func TestProxy_GetReplicateConfiguration_Unhealthy(t *testing.T) {
	node := &Proxy{}
	node.UpdateStateCode(commonpb.StateCode_Abnormal)

	resp, err := node.GetReplicateConfiguration(context.Background(), &milvuspb.GetReplicateConfigurationRequest{})
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp.GetStatus()))
}

func TestProxy_GetReplicateConfiguration_Error(t *testing.T) {
	replicateService := mock_streaming.NewMockReplicateService(t)
	replicateService.EXPECT().GetReplicateConfiguration(mock.Anything).Return(nil, errors.New("streaming error"))
	mockWAL := mock_streaming.NewMockWALAccesser(t)
	mockWAL.EXPECT().Replicate().Return(replicateService)
	streaming.SetWALForTest(mockWAL)
	defer streaming.SetWALForTest(nil)

	node := &Proxy{}
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	resp, err := node.GetReplicateConfiguration(context.Background(), &milvuspb.GetReplicateConfigurationRequest{})
	assert.NoError(t, err)
	assert.Error(t, merr.Error(resp.GetStatus()))
	assert.Nil(t, resp.GetConfiguration())
}

func TestHybridSearchRequestExprLogger_String(t *testing.T) {
	t.Run("empty requests", func(t *testing.T) {
		logger := &hybridSearchRequestExprLogger{
			req: &milvuspb.HybridSearchRequest{
				Requests: []*milvuspb.SearchRequest{},
			},
		}
		result := logger.String()
		assert.Equal(t, "", result)
	})

	t.Run("single request", func(t *testing.T) {
		logger := &hybridSearchRequestExprLogger{
			req: &milvuspb.HybridSearchRequest{
				Requests: []*milvuspb.SearchRequest{
					{Dsl: "id > 100"},
				},
			},
		}
		result := logger.String()
		assert.Equal(t, "[No.0 req, expr: id > 100]", result)
	})

	t.Run("multiple requests", func(t *testing.T) {
		logger := &hybridSearchRequestExprLogger{
			req: &milvuspb.HybridSearchRequest{
				Requests: []*milvuspb.SearchRequest{
					{Dsl: "id > 100"},
					{Dsl: "name == 'test'"},
					{Dsl: "age < 30"},
				},
			},
		}
		result := logger.String()
		expected := "[No.0 req, expr: id > 100][No.1 req, expr: name == 'test'][No.2 req, expr: age < 30]"
		assert.Equal(t, expected, result)
	})

	t.Run("request with empty dsl", func(t *testing.T) {
		logger := &hybridSearchRequestExprLogger{
			req: &milvuspb.HybridSearchRequest{
				Requests: []*milvuspb.SearchRequest{
					{Dsl: ""},
				},
			},
		}
		result := logger.String()
		assert.Equal(t, "[No.0 req, expr: ]", result)
	})

	t.Run("nil sub request in slice", func(t *testing.T) {
		logger := &hybridSearchRequestExprLogger{
			req: &milvuspb.HybridSearchRequest{
				Requests: []*milvuspb.SearchRequest{
					nil,
				},
			},
		}
		result := logger.String()
		assert.Equal(t, "[No.0 req, expr: ]", result)
	})
}

func TestProxy_BatchUpdateManifest(t *testing.T) {
	t.Run("unhealthy", func(t *testing.T) {
		mockey.PatchConvey("TestProxy_BatchUpdateManifest_unhealthy", t, func() {
			globalMetaCache = &MetaCache{}
			mockey.Mock(globalMetaCache.GetCollectionID).To(func(ctx context.Context, dbName, collectionName string) (UniqueID, error) {
				return UniqueID(0), nil
			}).Build()
			mockey.Mock(globalMetaCache.RemoveDatabase).To(func(ctx context.Context, dbName string) {}).Build()

			mockey.Mock(paramtable.Init).Return().Build()
			mockey.Mock((*paramtable.ComponentParam).Save).Return().Build()

			node := createTestProxy()
			defer node.sched.Close()

			node.UpdateStateCode(commonpb.StateCode_Abnormal)
			resp, err := node.BatchUpdateManifest(context.Background(), &milvuspb.BatchUpdateManifestRequest{
				CollectionName: "test_collection",
				Items: []*milvuspb.BatchUpdateManifestItem{
					{SegmentId: 1, ManifestVersion: 10},
				},
			})

			assert.NoError(t, err)
			assert.ErrorIs(t, merr.Error(resp), merr.ErrServiceNotReady)
		})
	})

	t.Run("success", func(t *testing.T) {
		mockey.PatchConvey("TestProxy_BatchUpdateManifest_success", t, func() {
			globalMetaCache = &MetaCache{}
			mockey.Mock((*MetaCache).GetCollectionID).To(func(m *MetaCache, ctx context.Context, dbName, collectionName string) (UniqueID, error) {
				return UniqueID(100), nil
			}).Build()
			mockey.Mock((*MetaCache).RemoveDatabase).To(func(m *MetaCache, ctx context.Context, dbName string) {}).Build()

			mockey.Mock(paramtable.Init).Return().Build()
			mockey.Mock((*paramtable.ComponentParam).Save).Return().Build()

			node := createTestProxy()
			defer node.sched.Close()

			mixcoord := &grpcmixcoordclient.Client{}
			node.mixCoord = mixcoord
			mockey.Mock((*grpcmixcoordclient.Client).BatchUpdateManifest).To(func(c *grpcmixcoordclient.Client, ctx context.Context, req *datapb.BatchUpdateManifestRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
				return merr.Success(), nil
			}).Build()

			resp, err := node.BatchUpdateManifest(context.Background(), &milvuspb.BatchUpdateManifestRequest{
				CollectionName: "test_collection",
				Items: []*milvuspb.BatchUpdateManifestItem{
					{SegmentId: 1, ManifestVersion: 10},
					{SegmentId: 2, ManifestVersion: 20},
				},
			})

			assert.NoError(t, err)
			assert.True(t, merr.Ok(resp))
		})
	})
}

func TestProxy_AlterCollectionSchema(t *testing.T) {
	t.Run("unhealthy node", func(t *testing.T) {
		proxy := &Proxy{}
		proxy.UpdateStateCode(commonpb.StateCode_Abnormal)
		resp, err := proxy.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{})
		assert.NoError(t, err)
		assert.Error(t, merr.Error(resp.GetAlterStatus()))
	})

	t.Run("DescribeCollection fails", func(t *testing.T) {
		mockey.PatchConvey("DescribeCollection fails", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Status(merr.ErrCollectionNotFound),
			}, nil).Build()

			resp, err := node.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{
				CollectionName: "test_coll",
			})
			assert.NoError(t, err)
			assert.Error(t, merr.Error(resp.GetAlterStatus()))
		})
	})

	t.Run("external collection rejected", func(t *testing.T) {
		mockey.PatchConvey("external collection", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: &schemapb.CollectionSchema{
					Name: "ext_col",
					Fields: []*schemapb.FieldSchema{
						{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "ext_id"},
					},
				},
			}, nil).Build()

			resp, err := node.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{
				CollectionName: "ext_col",
			})
			assert.NoError(t, err)
			assert.Error(t, merr.Error(resp.GetAlterStatus()))
			assert.Contains(t, resp.GetAlterStatus().GetReason(), "external collection")
		})
	})

	t.Run("does not query schema version stats before enqueue", func(t *testing.T) {
		mockey.PatchConvey("AlterCollectionSchema skips schema version stats gate", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: &schemapb.CollectionSchema{Name: "test_coll"},
			}, nil).Build()
			mockey.Mock((*Proxy).GetCollectionStatistics).To(
				func(*Proxy, context.Context, *milvuspb.GetCollectionStatisticsRequest) (*milvuspb.GetCollectionStatisticsResponse, error) {
					require.FailNow(t, "AlterCollectionSchema should not query collection statistics")
					return nil, errors.New("unexpected GetCollectionStatistics call")
				}).Build()
			mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
				_ = t.OnEnqueue()
				alterTask := t.(*alterCollectionSchemaTask)
				alterTask.AlterCollectionSchemaResponse = &milvuspb.AlterCollectionSchemaResponse{AlterStatus: merr.Success()}
				return nil
			}).Build()
			mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()

			resp, err := node.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{
				DbName:         "default",
				CollectionName: "test_coll",
			})
			assert.NoError(t, err)
			assert.True(t, merr.Ok(resp.GetAlterStatus()))
		})
	})

	t.Run("allows overlapping requests for same collection", func(t *testing.T) {
		mockey.PatchConvey("AlterCollectionSchema has no proxy-local in-flight gate", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: &schemapb.CollectionSchema{Name: "test_coll"},
			}, nil).Build()
			mockey.Mock((*Proxy).GetCollectionStatistics).Return(&milvuspb.GetCollectionStatisticsResponse{
				Status: merr.Success(),
			}, nil).Build()
			mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
				_ = t.OnEnqueue()
				alterTask := t.(*alterCollectionSchemaTask)
				alterTask.AlterCollectionSchemaResponse = &milvuspb.AlterCollectionSchemaResponse{AlterStatus: merr.Success()}
				return nil
			}).Build()

			firstWaitStarted := make(chan struct{})
			releaseFirst := make(chan struct{})
			var waitCalls atomic.Int32
			mockey.Mock((*TaskCondition).WaitToFinish).To(func(*TaskCondition) error {
				if waitCalls.Add(1) == 1 {
					close(firstWaitStarted)
					<-releaseFirst
				}
				return nil
			}).Build()

			firstDone := make(chan *milvuspb.AlterCollectionSchemaResponse, 1)
			go func() {
				resp, err := node.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{
					DbName:         "default",
					CollectionName: "test_coll",
				})
				require.NoError(t, err)
				firstDone <- resp
			}()
			<-firstWaitStarted

			resp, err := node.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{
				DbName:         "default",
				CollectionName: "test_coll",
			})
			assert.NoError(t, err)
			assert.True(t, merr.Ok(resp.GetAlterStatus()))
			assert.Equal(t, int32(2), waitCalls.Load())

			close(releaseFirst)
			assert.True(t, merr.Ok((<-firstDone).GetAlterStatus()))
		})
	})

	t.Run("enqueue fails", func(t *testing.T) {
		mockey.PatchConvey("enqueue fails", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: &schemapb.CollectionSchema{Name: "test_coll"},
			}, nil).Build()

			mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, _ task) error {
				return errors.New("queue full")
			}).Build()

			resp, err := node.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{
				CollectionName: "test_coll",
			})
			assert.NoError(t, err)
			assert.Error(t, merr.Error(resp.GetAlterStatus()))
		})
	})

	t.Run("WaitToFinish fails", func(t *testing.T) {
		mockey.PatchConvey("WaitToFinish fails", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: &schemapb.CollectionSchema{Name: "test_coll"},
			}, nil).Build()

			// Call OnEnqueue so task.Base is initialized (BeginTs/EndTs are logged after Enqueue).
			mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
				_ = t.OnEnqueue()
				return nil
			}).Build()

			mockey.Mock((*TaskCondition).WaitToFinish).Return(errors.New("timeout")).Build()

			resp, err := node.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{
				CollectionName: "test_coll",
			})
			assert.NoError(t, err)
			assert.Error(t, merr.Error(resp.GetAlterStatus()))
		})
	})

	t.Run("happy path", func(t *testing.T) {
		mockey.PatchConvey("happy path", t, func() {
			node := createTestProxy()
			defer node.sched.Close()

			mockey.Mock((*Proxy).DescribeCollection).Return(&milvuspb.DescribeCollectionResponse{
				Status: merr.Success(),
				Schema: &schemapb.CollectionSchema{Name: "test_coll"},
			}, nil).Build()

			// Call OnEnqueue so task.Base is initialized (BeginTs/EndTs are logged after Enqueue).
			mockey.Mock((*ddTaskQueue).Enqueue).To(func(_ *ddTaskQueue, t task) error {
				_ = t.OnEnqueue()
				return nil
			}).Build()

			mockey.Mock((*TaskCondition).WaitToFinish).Return(nil).Build()

			_, err := node.AlterCollectionSchema(context.Background(), &milvuspb.AlterCollectionSchemaRequest{
				CollectionName: "test_coll",
			})
			assert.NoError(t, err)
		})
	})
}

func TestProxy_RefreshExternalCollection_AtomicSourceSpec(t *testing.T) {
	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()
	node, err := NewProxy(ctx, factory)
	require.NoError(t, err)
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	cases := []struct {
		name       string
		src, spec  string
		wantSubstr string
	}{
		{"source only rejected", "s3://bucket/p", "", "both provided or both omitted"},
		{"spec only rejected", "", `{"format":"parquet"}`, "both provided or both omitted"},
		{"http scheme rejected", "http://169.254.169.254/metadata", `{"format":"parquet"}`, "scheme"},
		{"file scheme rejected", "file:///etc/passwd", `{"format":"parquet"}`, "scheme"},
		{"ftp scheme rejected", "ftp://internal/data", `{"format":"parquet"}`, "scheme"},
		{"unknown scheme rejected", "xyz://nope/", `{"format":"parquet"}`, "scheme"},
		{"userinfo rejected", "s3://ak:sk@bucket/prefix", `{"format":"parquet"}`, ""},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			resp, err := node.RefreshExternalCollection(ctx, &milvuspb.RefreshExternalCollectionRequest{
				CollectionName: "any_collection",
				ExternalSource: tc.src,
				ExternalSpec:   tc.spec,
			})
			require.NoError(t, err)
			require.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrParameterInvalid)
			if tc.wantSubstr != "" {
				assert.Contains(t, resp.GetStatus().GetReason(), tc.wantSubstr)
			}
		})
	}
}

func TestProxy_RefreshExternalCollection_ReusePathRequiresPersistedSourceSpec(t *testing.T) {
	factory := dependency.NewDefaultFactory(true)
	ctx := context.Background()
	node, err := NewProxy(ctx, factory)
	require.NoError(t, err)
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	mkInfo := func(src, spec string) *collectionInfo {
		schema := &schemapb.CollectionSchema{
			Name:           "demo",
			ExternalSource: src,
			ExternalSpec:   spec,
			Fields: []*schemapb.FieldSchema{
				{Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "id"},
			},
		}
		return &collectionInfo{
			collID: 1,
			schema: &schemaInfo{CollectionSchema: schema},
		}
	}

	cases := []struct {
		name      string
		persisted *collectionInfo
		wantErr   bool
	}{
		{"persisted both empty rejected", mkInfo("", ""), true},
		{"persisted source empty rejected", mkInfo("", `{"format":"parquet"}`), true},
		{"persisted spec empty rejected", mkInfo("s3://bucket/p", ""), true},
	}
	cacheBak := globalMetaCache
	defer func() { globalMetaCache = cacheBak }()
	globalMetaCache = &MetaCache{}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			m := mockey.Mock((*MetaCache).GetCollectionInfo).Return(tc.persisted, nil).Build()
			defer m.UnPatch()

			resp, err := node.RefreshExternalCollection(ctx, &milvuspb.RefreshExternalCollectionRequest{
				CollectionName: "demo",
			})
			require.NoError(t, err)
			if tc.wantErr {
				require.ErrorIs(t, merr.Error(resp.GetStatus()), merr.ErrParameterInvalid)
				assert.Contains(t, resp.GetStatus().GetReason(), "no persisted external_source")
			} else {
				require.NoError(t, merr.Error(resp.GetStatus()))
			}
		})
	}
}

func TestProxy_ListRefreshExternalCollectionJobs_ListAll(t *testing.T) {
	paramtable.Init()
	ctx := context.Background()
	node := &Proxy{mixCoord: &MixCoordMock{}}
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	cacheBak := globalMetaCache
	globalMetaCache = &MetaCache{}
	defer func() { globalMetaCache = cacheBak }()

	cacheCalled := false
	mockGetCollectionInfo := mockey.Mock((*MetaCache).GetCollectionInfo).To(
		func(_ *MetaCache, _ context.Context, _ string, _ string, _ int64) (*collectionInfo, error) {
			cacheCalled = true
			return nil, errors.New("collection cache should not be used for list all")
		}).Build()
	defer mockGetCollectionInfo.UnPatch()

	var capturedCollectionID int64
	mockList := mockey.Mock((*MixCoordMock).ListRefreshExternalCollectionJobs).To(
		func(_ *MixCoordMock, _ context.Context, req *datapb.ListRefreshExternalCollectionJobsRequest, _ ...grpc.CallOption) (*datapb.ListRefreshExternalCollectionJobsResponse, error) {
			capturedCollectionID = req.GetCollectionId()
			return &datapb.ListRefreshExternalCollectionJobsResponse{
				Status: merr.Success(),
				Jobs: []*datapb.ExternalCollectionRefreshJob{
					{
						JobId:          1,
						CollectionName: "allowed",
						State:          indexpb.JobState_JobStateFinished,
						Progress:       100,
					},
					{
						JobId:          2,
						CollectionName: "blocked",
						State:          indexpb.JobState_JobStateFinished,
						Progress:       100,
					},
				},
			}, nil
		}).Build()
	defer mockList.UnPatch()

	resp, err := node.ListRefreshExternalCollectionJobs(ctx, &milvuspb.ListRefreshExternalCollectionJobsRequest{
		DbName: "default",
	})

	require.NoError(t, err)
	require.True(t, merr.Ok(resp.GetStatus()))
	assert.Equal(t, int64(0), capturedCollectionID)
	assert.False(t, cacheCalled)
	require.Len(t, resp.GetJobs(), 2)
	assert.Equal(t, int64(1), resp.GetJobs()[0].GetJobId())
	assert.Equal(t, "allowed", resp.GetJobs()[0].GetCollectionName())
	assert.Equal(t, int64(2), resp.GetJobs()[1].GetJobId())
	assert.Equal(t, "blocked", resp.GetJobs()[1].GetCollectionName())
}

func TestProxy_ListRefreshExternalCollectionJobs_ByCollection(t *testing.T) {
	paramtable.Init()
	ctx := context.Background()
	node := &Proxy{mixCoord: &MixCoordMock{}}
	node.UpdateStateCode(commonpb.StateCode_Healthy)

	cacheBak := globalMetaCache
	globalMetaCache = &MetaCache{}
	defer func() { globalMetaCache = cacheBak }()

	schema := &schemapb.CollectionSchema{
		Name: "external_collection",
		Fields: []*schemapb.FieldSchema{
			{Name: "id", DataType: schemapb.DataType_Int64, ExternalField: "id"},
		},
	}
	mockGetCollectionInfo := mockey.Mock((*MetaCache).GetCollectionInfo).Return(&collectionInfo{
		collID: 101,
		schema: &schemaInfo{
			CollectionSchema: schema,
		},
	}, nil).Build()
	defer mockGetCollectionInfo.UnPatch()

	var capturedCollectionID int64
	mockList := mockey.Mock((*MixCoordMock).ListRefreshExternalCollectionJobs).To(
		func(_ *MixCoordMock, _ context.Context, req *datapb.ListRefreshExternalCollectionJobsRequest, _ ...grpc.CallOption) (*datapb.ListRefreshExternalCollectionJobsResponse, error) {
			capturedCollectionID = req.GetCollectionId()
			return &datapb.ListRefreshExternalCollectionJobsResponse{
				Status: merr.Success(),
				Jobs: []*datapb.ExternalCollectionRefreshJob{
					{
						JobId:          1,
						CollectionName: "external_collection",
						State:          indexpb.JobState_JobStateFinished,
						Progress:       100,
						ExternalSpec:   `{"format":"parquet","extfs":{"access_key_id":"AKIAEXAMPLE","access_key_value":"SUPERSECRET","region":"us-east-1"}}`,
					},
				},
			}, nil
		}).Build()
	defer mockList.UnPatch()

	resp, err := node.ListRefreshExternalCollectionJobs(ctx, &milvuspb.ListRefreshExternalCollectionJobsRequest{
		DbName:         "default",
		CollectionName: "external_collection",
	})

	require.NoError(t, err)
	require.True(t, merr.Ok(resp.GetStatus()))
	assert.Equal(t, int64(101), capturedCollectionID)
	require.Len(t, resp.GetJobs(), 1)
	redactedSpec := resp.GetJobs()[0].GetExternalSpec()
	assert.Contains(t, redactedSpec, `"access_key_id":"***"`)
	assert.Contains(t, redactedSpec, `"access_key_value":"***"`)
	assert.Contains(t, redactedSpec, `"region":"us-east-1"`)
	assert.NotContains(t, redactedSpec, "AKIAEXAMPLE")
	assert.NotContains(t, redactedSpec, "SUPERSECRET")
}