milvus-io--milvus/internal/querynodev2/delegator/delegator_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 delegator

import (
	"context"
	"io"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/bytedance/mockey"
	"github.com/cockroachdb/errors"
	"github.com/samber/lo"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
	"go.uber.org/atomic"
	"google.golang.org/protobuf/proto"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/internal/distributed/streaming"
	"github.com/milvus-io/milvus/internal/querynodev2/cluster"
	"github.com/milvus-io/milvus/internal/querynodev2/delegator/deletebuffer"
	"github.com/milvus-io/milvus/internal/querynodev2/segments"
	"github.com/milvus-io/milvus/internal/storage"
	"github.com/milvus-io/milvus/internal/util/function"
	"github.com/milvus-io/milvus/internal/util/streamrpc"
	"github.com/milvus-io/milvus/pkg/v3/common"
	"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
	"github.com/milvus-io/milvus/pkg/v3/proto/segcorepb"
	"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
	"github.com/milvus-io/milvus/pkg/v3/util/lifetime"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
	"github.com/milvus-io/milvus/pkg/v3/util/metric"
	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
	"github.com/milvus-io/milvus/pkg/v3/util/tsoutil"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

func TestMain(m *testing.M) {
	streaming.SetupNoopWALForTest()

	os.Exit(m.Run())
}

type DelegatorSuite struct {
	suite.Suite

	collectionID  int64
	partitionIDs  []int64
	replicaID     int64
	vchannelName  string
	version       int64
	workerManager *cluster.MockManager
	manager       *segments.Manager
	loader        *segments.MockLoader
	mq            *msgstream.MockMsgStream

	delegator    ShardDelegator
	chunkManager storage.ChunkManager
	rootPath     string
}

func (s *DelegatorSuite) SetupSuite() {
	paramtable.Init()
}

func (s *DelegatorSuite) TearDownSuite() {
}

func (s *DelegatorSuite) SetupTest() {
	paramtable.Get().Save(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key, "false")
	s.T().Cleanup(func() {
		paramtable.Get().Reset(paramtable.Get().CommonCfg.EnableGrowingSourceFlush.Key)
	})

	s.collectionID = 1000
	s.partitionIDs = []int64{500, 501}
	s.replicaID = 65535
	s.vchannelName = "rootcoord-dml_1000_v0"
	s.version = 2000
	s.workerManager = &cluster.MockManager{}
	s.manager = segments.NewManager()
	s.loader = &segments.MockLoader{}
	s.loader.EXPECT().
		Load(mock.Anything, s.collectionID, segments.SegmentTypeGrowing, int64(0), mock.Anything).
		Call.Return(func(ctx context.Context, collectionID int64, segmentType segments.SegmentType, version int64, infos ...*querypb.SegmentLoadInfo) []segments.Segment {
		return lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) segments.Segment {
			ms := &segments.MockSegment{}
			ms.EXPECT().ID().Return(info.GetSegmentID())
			ms.EXPECT().Type().Return(segments.SegmentTypeGrowing)
			ms.EXPECT().Partition().Return(info.GetPartitionID())
			ms.EXPECT().Collection().Return(info.GetCollectionID())
			ms.EXPECT().Indexes().Return(nil)
			ms.EXPECT().RowNum().Return(info.GetNumOfRows())
			ms.EXPECT().LoadInfo().Return(info)
			ms.EXPECT().Delete(mock.Anything, mock.Anything, mock.Anything).Return(nil)
			return ms
		})
	}, nil)

	// init schema
	s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
		Name: "TestCollection",
		Fields: []*schemapb.FieldSchema{
			{
				Name:         "id",
				FieldID:      100,
				IsPrimaryKey: true,
				DataType:     schemapb.DataType_Int64,
				AutoID:       true,
			},
			{
				Name:         "vector",
				FieldID:      101,
				IsPrimaryKey: false,
				DataType:     schemapb.DataType_BinaryVector,
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.DimKey,
						Value: "128",
					},
				},
			},
		},
	}, &segcorepb.CollectionIndexMeta{
		MaxIndexRowCount: 100,
		IndexMetas: []*segcorepb.FieldIndexMeta{
			{
				FieldID:      101,
				CollectionID: s.collectionID,
				IndexName:    "binary_index",
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.DimKey,
						Value: "128",
					},
				},
				IndexParams: []*commonpb.KeyValuePair{
					{
						Key:   common.IndexTypeKey,
						Value: "BIN_IVF_FLAT",
					},
					{
						Key:   common.MetricTypeKey,
						Value: metric.JACCARD,
					},
				},
			},
		},
	}, &querypb.LoadMetaInfo{
		PartitionIDs:    s.partitionIDs,
		SchemaBarrierTs: tsoutil.ComposeTSByTime(time.Now()),
	})

	s.mq = &msgstream.MockMsgStream{}
	s.rootPath = "delegator_test"

	// init chunkManager
	chunkManagerFactory := storage.NewTestChunkManagerFactory(paramtable.Get(), s.rootPath)
	s.chunkManager, _ = chunkManagerFactory.NewPersistentStorageChunkManager(context.Background())

	var err error
	//	s.delegator, err = NewShardDelegator(s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, s.manager, s.tsafeManager, s.loader)
	s.delegator, err = NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, s.manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
	s.Require().NoError(err)
}

func (s *DelegatorSuite) TearDownTest() {
	if s.delegator != nil {
		s.delegator.Close()
	}
	function.ReleaseFunctionRunners(s.collectionID, "WAL-"+s.vchannelName)
	s.delegator = nil
}

func (s *DelegatorSuite) TestCreateDelegatorWithFunction() {
	s.Run("init function failed", func() {
		manager := segments.NewManager()
		manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
			Name: "TestCollection",
			Fields: []*schemapb.FieldSchema{
				{
					Name:         "id",
					FieldID:      100,
					IsPrimaryKey: true,
					DataType:     schemapb.DataType_Int64,
					AutoID:       true,
				}, {
					Name:         "vector",
					FieldID:      101,
					IsPrimaryKey: false,
					DataType:     schemapb.DataType_SparseFloatVector,
				},
			},
			Functions: []*schemapb.FunctionSchema{{
				Type:           schemapb.FunctionType_BM25,
				InputFieldIds:  []int64{102},
				OutputFieldIds: []int64{101, 103}, // invalid output field
			}},
		}, nil, &querypb.LoadMetaInfo{SchemaBarrierTs: tsoutil.ComposeTSByTime(time.Now())})

		_, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
		s.Error(err)
	})

	s.Run("init function failed", func() {
		manager := segments.NewManager()
		manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
			Name: "TestCollection",
			Fields: []*schemapb.FieldSchema{
				{
					Name:         "id",
					FieldID:      100,
					IsPrimaryKey: true,
					DataType:     schemapb.DataType_Int64,
					AutoID:       true,
				}, {
					Name:         "vector",
					FieldID:      101,
					IsPrimaryKey: false,
					DataType:     schemapb.DataType_SparseFloatVector,
				}, {
					Name:     "text",
					FieldID:  102,
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{
							Key:   common.MaxLengthKey,
							Value: "256",
						},
					},
				},
			},
			Functions: []*schemapb.FunctionSchema{{
				Type:           schemapb.FunctionType_BM25,
				InputFieldIds:  []int64{102},
				OutputFieldIds: []int64{101},
			}},
		}, nil, &querypb.LoadMetaInfo{SchemaBarrierTs: tsoutil.ComposeTSByTime(time.Now())})

		delegator, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
		s.NoError(err)
		defer delegator.Close()
	})
}

func (s *DelegatorSuite) TestBasicInfo() {
	s.Equal(s.collectionID, s.delegator.Collection())
	s.Equal(s.version, s.delegator.Version())
	s.False(s.delegator.Serviceable())
	s.delegator.Start()
	s.True(s.delegator.Serviceable())
}

// TestDelegatorStateChecks tests the state checking methods added/modified in the delegator
func (s *DelegatorSuite) TestDelegatorStateChecks() {
	sd := s.delegator.(*shardDelegator)

	s.Run("test_state_methods_with_different_states", func() {
		// Test Initializing state
		sd.lifetime.SetState(lifetime.Initializing)

		// NotStopped should return nil for non-stopped states
		err := sd.NotStopped(sd.lifetime.GetState())
		s.NoError(err)

		// IsWorking should return error for non-working states
		err = sd.IsWorking(sd.lifetime.GetState())
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")
		s.Contains(err.Error(), "Initializing")

		// Serviceable should return false for non-working states
		s.False(sd.Serviceable())

		// Stopped should return false for non-stopped states
		s.False(sd.Stopped())

		// Test Working state
		sd.lifetime.SetState(lifetime.Working)

		// NotStopped should return nil for non-stopped states
		err = sd.NotStopped(sd.lifetime.GetState())
		s.NoError(err)

		// IsWorking should return nil for working state
		err = sd.IsWorking(sd.lifetime.GetState())
		s.NoError(err)

		// Serviceable should return true for working state
		s.True(sd.Serviceable())

		// Stopped should return false for non-stopped states
		s.False(sd.Stopped())

		// Test Stopped state
		sd.lifetime.SetState(lifetime.Stopped)

		// NotStopped should return error for stopped state
		err = sd.NotStopped(sd.lifetime.GetState())
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")
		s.Contains(err.Error(), "Stopped")

		// IsWorking should return error for stopped state
		err = sd.IsWorking(sd.lifetime.GetState())
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")
		s.Contains(err.Error(), "Stopped")

		// Serviceable should return false for stopped state
		s.False(sd.Serviceable())

		// Stopped should return true for stopped state
		s.True(sd.Stopped())
	})

	s.Run("test_state_methods_with_direct_state_parameter", func() {
		// Test NotStopped with different states
		err := sd.NotStopped(lifetime.Initializing)
		s.NoError(err)

		err = sd.NotStopped(lifetime.Working)
		s.NoError(err)

		err = sd.NotStopped(lifetime.Stopped)
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")
		s.Contains(err.Error(), sd.vchannelName)

		// Test IsWorking with different states
		err = sd.IsWorking(lifetime.Initializing)
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")
		s.Contains(err.Error(), sd.vchannelName)

		err = sd.IsWorking(lifetime.Working)
		s.NoError(err)

		err = sd.IsWorking(lifetime.Stopped)
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")
		s.Contains(err.Error(), sd.vchannelName)
	})

	s.Run("test_error_messages_contain_channel_name", func() {
		// Verify error messages contain the channel name for better debugging
		err := sd.NotStopped(lifetime.Stopped)
		s.Error(err)
		s.Contains(err.Error(), sd.vchannelName)

		err = sd.IsWorking(lifetime.Initializing)
		s.Error(err)
		s.Contains(err.Error(), sd.vchannelName)

		err = sd.IsWorking(lifetime.Stopped)
		s.Error(err)
		s.Contains(err.Error(), sd.vchannelName)
	})
}

func (s *DelegatorSuite) TestGetSegmentInfo() {
	sealed, growing := s.delegator.GetSegmentInfo(false)
	s.Equal(0, len(sealed))
	s.Equal(0, len(growing))

	s.delegator.SyncDistribution(context.Background(), SegmentEntry{
		NodeID:      1,
		SegmentID:   1001,
		PartitionID: 500,
		Version:     2001,
	})

	s.delegator.(*shardDelegator).distribution.Flush()
	sealed, growing = s.delegator.GetSegmentInfo(false)
	s.EqualValues([]SnapshotItem{
		{
			NodeID: 1,
			Segments: []SegmentEntry{
				{
					NodeID:        1,
					SegmentID:     1001,
					PartitionID:   500,
					Version:       2001,
					TargetVersion: unreadableTargetVersion,
				},
			},
		},
	}, sealed)
	s.Equal(0, len(growing))
}

// nodeID 1 => sealed segment 1000, 1001
// nodeID 1 => growing segment 1004
// nodeID 2 => sealed segment 1002, 1003
func (s *DelegatorSuite) initSegments() {
	s.delegator.LoadGrowing(context.Background(), []*querypb.SegmentLoadInfo{
		{
			SegmentID:    1004,
			CollectionID: s.collectionID,
			PartitionID:  500,
		},
	}, 0)
	s.delegator.SyncDistribution(context.Background(),
		SegmentEntry{
			NodeID:      1,
			SegmentID:   1000,
			PartitionID: 500,
			Version:     2001,
		},
		SegmentEntry{
			NodeID:      1,
			SegmentID:   1001,
			PartitionID: 501,
			Version:     2001,
		},
		SegmentEntry{
			NodeID:      2,
			SegmentID:   1002,
			PartitionID: 500,
			Version:     2001,
		},
		SegmentEntry{
			NodeID:      2,
			SegmentID:   1003,
			PartitionID: 501,
			Version:     2001,
		},
	)
	s.delegator.SyncTargetVersion(&querypb.SyncAction{
		TargetVersion:   2001,
		GrowingInTarget: []int64{1004},
		SealedSegmentRowCount: map[int64]int64{
			1000: 100,
			1001: 100,
			1002: 100,
			1003: 100,
		},
		DroppedInTarget: []int64{},
		Checkpoint:      &msgpb.MsgPosition{},
		DeleteCP:        &msgpb.MsgPosition{},
	}, []int64{500, 501})
}

func (s *DelegatorSuite) TestSearch() {
	s.delegator.Start()
	paramtable.SetNodeID(1)
	s.initSegments()
	s.Run("normal", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.EqualValues(1, req.Req.GetBase().GetTargetID())

				if req.GetScope() == querypb.DataScope_Streaming {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1004}, req.GetSegmentIDs())
				}
				if req.GetScope() == querypb.DataScope_Historical {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1000, 1001}, req.GetSegmentIDs())
				}
			}).Return(&internalpb.SearchResults{}, nil)
		worker2.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.SearchResults{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		results, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req:         &internalpb.SearchRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.NoError(err)
		s.Equal(3, len(results))
	})

	s.Run("partition_not_loaded", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Base: commonpbutil.NewMsgBase(),
				// not load partation -1,will return error
				PartitionIDs: []int64{-1},
			},
			DmlChannels: []string{s.vchannelName},
		})

		s.True(errors.Is(err, merr.ErrPartitionNotLoaded))
	})

	s.Run("worker_return_error", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).Return(nil, errors.New("mock error"))
		worker2.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.SearchResults{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req:         &internalpb.SearchRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("worker_return_failure_code", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).Return(&internalpb.SearchResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    "mocked error",
			},
		}, nil)
		worker2.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.SearchResults{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req:         &internalpb.SearchRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("wrong_channel", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		_, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req:         &internalpb.SearchRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{"non_exist_channel"},
		})

		s.Error(err)
	})

	s.Run("wait_tsafe_timeout", func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*100)
		defer cancel()

		_, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Base:               commonpbutil.NewMsgBase(),
				GuaranteeTimestamp: 10100,
			},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("tsafe_behind_max_lag", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		_, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Base:               commonpbutil.NewMsgBase(),
				GuaranteeTimestamp: uint64(paramtable.Get().QueryNodeCfg.MaxTimestampLag.GetAsDuration(time.Second)) + 10001,
			},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
		s.ErrorIs(err, merr.ErrChannelTSafeStalled)
		s.True(merr.IsRetryableErr(err))

		status := merr.Status(err)
		s.True(status.GetRetriable())
		s.Equal(commonpb.ErrorCode_TimeTickLongDelay, status.GetErrorCode())

		roundTripErr := merr.Error(status)
		s.ErrorIs(roundTripErr, merr.ErrChannelTSafeStalled)
		s.True(merr.IsRetryableErr(roundTripErr))
	})

	s.Run("downgrade_tsafe", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()
		pt := paramtable.Get()
		pt.Save(pt.QueryNodeCfg.DowngradeTsafe.Key, "true")
		defer pt.Reset(pt.QueryNodeCfg.DowngradeTsafe.Key)

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.EqualValues(1, req.Req.GetBase().GetTargetID())

				if req.GetScope() == querypb.DataScope_Streaming {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1004}, req.GetSegmentIDs())
				}
				if req.GetScope() == querypb.DataScope_Historical {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1000, 1001}, req.GetSegmentIDs())
				}
			}).Return(&internalpb.SearchResults{}, nil)
		worker2.EXPECT().SearchSegments(mock.Anything, mock.AnythingOfType("*querypb.SearchRequest")).
			Run(func(_ context.Context, req *querypb.SearchRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.SearchResults{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		results, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Base:               commonpbutil.NewMsgBase(),
				GuaranteeTimestamp: uint64(paramtable.Get().QueryNodeCfg.MaxTimestampLag.GetAsDuration(time.Second)) + 10001,
			},
			DmlChannels: []string{s.vchannelName},
		})

		s.NoError(err)
		s.Equal(3, len(results))
	})

	s.Run("distribution_not_serviceable", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		sd, ok := s.delegator.(*shardDelegator)
		s.Require().True(ok)
		sd.distribution.MarkOfflineSegments(1001)
		sd.distribution.Flush()

		_, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req: &internalpb.SearchRequest{
				Base: commonpbutil.NewMsgBase(),
			},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("cluster_not_serviceable", func() {
		s.delegator.Close()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		_, err := s.delegator.Search(ctx, &querypb.SearchRequest{
			Req:         &internalpb.SearchRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})
}

func (s *DelegatorSuite) TestQuery() {
	s.delegator.Start()
	paramtable.SetNodeID(1)
	s.initSegments()
	s.Run("normal", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().QuerySegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest")).
			Run(func(_ context.Context, req *querypb.QueryRequest) {
				s.EqualValues(1, req.Req.GetBase().GetTargetID())

				if req.GetScope() == querypb.DataScope_Streaming {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1004}, req.GetSegmentIDs())
				}
				if req.GetScope() == querypb.DataScope_Historical {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1000, 1001}, req.GetSegmentIDs())
				}
			}).Return(&internalpb.RetrieveResults{}, nil)
		worker2.EXPECT().QuerySegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest")).
			Run(func(_ context.Context, req *querypb.QueryRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.RetrieveResults{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		results, err := s.delegator.Query(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.NoError(err)
		s.Equal(3, len(results))
	})

	s.Run("partition_not_loaded", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_, err := s.delegator.Query(ctx, &querypb.QueryRequest{
			Req: &internalpb.RetrieveRequest{
				QueryLabel: "query",
				Base:       commonpbutil.NewMsgBase(),
				// not load partation -1,will return error
				PartitionIDs: []int64{-1},
			},
			DmlChannels: []string{s.vchannelName},
		})
		s.True(errors.Is(err, merr.ErrPartitionNotLoaded))
	})

	s.Run("worker_return_error", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().QuerySegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest")).Return(nil, errors.New("mock error"))
		worker2.EXPECT().QuerySegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest")).
			Run(func(_ context.Context, req *querypb.QueryRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.RetrieveResults{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_, err := s.delegator.Query(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("worker_return_failure_code", func() {
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().QuerySegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest")).Return(&internalpb.RetrieveResults{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    "mocked error",
			},
		}, nil)
		worker2.EXPECT().QuerySegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest")).
			Run(func(_ context.Context, req *querypb.QueryRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.RetrieveResults{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_, err := s.delegator.Query(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("wrong_channel", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		_, err := s.delegator.Query(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{"non_exist_channel"},
		})

		s.Error(err)
	})

	s.Run("distribution_not_serviceable", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		sd, ok := s.delegator.(*shardDelegator)
		s.Require().True(ok)
		sd.distribution.MarkOfflineSegments(1001)
		sd.distribution.Flush()

		_, err := s.delegator.Query(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("cluster_not_serviceable", func() {
		s.delegator.Close()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		_, err := s.delegator.Query(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})
}

func (s *DelegatorSuite) TestQueryStream() {
	s.delegator.Start()
	paramtable.SetNodeID(1)
	s.initSegments()

	s.Run("normal", func() {
		defer func() {
			s.workerManager.AssertExpectations(s.T())
			s.workerManager.ExpectedCalls = nil
		}()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		client := streamrpc.NewLocalQueryClient(ctx)
		server := client.CreateServer()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2
		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		worker1.EXPECT().QueryStreamSegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest"), mock.Anything).
			Run(func(ctx context.Context, req *querypb.QueryRequest, srv streamrpc.QueryStreamServer) {
				s.EqualValues(1, req.Req.GetBase().GetTargetID())

				if req.GetScope() == querypb.DataScope_Streaming {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1004}, req.GetSegmentIDs())
				}
				if req.GetScope() == querypb.DataScope_Historical {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1000, 1001}, req.GetSegmentIDs())
				}

				srv.Send(&internalpb.RetrieveResults{
					Status: merr.Success(),
					Ids: &schemapb.IDs{
						IdField: &schemapb.IDs_IntId{
							IntId: &schemapb.LongArray{Data: req.GetSegmentIDs()},
						},
					},
				})
			}).Return(nil)

		worker2.EXPECT().QueryStreamSegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest"), mock.Anything).
			Run(func(ctx context.Context, req *querypb.QueryRequest, srv streamrpc.QueryStreamServer) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
				srv.Send(&internalpb.RetrieveResults{
					Status: merr.Success(),
					Ids: &schemapb.IDs{
						IdField: &schemapb.IDs_IntId{
							IntId: &schemapb.LongArray{Data: req.GetSegmentIDs()},
						},
					},
				})
			}).Return(nil)

		// run stream function
		go func() {
			err := s.delegator.QueryStream(ctx, &querypb.QueryRequest{
				Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
				DmlChannels: []string{s.vchannelName},
			}, server)
			s.NoError(err)
			server.FinishSend(err)
		}()

		resultIDs := []int64{1000, 1001, 1002, 1003, 1004}
		recNum := 0
		for {
			result, err := client.Recv()
			if err == io.EOF {
				s.Equal(recNum, len(resultIDs))
				break
			}
			s.NoError(err)

			err = merr.Error(result.GetStatus())
			s.NoError(err)

			for _, segmentID := range result.Ids.GetIntId().Data {
				s.Less(recNum, len(resultIDs))
				lo.Contains[int64](resultIDs, segmentID)
				recNum++
			}
		}
	})

	s.Run("partition_not_loaded", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		client := streamrpc.NewLocalQueryClient(ctx)
		server := client.CreateServer()

		err := s.delegator.QueryStream(ctx, &querypb.QueryRequest{
			Req: &internalpb.RetrieveRequest{
				QueryLabel: "query",
				Base:       commonpbutil.NewMsgBase(),
				// not load partation -1,will return error
				PartitionIDs: []int64{-1},
			},
			DmlChannels: []string{s.vchannelName},
		}, server)
		s.True(errors.Is(err, merr.ErrPartitionNotLoaded))
	})

	s.Run("tsafe_behind_max_lag", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		client := streamrpc.NewLocalQueryClient(ctx)
		server := client.CreateServer()

		err := s.delegator.QueryStream(ctx, &querypb.QueryRequest{
			Req: &internalpb.RetrieveRequest{
				QueryLabel:         "query",
				Base:               commonpbutil.NewMsgBase(),
				GuaranteeTimestamp: uint64(paramtable.Get().QueryNodeCfg.MaxTimestampLag.GetAsDuration(time.Second)),
			},
			DmlChannels: []string{s.vchannelName},
		}, server)
		s.Error(err)
		s.ErrorIs(err, merr.ErrChannelTSafeStalled)
		s.True(merr.IsRetryableErr(err))
	})

	s.Run("get_worker_failed", func() {
		defer func() {
			s.workerManager.AssertExpectations(s.T())
			s.workerManager.ExpectedCalls = nil
		}()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		client := streamrpc.NewLocalQueryClient(ctx)
		server := client.CreateServer()

		mockErr := errors.New("mock error")

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(nil, mockErr)

		// run stream function
		err := s.delegator.QueryStream(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		}, server)
		s.Error(err)
		s.ErrorContains(err, "segments not loaded in any worker")
	})

	s.Run("worker_return_error", func() {
		defer func() {
			s.workerManager.AssertExpectations(s.T())
			s.workerManager.ExpectedCalls = nil
		}()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		client := streamrpc.NewLocalQueryClient(ctx)
		server := client.CreateServer()

		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}
		mockErr := errors.New("mock error")

		workers[1] = worker1
		workers[2] = worker2
		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		worker1.EXPECT().QueryStreamSegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest"), mock.Anything).
			Return(mockErr)

		worker2.EXPECT().QueryStreamSegments(mock.Anything, mock.AnythingOfType("*querypb.QueryRequest"), mock.Anything).
			Run(func(ctx context.Context, req *querypb.QueryRequest, srv streamrpc.QueryStreamServer) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())

				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
				srv.Send(&internalpb.RetrieveResults{
					Status: merr.Success(),
					Ids: &schemapb.IDs{
						IdField: &schemapb.IDs_IntId{
							IntId: &schemapb.LongArray{Data: req.GetSegmentIDs()},
						},
					},
				})
			}).Return(nil)

		// run stream function
		go func() {
			err := s.delegator.QueryStream(ctx, &querypb.QueryRequest{
				Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
				DmlChannels: []string{s.vchannelName},
			}, server)
			server.Send(&internalpb.RetrieveResults{
				Status: merr.Status(err),
			})
		}()

		resultIDs := []int64{1002, 1003}
		recNum := 0
		for {
			result, err := client.Recv()
			s.NoError(err)

			err = merr.Error(result.GetStatus())
			if err != nil {
				s.Equal(recNum, len(resultIDs))
				s.Equal(err.Error(), mockErr.Error())
				break
			}

			for _, segmentID := range result.Ids.GetIntId().GetData() {
				s.Less(recNum, len(resultIDs))
				lo.Contains[int64](resultIDs, segmentID)
				recNum++
			}
		}
	})

	s.Run("wrong_channel", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		client := streamrpc.NewLocalQueryClient(ctx)
		server := client.CreateServer()

		err := s.delegator.QueryStream(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{"non_exist_channel"},
		}, server)

		s.Error(err)
	})

	s.Run("distribution_not_serviceable", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		sd, ok := s.delegator.(*shardDelegator)
		s.Require().True(ok)
		sd.distribution.MarkOfflineSegments(1001)
		sd.distribution.Flush()

		client := streamrpc.NewLocalQueryClient(ctx)
		server := client.CreateServer()

		// run stream function
		err := s.delegator.QueryStream(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		}, server)
		s.Error(err)
	})

	s.Run("cluster_not_serviceable", func() {
		s.delegator.Close()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		client := streamrpc.NewLocalQueryClient(ctx)
		server := client.CreateServer()

		err := s.delegator.QueryStream(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		}, server)

		s.Error(err)
	})
}

func (s *DelegatorSuite) TestGetStats() {
	s.delegator.Start()
	// 1 => sealed segment 1000, 1001
	// 1 => growing segment 1004
	// 2 => sealed segment 1002, 1003
	paramtable.SetNodeID(1)
	s.initSegments()

	s.Run("normal", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().GetStatistics(mock.Anything, mock.AnythingOfType("*querypb.GetStatisticsRequest")).
			Run(func(_ context.Context, req *querypb.GetStatisticsRequest) {
				s.EqualValues(1, req.Req.GetBase().GetTargetID())
				s.True(req.GetFromShardLeader())
				if req.GetScope() == querypb.DataScope_Streaming {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1004}, req.GetSegmentIDs())
				}
				if req.GetScope() == querypb.DataScope_Historical {
					s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
					s.ElementsMatch([]int64{1000, 1001}, req.GetSegmentIDs())
				}
			}).Return(&internalpb.GetStatisticsResponse{}, nil)
		worker2.EXPECT().GetStatistics(mock.Anything, mock.AnythingOfType("*querypb.GetStatisticsRequest")).
			Run(func(_ context.Context, req *querypb.GetStatisticsRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())
				s.True(req.GetFromShardLeader())
				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.GetStatisticsResponse{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		results, err := s.delegator.GetStatistics(ctx, &querypb.GetStatisticsRequest{
			Req:         &internalpb.GetStatisticsRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.NoError(err)
		s.Equal(3, len(results))
	})

	s.Run("worker_return_error", func() {
		defer func() {
			s.workerManager.ExpectedCalls = nil
		}()
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().GetStatistics(mock.Anything, mock.AnythingOfType("*querypb.GetStatisticsRequest")).Return(nil, errors.New("mock error"))
		worker2.EXPECT().GetStatistics(mock.Anything, mock.AnythingOfType("*querypb.GetStatisticsRequest")).
			Run(func(_ context.Context, req *querypb.GetStatisticsRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())
				s.True(req.GetFromShardLeader())
				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.GetStatisticsResponse{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_, err := s.delegator.GetStatistics(ctx, &querypb.GetStatisticsRequest{
			Req:         &internalpb.GetStatisticsRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("worker_return_failure_code", func() {
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := &cluster.MockWorker{}
		worker2 := &cluster.MockWorker{}

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().GetStatistics(mock.Anything, mock.AnythingOfType("*querypb.GetStatisticsRequest")).Return(&internalpb.GetStatisticsResponse{
			Status: &commonpb.Status{
				ErrorCode: commonpb.ErrorCode_UnexpectedError,
				Reason:    "mocked error",
			},
		}, nil)
		worker2.EXPECT().GetStatistics(mock.Anything, mock.AnythingOfType("*querypb.GetStatisticsRequest")).
			Run(func(_ context.Context, req *querypb.GetStatisticsRequest) {
				s.EqualValues(2, req.Req.GetBase().GetTargetID())
				s.True(req.GetFromShardLeader())
				s.Equal(querypb.DataScope_Historical, req.GetScope())
				s.EqualValues([]string{s.vchannelName}, req.GetDmlChannels())
				s.ElementsMatch([]int64{1002, 1003}, req.GetSegmentIDs())
			}).Return(&internalpb.GetStatisticsResponse{}, nil)

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		_, err := s.delegator.GetStatistics(ctx, &querypb.GetStatisticsRequest{
			Req:         &internalpb.GetStatisticsRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("wrong_channel", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		_, err := s.delegator.GetStatistics(ctx, &querypb.GetStatisticsRequest{
			Req:         &internalpb.GetStatisticsRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{"non_exist_channel"},
		})

		s.Error(err)
	})

	s.Run("distribution_not_serviceable", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		sd, ok := s.delegator.(*shardDelegator)
		s.Require().True(ok)
		sd.distribution.MarkOfflineSegments(1001)
		sd.distribution.Flush()

		_, err := s.delegator.GetStatistics(ctx, &querypb.GetStatisticsRequest{
			Req:         &internalpb.GetStatisticsRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})

	s.Run("cluster_not_serviceable", func() {
		s.delegator.Close()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		_, err := s.delegator.GetStatistics(ctx, &querypb.GetStatisticsRequest{
			Req:         &internalpb.GetStatisticsRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})

		s.Error(err)
	})
}

func (s *DelegatorSuite) TestUpdateSchema() {
	s.delegator.Start()
	paramtable.SetNodeID(1)
	s.initSegments()

	s.Run("normal", func() {
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := cluster.NewMockWorker(s.T())
		worker2 := cluster.NewMockWorker(s.T())

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).RunAndReturn(func(ctx context.Context, usr *querypb.UpdateSchemaRequest) (*commonpb.Status, error) {
			s.Equal(int32(10), usr.GetSchema().GetVersion())
			s.Equal(uint64(100), usr.GetSchemaBarrierTs())
			return merr.Success(), nil
		}).Twice()

		worker2.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).RunAndReturn(func(ctx context.Context, usr *querypb.UpdateSchemaRequest) (*commonpb.Status, error) {
			s.Equal(int32(10), usr.GetSchema().GetVersion())
			s.Equal(uint64(100), usr.GetSchemaBarrierTs())
			return merr.Success(), nil
		}).Once()

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil).Times(3) // currently node 1 will be called twice for growing & sealed

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		schema := newFunctionRuntimeTestSchemaWithVersion(10)
		err := s.delegator.UpdateSchema(ctx, schema, 100)
		s.NoError(err)
	})

	s.Run("worker_return_error", func() {
		workers := make(map[int64]*cluster.MockWorker)
		worker1 := cluster.NewMockWorker(s.T())
		worker2 := cluster.NewMockWorker(s.T())

		workers[1] = worker1
		workers[2] = worker2

		worker1.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).RunAndReturn(func(ctx context.Context, usr *querypb.UpdateSchemaRequest) (*commonpb.Status, error) {
			return merr.Status(merr.WrapErrServiceInternal("mocked")), merr.WrapErrServiceInternal("mocked")
		}).Maybe()

		worker2.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).RunAndReturn(func(ctx context.Context, usr *querypb.UpdateSchemaRequest) (*commonpb.Status, error) {
			return merr.Success(), nil
		}).Maybe()

		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).Call.Return(func(_ context.Context, nodeID int64) cluster.Worker {
			return workers[nodeID]
		}, nil).Times(3)

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		err := s.delegator.UpdateSchema(ctx, newFunctionRuntimeTestSchemaWithVersion(s.nextSchemaVersion()), 100)
		s.Error(err)
	})

	s.Run("worker_manager_error", func() {
		s.workerManager.EXPECT().GetWorker(mock.Anything, mock.AnythingOfType("int64")).RunAndReturn(func(ctx context.Context, i int64) (cluster.Worker, error) {
			return nil, merr.WrapErrServiceInternal("mocked")
		})
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		err := s.delegator.UpdateSchema(ctx, newFunctionRuntimeTestSchemaWithVersion(s.nextSchemaVersion()), 100)
		s.Error(err)
	})

	s.Run("distribution_not_serviceable", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		sd, ok := s.delegator.(*shardDelegator)
		s.Require().True(ok)
		sd.distribution.MarkOfflineSegments(1001)
		sd.distribution.Flush()

		err := s.delegator.UpdateSchema(ctx, newFunctionRuntimeTestSchemaWithVersion(s.nextSchemaVersion()), 100)
		s.Error(err)
	})

	s.Run("cluster_not_serviceable", func() {
		s.delegator.Close()

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		err := s.delegator.UpdateSchema(ctx, &schemapb.CollectionSchema{}, 100)
		s.Error(err)
	})
}

func (s *DelegatorSuite) ResetDelegator() {
	var err error
	s.delegator.Close()
	s.delegator, err = NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, s.manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
	s.Require().NoError(err)
	s.allocFunctionRunnersForTest()
}

func (s *DelegatorSuite) allocFunctionRunnersForTest() {
	sd := s.delegator.(*shardDelegator)
	s.Require().NoError(function.UpdateFunctionRunners(s.collectionID, delegatorFunctionRunnerKey(s.vchannelName), sd.collection.Schema()))
}

func (s *DelegatorSuite) nextSchemaBarrierLoadMeta() *querypb.LoadMetaInfo {
	return &querypb.LoadMetaInfo{SchemaBarrierTs: uint64(s.nextSchemaVersion())}
}

func (s *DelegatorSuite) nextSchemaVersion() int32 {
	version := uint64(1)
	if collection := s.manager.Collection.Get(s.collectionID); collection != nil {
		version = collection.SchemaVersion() + 1
	}
	return int32(version)
}

func (s *DelegatorSuite) TestRunAnalyzer() {
	ctx := context.Background()
	s.TestCreateDelegatorWithFunction()
	s.Run("field analyzer not exist", func() {
		_, err := s.delegator.RunAnalyzer(ctx, &querypb.RunAnalyzerRequest{
			FieldId: 100,
		})
		s.Require().Error(err)
	})

	s.Run("normal analyer", func() {
		err := s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
			Version: s.nextSchemaVersion(),
			Fields: []*schemapb.FieldSchema{
				{
					FieldID:      103,
					Name:         "id",
					DataType:     schemapb.DataType_Int64,
					IsPrimaryKey: true,
					AutoID:       true,
				},
				{
					FieldID:  100,
					Name:     "text",
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{Key: common.MaxLengthKey, Value: "256"},
						{Key: "analyzer_params", Value: "{}"},
					},
				},
				{
					FieldID:  101,
					Name:     "sparse",
					DataType: schemapb.DataType_SparseFloatVector,
				},
			},
			Functions: []*schemapb.FunctionSchema{{
				Type:             schemapb.FunctionType_BM25,
				InputFieldNames:  []string{"text"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"sparse"},
				OutputFieldIds:   []int64{101},
			}},
		}, nil, s.nextSchemaBarrierLoadMeta())
		s.Require().NoError(err)
		s.ResetDelegator()

		result, err := s.delegator.RunAnalyzer(ctx, &querypb.RunAnalyzerRequest{
			FieldId:     100,
			Placeholder: [][]byte{[]byte("test doc")},
		})
		s.Require().NoError(err)
		s.Equal(2, len(result[0].GetTokens()))
	})

	s.Run("standalone field analyzer", func() {
		err := s.manager.Collection.PutOrRef(s.collectionID, newFunctionRuntimeTestSchemaWithVersion(s.nextSchemaVersion()), nil, s.nextSchemaBarrierLoadMeta())
		s.Require().NoError(err)
		s.ResetDelegator()

		result, err := s.delegator.RunAnalyzer(ctx, &querypb.RunAnalyzerRequest{
			FieldId:     103,
			Placeholder: [][]byte{[]byte("test doc")},
		})
		s.Require().NoError(err)
		s.Equal(2, len(result[0].GetTokens()))
	})

	s.Run("multi analyzer", func() {
		err := s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
			Version: s.nextSchemaVersion(),
			Fields: []*schemapb.FieldSchema{
				{
					FieldID:      103,
					Name:         "id",
					DataType:     schemapb.DataType_Int64,
					IsPrimaryKey: true,
					AutoID:       true,
				},
				{
					FieldID:  100,
					Name:     "text",
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{Key: common.MaxLengthKey, Value: "256"},
						{Key: "multi_analyzer_params", Value: `{
						"by_field": "analyzer",
    					"analyzers": {
							"default": {}
						}
					}`},
					},
				},
				{
					FieldID:  101,
					Name:     "sparse",
					DataType: schemapb.DataType_SparseFloatVector,
				},
				{
					FieldID:    102,
					Name:       "analyzer",
					DataType:   schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxLengthKey, Value: "256"}},
				},
			},
			Functions: []*schemapb.FunctionSchema{{
				Type:             schemapb.FunctionType_BM25,
				InputFieldNames:  []string{"text"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"sparse"},
				OutputFieldIds:   []int64{101},
			}},
		}, nil, s.nextSchemaBarrierLoadMeta())
		s.Require().NoError(err)
		s.ResetDelegator()

		result, err := s.delegator.RunAnalyzer(ctx, &querypb.RunAnalyzerRequest{
			FieldId:       100,
			Placeholder:   [][]byte{[]byte("test doc"), []byte("test doc2")},
			AnalyzerNames: []string{"default"},
		})
		s.Require().NoError(err)
		s.Equal(2, len(result[0].GetTokens()))
		s.Equal(2, len(result[1].GetTokens()))
	})

	s.Run("multi analyzer defaults analyzer name", func() {
		err := s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
			Version: s.nextSchemaVersion(),
			Fields: []*schemapb.FieldSchema{
				{
					FieldID:      103,
					Name:         "id",
					DataType:     schemapb.DataType_Int64,
					IsPrimaryKey: true,
					AutoID:       true,
				},
				{
					FieldID:  100,
					Name:     "text",
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{Key: common.MaxLengthKey, Value: "256"},
						{Key: "multi_analyzer_params", Value: `{
						"by_field": "analyzer",
    					"analyzers": {
							"default": {}
						}
					}`},
					},
				},
				{
					FieldID:  101,
					Name:     "sparse",
					DataType: schemapb.DataType_SparseFloatVector,
				},
				{
					FieldID:    102,
					Name:       "analyzer",
					DataType:   schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxLengthKey, Value: "256"}},
				},
			},
			Functions: []*schemapb.FunctionSchema{{
				Type:             schemapb.FunctionType_BM25,
				InputFieldNames:  []string{"text"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"sparse"},
				OutputFieldIds:   []int64{101},
			}},
		}, nil, s.nextSchemaBarrierLoadMeta())
		s.Require().NoError(err)
		s.ResetDelegator()

		_, err = s.delegator.RunAnalyzer(ctx, &querypb.RunAnalyzerRequest{
			FieldId:     100,
			Placeholder: [][]byte{[]byte("test doc")},
		})
		s.Require().NoError(err)

		_, err = s.delegator.RunAnalyzer(ctx, &querypb.RunAnalyzerRequest{
			FieldId:       100,
			Placeholder:   [][]byte{[]byte("test doc")},
			AnalyzerNames: []string{"default", "standard"},
		})
		s.Require().Error(err)
	})
}

func (s *DelegatorSuite) TestGetHighlight() {
	ctx := context.Background()
	s.TestCreateDelegatorWithFunction()

	s.Run("field analyzer not exist", func() {
		_, err := s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{1},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId: 999, // non-existent field
				},
			},
		})
		s.Require().Error(err)
	})

	s.Run("normal highlight with single analyzer", func() {
		err := s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
			Version: s.nextSchemaVersion(),
			Fields: []*schemapb.FieldSchema{
				{
					FieldID:      103,
					Name:         "id",
					DataType:     schemapb.DataType_Int64,
					IsPrimaryKey: true,
					AutoID:       true,
				},
				{
					FieldID:  100,
					Name:     "text",
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{Key: common.MaxLengthKey, Value: "256"},
						{Key: "analyzer_params", Value: "{}"},
					},
				},
				{
					FieldID:  101,
					Name:     "sparse",
					DataType: schemapb.DataType_SparseFloatVector,
				},
			},
			Functions: []*schemapb.FunctionSchema{{
				Type:             schemapb.FunctionType_BM25,
				InputFieldNames:  []string{"text"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"sparse"},
				OutputFieldIds:   []int64{101},
			}},
		}, nil, s.nextSchemaBarrierLoadMeta())
		s.Require().NoError(err)
		s.ResetDelegator()

		result, err := s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{2},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId:       100,
					Texts:         []string{"test", "this is a test document", "another test case"},
					SearchTextNum: 1,
					CorpusTextNum: 2,
				},
			},
		})
		s.Require().NoError(err)
		s.Require().Equal(2, len(result))
		// Check that we got highlight results
		s.Require().NotNil(result[0].Fragments)
		s.Require().NotNil(result[1].Fragments)
	})

	s.Run("highlight with standalone analyzer", func() {
		err := s.manager.Collection.PutOrRef(s.collectionID, newFunctionRuntimeTestSchemaWithVersion(s.nextSchemaVersion()), nil, s.nextSchemaBarrierLoadMeta())
		s.Require().NoError(err)
		s.ResetDelegator()

		result, err := s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{2},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId:       103,
					FieldName:     "standalone_text",
					Texts:         []string{"test", "this is a test document", "another test case"},
					SearchTextNum: 1,
					CorpusTextNum: 2,
				},
			},
		})
		s.Require().NoError(err)
		s.Require().Equal(2, len(result))
		s.Require().NotNil(result[0].Fragments)
		s.Require().NotNil(result[1].Fragments)
	})

	s.Run("highlight with multi analyzer", func() {
		err := s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
			Version: s.nextSchemaVersion(),
			Fields: []*schemapb.FieldSchema{
				{
					FieldID:      103,
					Name:         "id",
					DataType:     schemapb.DataType_Int64,
					IsPrimaryKey: true,
					AutoID:       true,
				},
				{
					FieldID:  100,
					Name:     "text",
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{Key: common.MaxLengthKey, Value: "256"},
						{Key: "multi_analyzer_params", Value: `{
						"by_field": "analyzer",
    					"analyzers": {
							"standard": {},
							"default": {}
						}
					}`},
					},
				},
				{
					FieldID:  101,
					Name:     "sparse",
					DataType: schemapb.DataType_SparseFloatVector,
				},
				{
					FieldID:    102,
					Name:       "analyzer",
					DataType:   schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{{Key: common.MaxLengthKey, Value: "256"}},
				},
			},
			Functions: []*schemapb.FunctionSchema{{
				Type:             schemapb.FunctionType_BM25,
				InputFieldNames:  []string{"text"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"sparse"},
				OutputFieldIds:   []int64{101},
			}},
		}, nil, s.nextSchemaBarrierLoadMeta())
		s.Require().NoError(err)
		s.ResetDelegator()

		// two target with two analyzer
		result, err := s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{1, 1},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId:       100,
					Texts:         []string{"test1", "test2", "this is a test1 document", "another test2 case"},
					AnalyzerNames: []string{"default", "standard", "default", "default"},
					SearchTextNum: 2,
					CorpusTextNum: 2,
				},
			},
		})
		s.Require().NoError(err)
		s.Require().Equal(2, len(result))

		result, err = s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{1},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId:       100,
					Texts:         []string{"test1", "test2", "this is a test1 document"},
					AnalyzerNames: []string{"default", "default", "default"},
					SearchTextNum: 1,
					CorpusTextNum: 1,
					Queries:       []*querypb.HighlightQuery{{Type: querypb.HighlightQueryType_TextMatch}},
				},
			},
		})
		s.Require().NoError(err)
		s.Require().Equal(1, len(result))

		_, err = s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{1},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId:       100,
					Texts:         []string{"test1", "test2", "this is a test1 document"},
					SearchTextNum: 1,
					CorpusTextNum: 1,
					Queries:       []*querypb.HighlightQuery{{Type: querypb.HighlightQueryType_TextMatch}},
				},
			},
		})
		s.Require().Error(err)

		_, err = s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{1},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId:       100,
					Texts:         []string{"test1", "test2", "this is a test1 document"},
					AnalyzerNames: []string{"standard"},
					SearchTextNum: 1,
					CorpusTextNum: 1,
					Queries:       []*querypb.HighlightQuery{{Type: querypb.HighlightQueryType_TextMatch}},
				},
			},
		})
		s.Require().Error(err)

		_, err = s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{1},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId:       100,
					Texts:         []string{"test1", "test2", "this is a test1 document"},
					AnalyzerNames: []string{"default", "standard"},
					SearchTextNum: 1,
					CorpusTextNum: 1,
					Queries:       []*querypb.HighlightQuery{{Type: querypb.HighlightQueryType_TextMatch}},
				},
			},
		})
		s.Require().Error(err)
	})

	s.Run("empty target texts", func() {
		err := s.manager.Collection.PutOrRef(s.collectionID, &schemapb.CollectionSchema{
			Version: s.nextSchemaVersion(),
			Fields: []*schemapb.FieldSchema{
				{
					FieldID:      103,
					Name:         "id",
					DataType:     schemapb.DataType_Int64,
					IsPrimaryKey: true,
					AutoID:       true,
				},
				{
					FieldID:  100,
					Name:     "text",
					DataType: schemapb.DataType_VarChar,
					TypeParams: []*commonpb.KeyValuePair{
						{Key: common.MaxLengthKey, Value: "256"},
						{Key: "analyzer_params", Value: "{}"},
					},
				},
				{
					FieldID:  101,
					Name:     "sparse",
					DataType: schemapb.DataType_SparseFloatVector,
				},
			},
			Functions: []*schemapb.FunctionSchema{{
				Type:             schemapb.FunctionType_BM25,
				InputFieldNames:  []string{"text"},
				InputFieldIds:    []int64{100},
				OutputFieldNames: []string{"sparse"},
				OutputFieldIds:   []int64{101},
			}},
		}, nil, s.nextSchemaBarrierLoadMeta())
		s.Require().NoError(err)
		s.ResetDelegator()

		result, err := s.delegator.GetHighlight(ctx, &querypb.GetHighlightRequest{
			Topks: []int64{1},
			Tasks: []*querypb.HighlightTask{
				{
					FieldId:       100,
					Texts:         []string{"test document"},
					SearchTextNum: 0,
					CorpusTextNum: 1,
				},
			},
		})
		s.Require().NoError(err)
		s.Require().NotNil(result)
	})
}

// TestDelegatorLifetimeIntegration tests the integration of lifetime state checks with main delegator methods
func (s *DelegatorSuite) TestDelegatorLifetimeIntegration() {
	sd := s.delegator.(*shardDelegator)
	ctx := context.Background()

	s.Run("test_methods_fail_when_not_working", func() {
		// Set delegator to Initializing state (not ready)
		sd.lifetime.SetState(lifetime.Initializing)

		// Search should fail when not ready
		_, err := sd.Search(ctx, &querypb.SearchRequest{
			Req:         &internalpb.SearchRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")

		// Query should fail when not ready
		_, err = sd.Query(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")

		// GetStatistics should fail when not ready
		_, err = sd.GetStatistics(ctx, &querypb.GetStatisticsRequest{
			Req:         &internalpb.GetStatisticsRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")

		// UpdateSchema should fail when not ready
		err = sd.UpdateSchema(ctx, &schemapb.CollectionSchema{Name: "test"}, 1)
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")
	})

	s.Run("test_methods_fail_when_stopped", func() {
		// Set delegator to Stopped state
		sd.lifetime.SetState(lifetime.Stopped)

		// Search should fail when stopped
		_, err := sd.Search(ctx, &querypb.SearchRequest{
			Req:         &internalpb.SearchRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")

		// Query should fail when stopped
		_, err = sd.Query(ctx, &querypb.QueryRequest{
			Req:         &internalpb.RetrieveRequest{QueryLabel: "query", Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")

		// GetStatistics should fail when stopped
		_, err = sd.GetStatistics(ctx, &querypb.GetStatisticsRequest{
			Req:         &internalpb.GetStatisticsRequest{Base: commonpbutil.NewMsgBase()},
			DmlChannels: []string{s.vchannelName},
		})
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")

		// UpdateSchema should fail when stopped
		err = sd.UpdateSchema(ctx, &schemapb.CollectionSchema{Name: "test"}, 1)
		s.Error(err)
		s.Contains(err.Error(), "delegator is not ready")
	})
}

// TestDelegatorStateTransitions tests state transitions and edge cases
func (s *DelegatorSuite) TestDelegatorStateTransitions() {
	sd := s.delegator.(*shardDelegator)

	s.Run("test_state_transition_sequence", func() {
		// Test normal state transition sequence

		// Start from Initializing
		sd.lifetime.SetState(lifetime.Initializing)
		s.False(sd.Serviceable())
		s.False(sd.Stopped())

		// Transition to Working
		sd.Start() // This calls lifetime.SetState(lifetime.Working)
		s.True(sd.Serviceable())
		s.False(sd.Stopped())

		// Transition to Stopped
		sd.lifetime.SetState(lifetime.Stopped)
		s.False(sd.Serviceable())
		s.True(sd.Stopped())
	})

	s.Run("test_multiple_start_calls", func() {
		// Test that multiple Start() calls don't cause issues
		sd.lifetime.SetState(lifetime.Initializing)
		s.False(sd.Serviceable())

		// Call Start multiple times
		sd.Start()
		s.True(sd.Serviceable())

		sd.Start()
		sd.Start()
		s.True(sd.Serviceable()) // Should remain serviceable
	})

	s.Run("test_start_after_stopped", func() {
		// Test starting after being stopped
		sd.lifetime.SetState(lifetime.Stopped)
		s.True(sd.Stopped())
		s.False(sd.Serviceable())

		// Start again
		sd.Start()
		s.False(sd.Stopped())
		s.True(sd.Serviceable())
	})

	s.Run("test_consistency_between_methods", func() {
		// Test consistency between Serviceable() and Stopped() methods

		// In Initializing state
		sd.lifetime.SetState(lifetime.Initializing)
		serviceable := sd.Serviceable()
		stopped := sd.Stopped()
		s.False(serviceable)
		s.False(stopped)

		// In Working state
		sd.lifetime.SetState(lifetime.Working)
		serviceable = sd.Serviceable()
		stopped = sd.Stopped()
		s.True(serviceable)
		s.False(stopped)

		// In Stopped state
		sd.lifetime.SetState(lifetime.Stopped)
		serviceable = sd.Serviceable()
		stopped = sd.Stopped()
		s.False(serviceable)
		s.True(stopped)
	})

	s.Run("test_error_types_and_wrapping", func() {
		// Test that errors are properly wrapped with channel information

		// Test NotStopped error
		err := sd.NotStopped(lifetime.Stopped)
		s.Error(err)
		s.True(errors.Is(err, merr.ErrChannelNotAvailable))

		// Test IsWorking error
		err = sd.IsWorking(lifetime.Initializing)
		s.Error(err)
		s.True(errors.Is(err, merr.ErrChannelNotAvailable))

		err = sd.IsWorking(lifetime.Stopped)
		s.Error(err)
		s.True(errors.Is(err, merr.ErrChannelNotAvailable))
	})
}

func TestDelegatorSuite(t *testing.T) {
	suite.Run(t, new(DelegatorSuite))
}

func newFunctionRuntimeTestSchema(functions ...*schemapb.FunctionSchema) *schemapb.CollectionSchema {
	return &schemapb.CollectionSchema{
		Name: "TestCollection",
		Fields: []*schemapb.FieldSchema{
			{
				Name:         "id",
				FieldID:      100,
				IsPrimaryKey: true,
				DataType:     schemapb.DataType_Int64,
				AutoID:       true,
			},
			{
				Name:     "text",
				FieldID:  101,
				DataType: schemapb.DataType_VarChar,
				TypeParams: []*commonpb.KeyValuePair{
					{Key: common.MaxLengthKey, Value: "256"},
					{Key: common.EnableAnalyzerKey, Value: "true"},
					{Key: common.AnalyzerParamKey, Value: "{}"},
				},
			},
			{
				Name:     "sparse",
				FieldID:  102,
				DataType: schemapb.DataType_SparseFloatVector,
			},
			{
				Name:     "standalone_text",
				FieldID:  103,
				DataType: schemapb.DataType_VarChar,
				TypeParams: []*commonpb.KeyValuePair{
					{Key: common.MaxLengthKey, Value: "256"},
					{Key: common.EnableAnalyzerKey, Value: "true"},
					{Key: common.AnalyzerParamKey, Value: "{}"},
				},
			},
			{
				Name:     "minhash",
				FieldID:  104,
				DataType: schemapb.DataType_BinaryVector,
				TypeParams: []*commonpb.KeyValuePair{
					{Key: common.DimKey, Value: "128"},
				},
			},
			{
				Name:     "sparse_additive",
				FieldID:  105,
				DataType: schemapb.DataType_SparseFloatVector,
			},
		},
		Functions: functions,
	}
}

func newFunctionRuntimeTestSchemaWithVersion(version int32, functions ...*schemapb.FunctionSchema) *schemapb.CollectionSchema {
	schema := newFunctionRuntimeTestSchema(functions...)
	schema.Version = version
	return schema
}

func newBM25FunctionSchema() *schemapb.FunctionSchema {
	return &schemapb.FunctionSchema{
		Type:             schemapb.FunctionType_BM25,
		InputFieldNames:  []string{"text"},
		InputFieldIds:    []int64{101},
		OutputFieldNames: []string{"sparse"},
		OutputFieldIds:   []int64{102},
	}
}

func newAdditionalBM25FunctionSchema() *schemapb.FunctionSchema {
	return &schemapb.FunctionSchema{
		Type:             schemapb.FunctionType_BM25,
		InputFieldNames:  []string{"text"},
		InputFieldIds:    []int64{101},
		OutputFieldNames: []string{"sparse_additive"},
		OutputFieldIds:   []int64{105},
	}
}

func newMinHashFunctionSchema() *schemapb.FunctionSchema {
	return &schemapb.FunctionSchema{
		Type:             schemapb.FunctionType_MinHash,
		InputFieldNames:  []string{"text"},
		InputFieldIds:    []int64{101},
		OutputFieldNames: []string{"minhash"},
		OutputFieldIds:   []int64{104},
	}
}

func TestBuildFunctionRuntimeStateAddsBM25FunctionType(t *testing.T) {
	paramtable.Init()
	schema := newFunctionRuntimeTestSchema(newBM25FunctionSchema())

	state, err := buildFunctionRuntimeState(schema)
	require.NoError(t, err)

	runnerCount, functionTypeCount, analyzerCount := state.counts()
	assert.Equal(t, 0, runnerCount)
	assert.Equal(t, 1, functionTypeCount)
	assert.Equal(t, 0, analyzerCount)
	assert.True(t, state.hasFunctionType(102, schemapb.FunctionType_BM25))
}

func TestBuildFunctionRuntimeStateAddsMinHashFunctionType(t *testing.T) {
	paramtable.Init()
	schema := newFunctionRuntimeTestSchema(newMinHashFunctionSchema())

	state, err := buildFunctionRuntimeState(schema)
	require.NoError(t, err)

	runnerCount, functionTypeCount, analyzerCount := state.counts()
	assert.Equal(t, 0, runnerCount)
	assert.Equal(t, 1, functionTypeCount)
	assert.Equal(t, 0, analyzerCount)
	assert.True(t, state.hasFunctionType(104, schemapb.FunctionType_MinHash))
}

func TestBM25FunctionSetAdditiveValidation(t *testing.T) {
	base := newBM25FunctionSet(newFunctionRuntimeTestSchema(newBM25FunctionSchema()))

	assert.True(t, newBM25FunctionSet(newFunctionRuntimeTestSchema(newBM25FunctionSchema())).IsSupersetOf(base))
	assert.True(t, newBM25FunctionSet(newFunctionRuntimeTestSchema(newBM25FunctionSchema(), newAdditionalBM25FunctionSchema())).IsSupersetOf(base))

	changed := proto.Clone(newBM25FunctionSchema()).(*schemapb.FunctionSchema)
	changed.InputFieldIds = []int64{103}
	changedSet := newBM25FunctionSet(newFunctionRuntimeTestSchema(changed))
	assert.False(t, changedSet.IsSupersetOf(base))
	assert.True(t, changedSet.HasIncompatibleCommonFunction(base))

	droppedSet := newBM25FunctionSet(newFunctionRuntimeTestSchema())
	assert.False(t, droppedSet.IsSupersetOf(base))
	assert.False(t, droppedSet.HasIncompatibleCommonFunction(base))

	readded := proto.Clone(newBM25FunctionSchema()).(*schemapb.FunctionSchema)
	readded.OutputFieldIds = []int64{104}
	readdedSet := newBM25FunctionSet(newFunctionRuntimeTestSchema(readded))
	assert.False(t, readdedSet.IsSupersetOf(base))
	assert.False(t, readdedSet.HasIncompatibleCommonFunction(base))
}

func TestBM25FunctionSetIgnoresAnalyzerParams(t *testing.T) {
	baseSchema := newFunctionRuntimeTestSchema(newBM25FunctionSchema())
	changedSchema := proto.Clone(baseSchema).(*schemapb.CollectionSchema)
	changedSchema.Fields[1].TypeParams = []*commonpb.KeyValuePair{
		{Key: common.MaxLengthKey, Value: "256"},
		{Key: common.EnableAnalyzerKey, Value: "true"},
		{Key: common.AnalyzerParamKey, Value: `{"tokenizer":"standard"}`},
	}

	base := newBM25FunctionSet(baseSchema)
	changed := newBM25FunctionSet(changedSchema)
	assert.True(t, changed.IsSupersetOf(base))
}

func TestBM25FunctionSetNormalizesFunctionParamOrder(t *testing.T) {
	baseFunction := newBM25FunctionSchema()
	baseFunction.Params = []*commonpb.KeyValuePair{
		{Key: "k1", Value: "v1"},
		{Key: "k2", Value: "v2"},
	}
	changedFunction := proto.Clone(baseFunction).(*schemapb.FunctionSchema)
	changedFunction.Params = []*commonpb.KeyValuePair{
		{Key: "k2", Value: "v2"},
		{Key: "k1", Value: "v1"},
	}

	base := newBM25FunctionSet(newFunctionRuntimeTestSchema(baseFunction))
	changed := newBM25FunctionSet(newFunctionRuntimeTestSchema(changedFunction))
	assert.True(t, changed.IsSupersetOf(base))
}

func TestUpdateSchemaRejectsIncompatibleBM25FunctionChange(t *testing.T) {
	paramtable.Init()
	paramtable.SetNodeID(1)
	workerManager := cluster.NewMockManager(t)
	oldSchema := newFunctionRuntimeTestSchema(newBM25FunctionSchema())
	oldOracle := NewIDFOracle("test-channel", oldSchema.GetFunctions())
	sd := &shardDelegator{
		collectionID:  1000,
		vchannelName:  "test-channel",
		collection:    segments.NewCollectionWithoutSegcoreForTest(1000, oldSchema),
		lifetime:      lifetime.NewLifetime(lifetime.Working),
		distribution:  NewDistribution("test-channel", NewChannelQueryView(nil, nil, nil, initialTargetVersion)),
		workerManager: workerManager,
		functionState: &functionRuntimeState{
			functionFieldType: map[int64]schemapb.FunctionType{102: schemapb.FunctionType_BM25},
		},
		deleteBuffer:               deletebuffer.NewListDeleteBuffer[*deletebuffer.Item](0, 0, []string{"1", "test-channel"}),
		tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
		latestRequiredMVCCTimeTick: atomic.NewUint64(0),
	}
	sd.publishIDFOracle(oldOracle)
	defer sd.Close()

	changed := proto.Clone(newBM25FunctionSchema()).(*schemapb.FunctionSchema)
	changed.InputFieldIds = []int64{103}
	err := sd.UpdateSchema(context.Background(), newFunctionRuntimeTestSchemaWithVersion(1, changed), 100)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "unsupported incompatible BM25 function schema change")
	assert.True(t, sd.functionState.hasFunctionType(102, schemapb.FunctionType_BM25))
	assert.Same(t, oldOracle, sd.getIDFOracle())
}

func TestUpdateSchemaSyncsAdditiveIDFOracleFunctions(t *testing.T) {
	paramtable.Init()
	paramtable.SetNodeID(1)
	worker := cluster.NewMockWorker(t)
	worker.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).Return(merr.Success(), nil).Once()
	workerManager := cluster.NewMockManager(t)
	workerManager.EXPECT().GetWorker(mock.Anything, int64(1)).Return(worker, nil).Once()
	oldSchema := newFunctionRuntimeTestSchema(newBM25FunctionSchema())
	sd := &shardDelegator{
		collectionID:  1000,
		vchannelName:  "test-channel",
		collection:    segments.NewCollectionWithoutSegcoreForTest(1000, oldSchema),
		lifetime:      lifetime.NewLifetime(lifetime.Working),
		distribution:  NewDistribution("test-channel", NewChannelQueryView(nil, nil, nil, initialTargetVersion)),
		workerManager: workerManager,
		functionState: &functionRuntimeState{
			functionFieldType: map[int64]schemapb.FunctionType{},
		},
		deleteBuffer:               deletebuffer.NewListDeleteBuffer[*deletebuffer.Item](0, 0, []string{"1", "test-channel"}),
		tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
		latestRequiredMVCCTimeTick: atomic.NewUint64(0),
	}
	oldOracle := NewIDFOracle("test-channel", oldSchema.GetFunctions())
	sd.publishIDFOracle(oldOracle)
	defer sd.Close()

	newSchema := newFunctionRuntimeTestSchemaWithVersion(1, newBM25FunctionSchema(), newAdditionalBM25FunctionSchema())
	collectionManager := segments.NewMockCollectionManager(t)
	collectionManager.EXPECT().UpdateSchema(int64(1000), newSchema, uint64(100)).Return(nil).Once()
	sd.collectionManager = collectionManager

	err := sd.UpdateSchema(context.Background(), newSchema, 100)
	require.NoError(t, err)

	_, _, err = sd.getIDFOracle().BuildIDF(105, &schemapb.SparseFloatArray{Contents: [][]byte{typeutil.CreateAndSortSparseFloatRow(map[uint32]float32{7: 1})}, Dim: 1})
	require.NoError(t, err)
}

func TestUpdateSchemaDoesNotSyncIDFOracleWhenWorkerUpdateFails(t *testing.T) {
	paramtable.Init()
	paramtable.SetNodeID(1)
	worker := cluster.NewMockWorker(t)
	worker.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).Return(merr.Status(merr.WrapErrServiceInternal("mocked")), merr.WrapErrServiceInternal("mocked")).Once()
	workerManager := cluster.NewMockManager(t)
	workerManager.EXPECT().GetWorker(mock.Anything, int64(1)).Return(worker, nil).Once()
	oldSchema := newFunctionRuntimeTestSchema(newBM25FunctionSchema())
	sd := &shardDelegator{
		collectionID:  1000,
		vchannelName:  "test-channel",
		collection:    segments.NewCollectionWithoutSegcoreForTest(1000, oldSchema),
		lifetime:      lifetime.NewLifetime(lifetime.Working),
		distribution:  NewDistribution("test-channel", NewChannelQueryView(nil, nil, nil, initialTargetVersion)),
		workerManager: workerManager,
		functionState: &functionRuntimeState{
			functionFieldType: map[int64]schemapb.FunctionType{},
		},
		deleteBuffer:               deletebuffer.NewListDeleteBuffer[*deletebuffer.Item](0, 0, []string{"1", "test-channel"}),
		tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
		latestRequiredMVCCTimeTick: atomic.NewUint64(0),
	}
	oldOracle := NewIDFOracle("test-channel", oldSchema.GetFunctions())
	sd.publishIDFOracle(oldOracle)
	defer sd.Close()

	newSchema := newFunctionRuntimeTestSchemaWithVersion(1, newBM25FunctionSchema(), newAdditionalBM25FunctionSchema())
	err := sd.UpdateSchema(context.Background(), newSchema, 100)
	require.Error(t, err)

	_, _, err = sd.getIDFOracle().BuildIDF(105, &schemapb.SparseFloatArray{Contents: [][]byte{typeutil.CreateAndSortSparseFloatRow(map[uint32]float32{7: 1})}, Dim: 1})
	require.Error(t, err)
}

func TestUpdateSchemaInitializesIDFOracleWhenBM25Added(t *testing.T) {
	paramtable.Init()
	paramtable.SetNodeID(1)
	worker := cluster.NewMockWorker(t)
	worker.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).Return(merr.Success(), nil).Once()
	workerManager := cluster.NewMockManager(t)
	workerManager.EXPECT().GetWorker(mock.Anything, int64(1)).Return(worker, nil).Once()
	oldSchema := newFunctionRuntimeTestSchema()
	sd := &shardDelegator{
		collectionID:  1000,
		vchannelName:  "test-channel",
		collection:    segments.NewCollectionWithoutSegcoreForTest(1000, oldSchema),
		lifetime:      lifetime.NewLifetime(lifetime.Working),
		distribution:  NewDistribution("test-channel", NewChannelQueryView(nil, nil, nil, initialTargetVersion)),
		workerManager: workerManager,
		functionState: &functionRuntimeState{
			functionFieldType: map[int64]schemapb.FunctionType{},
		},
		deleteBuffer:               deletebuffer.NewListDeleteBuffer[*deletebuffer.Item](0, 0, []string{"1", "test-channel"}),
		tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
		latestRequiredMVCCTimeTick: atomic.NewUint64(0),
	}
	defer sd.Close()

	newSchema := newFunctionRuntimeTestSchemaWithVersion(1, newBM25FunctionSchema())
	collectionManager := segments.NewMockCollectionManager(t)
	collectionManager.EXPECT().UpdateSchema(int64(1000), newSchema, uint64(100)).Return(nil).Once()
	sd.collectionManager = collectionManager

	err := sd.UpdateSchema(context.Background(), newSchema, 100)
	require.NoError(t, err)
	require.NotNil(t, sd.getIDFOracle())

	_, _, err = sd.getIDFOracle().BuildIDF(102, &schemapb.SparseFloatArray{Contents: [][]byte{typeutil.CreateAndSortSparseFloatRow(map[uint32]float32{7: 1})}, Dim: 1})
	require.NoError(t, err)
}

func TestUpdateSchemaRefreshesCollectionBaselineForSequentialBM25Validation(t *testing.T) {
	paramtable.Init()
	paramtable.SetNodeID(1)
	manager := segments.NewManager()
	oldSchema := newFunctionRuntimeTestSchema()
	require.NoError(t, manager.Collection.PutOrRef(1000, oldSchema, nil, &querypb.LoadMetaInfo{SchemaBarrierTs: 1}))
	defer manager.Collection.Unref(1000, 1)

	worker := cluster.NewMockWorker(t)
	worker.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).Return(merr.Success(), nil).Once()
	workerManager := cluster.NewMockManager(t)
	workerManager.EXPECT().GetWorker(mock.Anything, int64(1)).Return(worker, nil).Once()

	functionState, err := buildFunctionRuntimeState(oldSchema)
	require.NoError(t, err)
	sd := &shardDelegator{
		collectionID:               1000,
		vchannelName:               "test-channel",
		collection:                 manager.Collection.Get(1000),
		collectionManager:          manager.Collection,
		lifetime:                   lifetime.NewLifetime(lifetime.Working),
		distribution:               NewDistribution("test-channel", NewChannelQueryView(nil, nil, nil, initialTargetVersion)),
		workerManager:              workerManager,
		functionState:              functionState,
		deleteBuffer:               deletebuffer.NewListDeleteBuffer[*deletebuffer.Item](0, 0, []string{"1", "test-channel"}),
		tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
		latestRequiredMVCCTimeTick: atomic.NewUint64(0),
	}
	defer sd.Close()

	firstSchema := newFunctionRuntimeTestSchemaWithVersion(1, newBM25FunctionSchema())
	err = sd.UpdateSchema(context.Background(), firstSchema, 100)
	require.NoError(t, err)
	require.Equal(t, uint64(1), sd.collection.SchemaVersion())

	changed := proto.Clone(newBM25FunctionSchema()).(*schemapb.FunctionSchema)
	changed.InputFieldIds = []int64{103}
	err = sd.UpdateSchema(context.Background(), newFunctionRuntimeTestSchemaWithVersion(2, changed), 200)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "unsupported incompatible BM25 function schema change")
	require.Equal(t, uint64(1), sd.collection.SchemaVersion())
}

func TestUpdateSchemaSyncsFunctionRuntimeMetadata(t *testing.T) {
	paramtable.Init()
	paramtable.SetNodeID(1)
	worker := cluster.NewMockWorker(t)
	worker.EXPECT().UpdateSchema(mock.Anything, mock.AnythingOfType("*querypb.UpdateSchemaRequest")).Return(merr.Success(), nil).Once()
	workerManager := cluster.NewMockManager(t)
	workerManager.EXPECT().GetWorker(mock.Anything, int64(1)).Return(worker, nil).Once()
	sd := &shardDelegator{
		collectionID:  1000,
		vchannelName:  "test-channel",
		collection:    segments.NewCollectionWithoutSegcoreForTest(1000, newFunctionRuntimeTestSchema()),
		lifetime:      lifetime.NewLifetime(lifetime.Working),
		distribution:  NewDistribution("test-channel", NewChannelQueryView(nil, nil, nil, initialTargetVersion)),
		workerManager: workerManager,
		functionState: &functionRuntimeState{
			functionFieldType: map[int64]schemapb.FunctionType{999: schemapb.FunctionType_BM25},
		},
		deleteBuffer:               deletebuffer.NewListDeleteBuffer[*deletebuffer.Item](0, 0, []string{"1", "test-channel"}),
		tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
		latestRequiredMVCCTimeTick: atomic.NewUint64(0),
	}
	defer sd.Close()

	newSchema := newFunctionRuntimeTestSchemaWithVersion(1, newBM25FunctionSchema(), newMinHashFunctionSchema())
	collectionManager := segments.NewMockCollectionManager(t)
	collectionManager.EXPECT().UpdateSchema(int64(1000), newSchema, uint64(100)).Return(nil).Once()
	sd.collectionManager = collectionManager

	err := sd.UpdateSchema(context.Background(), newSchema, 100)
	require.NoError(t, err)

	assert.True(t, sd.functionState.hasFunctionType(102, schemapb.FunctionType_BM25))
	assert.True(t, sd.functionState.hasFunctionType(104, schemapb.FunctionType_MinHash))
	assert.False(t, sd.functionState.hasFunctionType(999, schemapb.FunctionType_BM25))
}

func TestUpdateSchemaSkipsStaleSchemaBeforeSideEffects(t *testing.T) {
	paramtable.Init()
	paramtable.SetNodeID(1)
	currentSchema := newFunctionRuntimeTestSchemaWithVersion(2, newMinHashFunctionSchema())
	currentFunctionState, err := buildFunctionRuntimeState(currentSchema)
	require.NoError(t, err)
	workerManager := cluster.NewMockManager(t)
	collectionManager := segments.NewMockCollectionManager(t)
	sd := &shardDelegator{
		collectionID:               1000,
		vchannelName:               "test-channel",
		collection:                 segments.NewCollectionWithoutSegcoreForTest(1000, currentSchema),
		collectionManager:          collectionManager,
		lifetime:                   lifetime.NewLifetime(lifetime.Working),
		distribution:               NewDistribution("test-channel", NewChannelQueryView(nil, nil, nil, initialTargetVersion)),
		workerManager:              workerManager,
		functionState:              currentFunctionState,
		schemaBarrierTs:            100,
		deleteBuffer:               deletebuffer.NewListDeleteBuffer[*deletebuffer.Item](0, 0, []string{"1", "test-channel"}),
		tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
		latestRequiredMVCCTimeTick: atomic.NewUint64(0),
	}
	defer sd.Close()

	staleSchema := newFunctionRuntimeTestSchemaWithVersion(1, newBM25FunctionSchema())
	err = sd.UpdateSchema(context.Background(), staleSchema, 200)
	require.NoError(t, err)

	assert.Equal(t, uint64(100), sd.schemaBarrierTs)
	assert.True(t, sd.functionState.hasFunctionType(104, schemapb.FunctionType_MinHash))
	assert.False(t, sd.functionState.hasFunctionType(102, schemapb.FunctionType_BM25))
}

func TestDelegatorSearchBM25InvalidMetricType(t *testing.T) {
	paramtable.Init()
	searchReq := &querypb.SearchRequest{
		Req: &internalpb.SearchRequest{
			Base: commonpbutil.NewMsgBase(),
		},
	}

	searchReq.Req.FieldId = 101
	searchReq.Req.MetricType = metric.IP

	sd := &shardDelegator{
		functionState: &functionRuntimeState{
			functionFieldType: map[int64]schemapb.FunctionType{101: schemapb.FunctionType_BM25},
		},
		latestRequiredMVCCTimeTick: atomic.NewUint64(0),
	}

	_, err := sd.search(context.Background(), searchReq, []SnapshotItem{}, []SegmentEntry{}, map[int64]int64{})
	require.Error(t, err)
	assert.Contains(t, err.Error(), "must use BM25 metric type when searching against BM25 Function output field")
}

func TestNewRowCountBasedEvaluator_SearchAndQueryTasks(t *testing.T) {
	mockey.PatchConvey("TestNewRowCountBasedEvaluator_SearchAndQueryTasks", t, func() {
		sealedRowCount := map[int64]int64{
			1: 1000,
			2: 2000,
			3: 3000,
		}

		// Mock paramtable configuration
		mockParamTable := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsFloat")).Return(0.8).Build()
		defer mockParamTable.UnPatch()

		evaluator := NewRowCountBasedEvaluator(sealedRowCount)

		// Test Search task - success segments meet required ratio
		successSegments := typeutil.NewSet[int64]()
		successSegments.Insert(1, 2) // 3000 rows out of 6000 total = 0.5
		failureSegments := []int64{3}
		testErrors := []error{errors.New("segment 3 failed")}

		shouldReturn, accessedRatio := evaluator("Search", successSegments, failureSegments, testErrors)
		assert.False(t, shouldReturn) // 0.5 < 0.8, should not return partial
		assert.Equal(t, 0.5, accessedRatio)

		// Test Query task - Query never allows partial results (only Search does)
		successSegments = typeutil.NewSet[int64]()
		successSegments.Insert(1, 2) // 3000 rows out of 6000 total = 0.5
		failureSegments = []int64{3}
		testErrors = []error{errors.New("segment 3 failed")}

		shouldReturn, accessedRatio = evaluator("Query", successSegments, failureSegments, testErrors)
		assert.False(t, shouldReturn, "Query should never return partial results")
		assert.Equal(t, 0.0, accessedRatio) // ratio forced to 1.0, early return
	})
}

func TestNewRowCountBasedEvaluator_RequiredRatioConfiguration(t *testing.T) {
	mockey.PatchConvey("TestNewRowCountBasedEvaluator_RequiredRatioConfiguration", t, func() {
		sealedRowCount := map[int64]int64{
			1: 1000,
			2: 2000,
		}

		// Test with required ratio >= 1.0 (should always return false for partial results)
		mockParamTable := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsFloat")).Return(1.0).Build()
		defer mockParamTable.UnPatch()

		evaluator := NewRowCountBasedEvaluator(sealedRowCount)
		successSegments := typeutil.NewSet[int64]()
		successSegments.Insert(1) // Only partial data available

		shouldReturn, accessedRatio := evaluator("Search", successSegments, []int64{2}, []error{errors.New("test error")})
		assert.False(t, shouldReturn) // Should not return partial when ratio >= 1.0
		assert.Equal(t, 0.0, accessedRatio)

		// Test non-Search/Query task (should use ratio 1.0)
		shouldReturn, accessedRatio = evaluator("GetStatistics", successSegments, []int64{2}, []error{errors.New("test error")})
		assert.False(t, shouldReturn)
		assert.Equal(t, 0.0, accessedRatio)
	})
}

func TestNewRowCountBasedEvaluator_PartialResultAcceptance(t *testing.T) {
	mockey.PatchConvey("TestNewRowCountBasedEvaluator_PartialResultAcceptance", t, func() {
		sealedRowCount := map[int64]int64{
			1: 1000, // 1000 rows
			2: 2000, // 2000 rows
			3: 3000, // 3000 rows
			4: 4000, // 4000 rows
		}
		// Total: 10000 rows

		// Mock paramtable to require 70% data availability
		mockParamTable := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsFloat")).Return(0.7).Build()
		defer mockParamTable.UnPatch()

		evaluator := NewRowCountBasedEvaluator(sealedRowCount)

		// Test case: 80% data available (should accept partial result)
		successSegments := typeutil.NewSet[int64]()
		successSegments.Insert(1, 2, 3) // 6000 rows out of 10000 = 0.6
		failureSegments := []int64{4}
		testErrors := []error{errors.New("segment 4 failed")}

		shouldReturn, accessedRatio := evaluator("Search", successSegments, failureSegments, testErrors)
		assert.False(t, shouldReturn) // 0.6 < 0.7, should not return partial
		assert.Equal(t, 0.6, accessedRatio)

		// Test case: 90% data available with Search (should accept partial result)
		successSegments = typeutil.NewSet[int64]()
		successSegments.Insert(2, 3, 4) // 9000 rows out of 10000 = 0.9
		failureSegments = []int64{1}
		testErrors = []error{errors.New("segment 1 failed")}

		shouldReturn, accessedRatio = evaluator("Search", successSegments, failureSegments, testErrors)
		assert.True(t, shouldReturn) // 0.9 >= 0.7, should return partial for Search
		assert.Equal(t, 0.9, accessedRatio)

		// Test case: same 90% with Query (should NOT accept partial result)
		shouldReturn, accessedRatio = evaluator("Query", successSegments, failureSegments, testErrors)
		assert.False(t, shouldReturn, "Query should never return partial results")
		assert.Equal(t, 0.0, accessedRatio)
	})
}

func TestNewRowCountBasedEvaluator_ZeroRatioBoundary(t *testing.T) {
	mockey.PatchConvey("TestNewRowCountBasedEvaluator_ZeroRatioBoundary", t, func() {
		sealedRowCount := map[int64]int64{
			1: 1000,
			2: 2000,
			3: 3000,
			4: 4000,
		}
		// Total: 10000 rows

		// Mock ratio=0.0: search should never fail, even with empty results
		mockParamTable := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsFloat")).Return(0.0).Build()
		defer mockParamTable.UnPatch()

		evaluator := NewRowCountBasedEvaluator(sealedRowCount)

		// All segments failed (accessedDataRatio=0.0, ratio=0.0)
		// 0.0 >= 0.0 should return true
		successSegments := typeutil.NewSet[int64]()
		failureSegments := []int64{1, 2, 3, 4}
		testErrors := []error{
			errors.New("node 1 failed"),
			errors.New("node 2 failed"),
			errors.New("node 3 failed"),
			errors.New("node 4 failed"),
		}

		shouldReturn, accessedRatio := evaluator("Search", successSegments, failureSegments, testErrors)
		assert.True(t, shouldReturn, "ratio=0.0 with all segments failed should still return partial result")
		assert.Equal(t, 0.0, accessedRatio)

		// Some segments succeeded (accessedDataRatio=0.3 > 0.0)
		successSegments = typeutil.NewSet[int64]()
		successSegments.Insert(1, 2) // 3000/10000 = 0.3
		failureSegments = []int64{3, 4}
		testErrors = []error{errors.New("node 3 failed"), errors.New("node 4 failed")}

		shouldReturn, accessedRatio = evaluator("Search", successSegments, failureSegments, testErrors)
		assert.True(t, shouldReturn)
		assert.Equal(t, 0.3, accessedRatio)

		// Query task with ratio=0.0 and all failed - Query never allows partial
		successSegments = typeutil.NewSet[int64]()
		failureSegments = []int64{1, 2, 3, 4}
		testErrors = []error{errors.New("all failed")}

		shouldReturn, accessedRatio = evaluator("Query", successSegments, failureSegments, testErrors)
		assert.False(t, shouldReturn, "Query should never return partial results, even with ratio=0.0")
		assert.Equal(t, 0.0, accessedRatio)
	})
}

func TestNewRowCountBasedEvaluator_ExactRatioBoundary(t *testing.T) {
	mockey.PatchConvey("TestNewRowCountBasedEvaluator_ExactRatioBoundary", t, func() {
		sealedRowCount := map[int64]int64{
			1: 5000,
			2: 5000,
		}
		// Total: 10000 rows

		// Mock ratio=0.5: exactly 50% available should pass
		mockParamTable := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsFloat")).Return(0.5).Build()
		defer mockParamTable.UnPatch()

		evaluator := NewRowCountBasedEvaluator(sealedRowCount)

		// Exactly 50% data available (accessedDataRatio=0.5, ratio=0.5)
		// 0.5 >= 0.5 should return true
		successSegments := typeutil.NewSet[int64]()
		successSegments.Insert(1) // 5000/10000 = 0.5
		failureSegments := []int64{2}
		testErrors := []error{errors.New("segment 2 failed")}

		shouldReturn, accessedRatio := evaluator("Search", successSegments, failureSegments, testErrors)
		assert.True(t, shouldReturn, "exact ratio boundary (0.5 >= 0.5) should return partial result")
		assert.Equal(t, 0.5, accessedRatio)
	})
}

func TestNewRowCountBasedEvaluator_QueryStreamDisablesPartialResult(t *testing.T) {
	mockey.PatchConvey("TestNewRowCountBasedEvaluator_QueryStreamDisablesPartialResult", t, func() {
		sealedRowCount := map[int64]int64{
			1: 5000,
			2: 5000,
		}
		// Total: 10000 rows

		// Even with ratio=0.0, QueryStream should NOT return partial results
		mockParamTable := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsFloat")).Return(0.0).Build()
		defer mockParamTable.UnPatch()

		evaluator := NewRowCountBasedEvaluator(sealedRowCount)

		// 50% data available, ratio=0.0, but taskType=QueryStream
		// QueryStream must NOT allow partial results (used by delete-by-expr)
		successSegments := typeutil.NewSet[int64]()
		successSegments.Insert(1) // 5000/10000 = 0.5
		failureSegments := []int64{2}
		testErrors := []error{errors.New("segment 2 failed")}

		shouldReturn, _ := evaluator("QueryStream", successSegments, failureSegments, testErrors)
		assert.False(t, shouldReturn, "QueryStream should never return partial results, even with ratio=0.0")

		// Query also does NOT allow partial results (only Search does)
		shouldReturn, _ = evaluator("Query", successSegments, failureSegments, testErrors)
		assert.False(t, shouldReturn, "Query should never return partial results, even with ratio=0.0")

		// Only Search allows partial results
		shouldReturn, accessedRatio := evaluator("Search", successSegments, failureSegments, testErrors)
		assert.True(t, shouldReturn, "Search should allow partial results with ratio=0.0")
		assert.Equal(t, 0.5, accessedRatio)
	})
}

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

	t.Run("initial state is catching up", func(t *testing.T) {
		// Create a minimal delegator to test CatchingUpStreamingData
		sd := &shardDelegator{
			catchingUpStreamingData:    atomic.NewBool(true),
			latestRequiredMVCCTimeTick: atomic.NewUint64(0),
		}
		assert.True(t, sd.CatchingUpStreamingData())
	})

	t.Run("state changes to caught up when lag is small", func(t *testing.T) {
		// Mock the config to return 5 seconds threshold
		mockParam := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsDurationByParse")).Return(5 * time.Second).Build()
		defer mockParam.UnPatch()

		sd := &shardDelegator{
			vchannelName:               "test-channel",
			latestTsafe:                atomic.NewUint64(0),
			catchingUpStreamingData:    atomic.NewBool(true),
			tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
			latestRequiredMVCCTimeTick: atomic.NewUint64(0),
		}

		// Initially catching up
		assert.True(t, sd.CatchingUpStreamingData())

		// Update tsafe with a recent timestamp (lag < 5s)
		recentTs := tsoutil.ComposeTSByTime(time.Now())
		sd.UpdateTSafe(recentTs)

		// Should now be caught up
		assert.False(t, sd.CatchingUpStreamingData())
	})

	t.Run("state remains catching up when lag is large", func(t *testing.T) {
		// Mock the config to return 5 seconds threshold
		mockParam := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsDurationByParse")).Return(5 * time.Second).Build()
		defer mockParam.UnPatch()

		sd := &shardDelegator{
			vchannelName:               "test-channel",
			latestTsafe:                atomic.NewUint64(0),
			catchingUpStreamingData:    atomic.NewBool(true),
			tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
			latestRequiredMVCCTimeTick: atomic.NewUint64(0),
		}

		// Initially catching up
		assert.True(t, sd.CatchingUpStreamingData())

		// Update tsafe with an old timestamp (lag > 5s)
		oldTs := tsoutil.ComposeTSByTime(time.Now().Add(-10 * time.Second))
		sd.UpdateTSafe(oldTs)

		// Should still be catching up
		assert.True(t, sd.CatchingUpStreamingData())
	})

	t.Run("threshold disabled when set to 0", func(t *testing.T) {
		// Mock the config to return 0 (disabled)
		mockParam := mockey.Mock(mockey.GetMethod(&paramtable.ParamItem{}, "GetAsDurationByParse")).Return(0 * time.Second).Build()
		defer mockParam.UnPatch()

		sd := &shardDelegator{
			vchannelName:               "test-channel",
			latestTsafe:                atomic.NewUint64(0),
			catchingUpStreamingData:    atomic.NewBool(true),
			tsCond:                     syncutil.NewContextCond(&sync.Mutex{}),
			latestRequiredMVCCTimeTick: atomic.NewUint64(0),
		}

		// Initially catching up
		assert.True(t, sd.CatchingUpStreamingData())

		// Update tsafe with a recent timestamp
		recentTs := tsoutil.ComposeTSByTime(time.Now())
		sd.UpdateTSafe(recentTs)

		// Should still be catching up (threshold disabled)
		assert.True(t, sd.CatchingUpStreamingData())
	})
}

// MinHash Function test
func (s *DelegatorSuite) TestDelegatorSearchWithMinHashFunction() {
	// miss parametres
	minHashFunctionSchema := &schemapb.FunctionSchema{
		Type:           schemapb.FunctionType_MinHash,
		InputFieldIds:  []int64{102},
		OutputFieldIds: []int64{101, 102}, // invalid output field
	}
	schema1 := &schemapb.CollectionSchema{
		Name: "TestCollection",
		Fields: []*schemapb.FieldSchema{
			{
				Name:         "id",
				FieldID:      100,
				IsPrimaryKey: true,
				DataType:     schemapb.DataType_Int64,
				AutoID:       true,
			}, {
				Name:         "binary_vector",
				FieldID:      101,
				IsPrimaryKey: false,
				DataType:     schemapb.DataType_BinaryVector,
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.DimKey,
						Value: "1024",
					},
				},
			}, {
				Name:     "text",
				FieldID:  102,
				DataType: schemapb.DataType_VarChar,
				TypeParams: []*commonpb.KeyValuePair{
					{
						Key:   common.MaxLengthKey,
						Value: "256",
					},
				},
			},
		},
		Functions: []*schemapb.FunctionSchema{minHashFunctionSchema},
	}

	s.Run("alloc function failed", func() {
		manager := segments.NewManager()
		manager.Collection.PutOrRef(s.collectionID, schema1, nil, &querypb.LoadMetaInfo{SchemaBarrierTs: tsoutil.ComposeTSByTime(time.Now())})

		delegator, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
		s.Require().NoError(err)
		defer delegator.Close()

		function.ReleaseFunctionRunners(s.collectionID, "WAL-"+s.vchannelName)
		s.NoError(function.AllocFunctionRunners(s.collectionID, "WAL-"+s.vchannelName, schema1))

		changed, err := function.FillFunctionData(context.Background(), s.collectionID, schema1, &msgpb.InsertRequest{
			FieldsData: []*schemapb.FieldData{{
				Type:    schemapb.DataType_VarChar,
				FieldId: 102,
				Field: &schemapb.FieldData_Scalars{
					Scalars: &schemapb.ScalarField{
						Data: &schemapb.ScalarField_StringData{
							StringData: &schemapb.StringArray{Data: []string{"test minhash data"}},
						},
					},
				},
			}},
			NumRows: 1,
		})
		s.False(changed)
		s.ErrorContains(err, "minhash function should only have one output field")
	})

	s.Run("init function ", func() {
		minHashFunctionSchema.OutputFieldIds = []int64{101}
		manager := segments.NewManager()
		manager.Collection.PutOrRef(s.collectionID, schema1, nil, &querypb.LoadMetaInfo{SchemaBarrierTs: tsoutil.ComposeTSByTime(time.Now())})

		delegator, err := NewShardDelegator(context.Background(), s.collectionID, s.replicaID, s.vchannelName, s.version, s.workerManager, manager, s.loader, 10000, nil, s.chunkManager, NewChannelQueryView(nil, nil, nil, initialTargetVersion), nil)
		s.NoError(err)
		defer delegator.Close()
	})
}