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

import (
	"context"
	"crypto/rand"
	"sync"
	"testing"
	"time"

	"github.com/cockroachdb/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"

	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
	"github.com/milvus-io/milvus/internal/datanode/allocator"
	"github.com/milvus-io/milvus/internal/datanode/broker"
	"github.com/milvus-io/milvus/internal/proto/datapb"
	"github.com/milvus-io/milvus/internal/proto/etcdpb"
	"github.com/milvus-io/milvus/internal/storage"
	"github.com/milvus-io/milvus/pkg/util/merr"
	"github.com/milvus-io/milvus/pkg/util/retry"
)

var flushTestDir = "/tmp/milvus_test/flush"

type emptyFlushTask struct{}

func (t *emptyFlushTask) flushInsertData() error {
	return nil
}

func (t *emptyFlushTask) flushDeleteData() error {
	return nil
}

type errFlushTask struct{}

func (t *errFlushTask) flushInsertData() error {
	return errors.New("mocked error")
}

func (t *errFlushTask) flushDeleteData() error {
	return errors.New("mocked error")
}

func TestOrderFlushQueue_Execute(t *testing.T) {
	counter := atomic.Int64{}
	finish := sync.WaitGroup{}

	size := 1000
	finish.Add(size)
	q := newOrderFlushQueue(1, "", func(*segmentFlushPack) {
		counter.Inc()
		finish.Done()
	})

	q.init()
	ids := make([][]byte, 0, size)
	for i := 0; i < size; i++ {
		id := make([]byte, 10)
		rand.Read(id)
		ids = append(ids, id)
	}

	wg := sync.WaitGroup{}
	wg.Add(2 * size)
	for i := 0; i < size; i++ {
		go func(id []byte) {
			q.enqueueDelFlush(&emptyFlushTask{}, &DelDataBuf{}, &msgpb.MsgPosition{
				MsgID: id,
			})
			wg.Done()
		}(ids[i])
		go func(id []byte) {
			q.enqueueInsertFlush(&emptyFlushTask{}, map[UniqueID]*datapb.Binlog{}, map[UniqueID]*datapb.Binlog{}, false, false, &msgpb.MsgPosition{
				MsgID: id,
			})
			wg.Done()
		}(ids[i])
	}
	wg.Wait()
	finish.Wait()

	assert.EqualValues(t, size, counter.Load())
}

func TestOrderFlushQueue_Order(t *testing.T) {
	counter := atomic.Int64{}
	finish := sync.WaitGroup{}

	size := 1000
	finish.Add(size)
	resultList := make([][]byte, 0, size)
	q := newOrderFlushQueue(1, "", func(pack *segmentFlushPack) {
		counter.Inc()
		resultList = append(resultList, pack.pos.MsgID)
		finish.Done()
	})

	q.init()
	ids := make([][]byte, 0, size)
	for i := 0; i < size; i++ {
		id := make([]byte, 10)
		rand.Read(id)
		ids = append(ids, id)
	}

	wg := sync.WaitGroup{}
	wg.Add(size)
	for i := 0; i < size; i++ {
		q.enqueueDelFlush(&emptyFlushTask{}, &DelDataBuf{}, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		q.enqueueInsertFlush(&emptyFlushTask{}, map[UniqueID]*datapb.Binlog{}, map[UniqueID]*datapb.Binlog{}, false, false, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		wg.Done()
	}
	wg.Wait()
	finish.Wait()

	assert.EqualValues(t, size, counter.Load())

	require.Equal(t, size, len(resultList))
	for i := 0; i < size; i++ {
		assert.EqualValues(t, ids[i], resultList[i])
	}
}

func newTestChannel() *ChannelMeta {
	return &ChannelMeta{
		segments:     make(map[UniqueID]*Segment),
		collectionID: 1,
		collSchema: &schemapb.CollectionSchema{
			Fields: []*schemapb.FieldSchema{{
				FieldID:      100,
				DataType:     schemapb.DataType_Int64,
				IsPrimaryKey: true,
			}},
		},
	}
}

func TestRendezvousFlushManager(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
	defer cm.RemoveWithPrefix(ctx, cm.RootPath())

	size := 1000
	var counter atomic.Int64
	alloc := allocator.NewMockAllocator(t)
	channel := newTestChannel()
	hisotricalStats := storage.NewPrimaryKeyStats(100, 5, 10)
	testSeg := &Segment{
		collectionID: 1,
		segmentID:    3,
		// flush segment will merge all historial stats
		historyStats: []*storage.PkStatistics{
			{
				PkFilter: hisotricalStats.BF,
				MinPK:    hisotricalStats.MinPk,
				MaxPK:    hisotricalStats.MaxPk,
			},
		},
	}
	testSeg.setType(datapb.SegmentType_New)
	channel.segments[testSeg.segmentID] = testSeg

	m := NewRendezvousFlushManager(alloc, cm, channel, func(pack *segmentFlushPack) {
		counter.Inc()
	}, emptyFlushAndDropFunc)

	ids := make([][]byte, 0, size)
	for i := 0; i < size; i++ {
		id := make([]byte, 10)
		rand.Read(id)
		ids = append(ids, id)
	}

	for i := 0; i < size; i++ {
		err := m.flushDelData(nil, testSeg.segmentID, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)

		_, err = m.flushBufferData(nil, testSeg.segmentID, true, false, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)
	}
	assert.Eventually(t, func() bool { return counter.Load() == int64(size) }, 3*time.Second, 100*time.Millisecond)

	_, _, err := m.serializePkStatsLog(0, false, nil, &storage.InsertCodec{Schema: &etcdpb.CollectionMeta{Schema: &schemapb.CollectionSchema{}, ID: 0}})
	assert.Error(t, err)
}

func TestRendezvousFlushManager_Inject(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
	defer cm.RemoveWithPrefix(ctx, cm.RootPath())

	size := 1000
	var counter atomic.Int64
	var packMut sync.Mutex
	packs := make([]*segmentFlushPack, 0, size+3)
	alloc := allocator.NewMockAllocator(t)

	channel := newTestChannel()
	segments := []*Segment{{
		collectionID: 1,
		segmentID:    1,
	}, {
		collectionID: 1,
		segmentID:    2,
	}, {
		collectionID: 1,
		segmentID:    3,
	}}

	for _, seg := range segments {
		seg.setType(datapb.SegmentType_New)
		channel.segments[seg.segmentID] = seg
	}

	m := NewRendezvousFlushManager(alloc, cm, channel, func(pack *segmentFlushPack) {
		packMut.Lock()
		packs = append(packs, pack)
		packMut.Unlock()
		counter.Inc()
	}, emptyFlushAndDropFunc)

	ti := newTaskInjection(1, func(*segmentFlushPack) {})
	m.injectFlush(ti, 1)
	<-ti.injected
	ti.injectDone(true)

	ids := make([][]byte, 0, size)
	for i := 0; i < size; i++ {
		id := make([]byte, 10)
		rand.Read(id)
		ids = append(ids, id)
	}

	for i := 0; i < size; i++ {
		err := m.flushDelData(nil, 1, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)

		_, err = m.flushBufferData(nil, 1, true, false, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)
	}
	assert.Eventually(t, func() bool { return counter.Load() == int64(size) }, 3*time.Second, 100*time.Millisecond)

	id := make([]byte, 10)
	rand.Read(id)
	id2 := make([]byte, 10)
	rand.Read(id2)
	_, err := m.flushBufferData(nil, 2, true, false, &msgpb.MsgPosition{
		MsgID: id,
	})

	assert.NoError(t, err)
	_, err = m.flushBufferData(nil, 3, true, false, &msgpb.MsgPosition{
		MsgID: id2,
	})
	assert.NoError(t, err)

	ti = newTaskInjection(2, func(pack *segmentFlushPack) {
		pack.segmentID = 4
	})
	m.injectFlush(ti, 2, 3)

	err = m.flushDelData(nil, 2, &msgpb.MsgPosition{
		MsgID: id,
	})
	assert.NoError(t, err)

	err = m.flushDelData(nil, 3, &msgpb.MsgPosition{
		MsgID: id2,
	})
	assert.NoError(t, err)

	<-ti.Injected()
	ti.injectDone(true)

	assert.Eventually(t, func() bool { return counter.Load() == int64(size+2) }, 3*time.Second, 100*time.Millisecond)
	assert.EqualValues(t, 4, packs[size].segmentID)

	rand.Read(id)

	_, err = m.flushBufferData(nil, 2, false, false, &msgpb.MsgPosition{
		MsgID: id,
	})
	assert.NoError(t, err)

	ti = newTaskInjection(1, func(pack *segmentFlushPack) {
		pack.segmentID = 5
	})
	go func() {
		<-ti.injected
		ti.injectDone(false) // inject fail, segment id shall not be changed to 5
	}()
	m.injectFlush(ti, 2)

	m.flushDelData(nil, 2, &msgpb.MsgPosition{
		MsgID: id,
	})
	assert.Eventually(t, func() bool { return counter.Load() == int64(size+3) }, 3*time.Second, 100*time.Millisecond)
	assert.EqualValues(t, 4, packs[size+1].segmentID)
}

func TestRendezvousFlushManager_getSegmentMeta(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
	defer cm.RemoveWithPrefix(ctx, cm.RootPath())

	channel := newTestChannel()
	channel.collSchema = &schemapb.CollectionSchema{}
	fm := NewRendezvousFlushManager(allocator.NewMockAllocator(t), cm, channel, func(*segmentFlushPack) {
	}, emptyFlushAndDropFunc)

	// non exists segment
	_, _, _, err := fm.getSegmentMeta(-1, &msgpb.MsgPosition{})
	assert.Error(t, err)

	seg0 := Segment{segmentID: -1}
	seg1 := Segment{segmentID: 1}
	seg0.setType(datapb.SegmentType_New)
	seg1.setType(datapb.SegmentType_New)

	channel.segments[-1] = &seg0
	channel.segments[1] = &seg1

	// // injected get part/coll id error
	// _, _, _, err = fm.getSegmentMeta(-1, &msgpb.MsgPosition{})
	// assert.Error(t, err)
	// // injected get schema  error
	// _, _, _, err = fm.getSegmentMeta(1, &msgpb.MsgPosition{})
	// assert.Error(t, err)
}

func TestRendezvousFlushManager_waitForAllFlushQueue(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
	defer cm.RemoveWithPrefix(ctx, cm.RootPath())

	channel := newTestChannel()
	testSeg := &Segment{
		collectionID: 1,
		segmentID:    1,
	}
	testSeg.setType(datapb.SegmentType_New)
	channel.segments[testSeg.segmentID] = testSeg

	size := 1000
	var counter atomic.Int64
	m := NewRendezvousFlushManager(allocator.NewMockAllocator(t), cm, channel, func(pack *segmentFlushPack) {
		counter.Inc()
	}, emptyFlushAndDropFunc)

	ids := make([][]byte, 0, size)
	for i := 0; i < size; i++ {
		id := make([]byte, 10)
		rand.Read(id)
		ids = append(ids, id)
	}

	for i := 0; i < size; i++ {
		err := m.flushDelData(nil, 1, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)
	}

	var finished bool
	var mut sync.RWMutex
	signal := make(chan struct{})

	go func() {
		m.waitForAllFlushQueue()
		mut.Lock()
		finished = true
		mut.Unlock()
		close(signal)
	}()

	mut.RLock()
	assert.False(t, finished)
	mut.RUnlock()

	for i := 0; i < size/2; i++ {
		_, err := m.flushBufferData(nil, 1, true, false, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)
	}

	mut.RLock()
	assert.False(t, finished)
	mut.RUnlock()

	for i := size / 2; i < size; i++ {
		_, err := m.flushBufferData(nil, 1, true, false, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)
	}

	select {
	case <-time.After(time.Second):
		t.FailNow()
	case <-signal:
	}

	mut.RLock()
	assert.True(t, finished)
	mut.RUnlock()
}

func TestRendezvousFlushManager_dropMode(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	t.Run("test drop mode", func(t *testing.T) {
		cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
		defer cm.RemoveWithPrefix(ctx, cm.RootPath())

		var mut sync.Mutex
		var result []*segmentFlushPack
		signal := make(chan struct{})

		channel := newTestChannel()
		targets := make(map[int64]struct{})
		// init failed segment
		testSeg := &Segment{
			collectionID: 1,
			segmentID:    -1,
		}
		testSeg.setType(datapb.SegmentType_New)
		channel.segments[testSeg.segmentID] = testSeg

		// init target segment
		for i := 1; i < 11; i++ {
			targets[int64(i)] = struct{}{}
			testSeg := &Segment{
				collectionID: 1,
				segmentID:    int64(i),
			}
			testSeg.setType(datapb.SegmentType_New)
			channel.segments[testSeg.segmentID] = testSeg
		}

		// init flush manager
		m := NewRendezvousFlushManager(allocator.NewMockAllocator(t), cm, channel, func(pack *segmentFlushPack) {
		}, func(packs []*segmentFlushPack) {
			mut.Lock()
			result = packs
			mut.Unlock()
			close(signal)
		})

		halfMsgID := []byte{1, 1, 1}
		_, err := m.flushBufferData(nil, -1, true, false, &msgpb.MsgPosition{
			MsgID: halfMsgID,
		})
		assert.NoError(t, err)

		m.startDropping()
		// half normal, half drop mode, should not appear in final packs
		err = m.flushDelData(nil, -1, &msgpb.MsgPosition{
			MsgID: halfMsgID,
		})
		assert.NoError(t, err)

		for target := range targets {
			_, err := m.flushBufferData(nil, target, true, false, &msgpb.MsgPosition{
				MsgID: []byte{byte(target)},
			})
			assert.NoError(t, err)

			err = m.flushDelData(nil, target, &msgpb.MsgPosition{
				MsgID: []byte{byte(target)},
			})
			assert.NoError(t, err)
		}

		m.notifyAllFlushed()

		<-signal
		mut.Lock()
		defer mut.Unlock()

		output := make(map[int64]struct{})
		for _, pack := range result {
			assert.NotEqual(t, -1, pack.segmentID)
			output[pack.segmentID] = struct{}{}
			_, has := targets[pack.segmentID]
			assert.True(t, has)
		}
		assert.Equal(t, len(targets), len(output))
	})

	t.Run("test drop mode with injection", func(t *testing.T) {
		cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
		defer cm.RemoveWithPrefix(ctx, cm.RootPath())

		var mut sync.Mutex
		var result []*segmentFlushPack
		signal := make(chan struct{})
		channel := newTestChannel()
		// init failed segment
		testSeg := &Segment{
			collectionID: 1,
			segmentID:    -1,
		}
		testSeg.setType(datapb.SegmentType_New)
		channel.segments[testSeg.segmentID] = testSeg

		// init target segment
		for i := 1; i < 11; i++ {
			seg := &Segment{
				collectionID: 1,
				segmentID:    int64(i),
			}
			seg.setType(datapb.SegmentType_New)
			channel.segments[seg.segmentID] = seg
		}

		m := NewRendezvousFlushManager(allocator.NewMockAllocator(t), cm, channel, func(pack *segmentFlushPack) {
		}, func(packs []*segmentFlushPack) {
			mut.Lock()
			result = packs
			mut.Unlock()
			close(signal)
		})

		// flush failed segment before start drop mode
		halfMsgID := []byte{1, 1, 1}
		_, err := m.flushBufferData(nil, -1, true, false, &msgpb.MsgPosition{
			MsgID: halfMsgID,
		})
		assert.NoError(t, err)

		// inject target segment
		injFunc := func(pack *segmentFlushPack) {
			pack.segmentID = 100
		}
		for i := 1; i < 11; i++ {
			it := newTaskInjection(1, injFunc)
			m.injectFlush(it, int64(i))
			<-it.Injected()
			it.injectDone(true)
		}

		m.startDropping()
		// half normal, half drop mode, should not appear in final packs
		err = m.flushDelData(nil, -1, &msgpb.MsgPosition{
			MsgID: halfMsgID,
		})
		assert.NoError(t, err)

		for i := 1; i < 11; i++ {
			_, err = m.flushBufferData(nil, int64(i), true, false, &msgpb.MsgPosition{
				MsgID: []byte{byte(i)},
			})
			assert.NoError(t, err)

			err = m.flushDelData(nil, int64(i), &msgpb.MsgPosition{
				MsgID: []byte{byte(i)},
			})
			assert.NoError(t, err)
		}

		m.notifyAllFlushed()
		mut.Lock()
		defer mut.Unlock()

		for _, pack := range result {
			assert.NotEqual(t, -1, pack.segmentID)
			assert.Equal(t, int64(100), pack.segmentID)
		}
	})
}

func TestRendezvousFlushManager_close(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
	defer cm.RemoveWithPrefix(ctx, cm.RootPath())

	size := 1000
	var counter atomic.Int64

	channel := newTestChannel()

	// init test segment
	testSeg := &Segment{
		collectionID: 1,
		segmentID:    1,
	}
	testSeg.setType(datapb.SegmentType_New)
	channel.segments[testSeg.segmentID] = testSeg

	m := NewRendezvousFlushManager(allocator.NewMockAllocator(t), cm, channel, func(pack *segmentFlushPack) {
		counter.Inc()
	}, emptyFlushAndDropFunc)

	ids := make([][]byte, 0, size)
	for i := 0; i < size; i++ {
		id := make([]byte, 10)
		rand.Read(id)
		ids = append(ids, id)
	}

	for i := 0; i < size; i++ {
		err := m.flushDelData(nil, 1, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)

		_, err = m.flushBufferData(nil, 1, true, false, &msgpb.MsgPosition{
			MsgID: ids[i],
		})
		assert.NoError(t, err)
	}

	m.close()
	assert.Eventually(t, func() bool { return counter.Load() == int64(size) }, 3*time.Second, 100*time.Millisecond)
}

func TestFlushNotifyFunc(t *testing.T) {
	meta := NewMetaFactory().GetCollectionMeta(1, "testCollection", schemapb.DataType_Int64)
	broker := broker.NewMockBroker(t)
	broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything, mock.Anything).
		Return(&milvuspb.DescribeCollectionResponse{
			Status: merr.Status(nil),
			Schema: meta.GetSchema(),
		}, nil).Maybe()
	broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).Return(nil).Maybe()
	broker.EXPECT().DropVirtualChannel(mock.Anything, mock.Anything).Return(nil, nil).Maybe()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
	defer cm.RemoveWithPrefix(ctx, cm.RootPath())

	channel := newChannel("channel", 1, nil, broker, cm)

	flushingCache := newCache()
	dsService := &dataSyncService{
		collectionID:     1,
		channel:          channel,
		broker:           broker,
		flushingSegCache: flushingCache,
	}
	notifyFunc := flushNotifyFunc(dsService, retry.Attempts(1))

	t.Run("normal run", func(t *testing.T) {
		assert.NotPanics(t, func() {
			notifyFunc(&segmentFlushPack{
				insertLogs: map[UniqueID]*datapb.Binlog{1: {LogPath: "/dev/test/id"}},
				statsLogs:  map[UniqueID]*datapb.Binlog{1: {LogPath: "/dev/test/id-stats"}},
				deltaLogs:  []*datapb.Binlog{{LogPath: "/dev/test/del"}},
				flushed:    true,
			})
		})
	})

	t.Run("pack has error", func(t *testing.T) {
		assert.Panics(t, func() {
			notifyFunc(&segmentFlushPack{
				err: errors.New("mocked pack error"),
			})
		})
	})

	t.Run("datacoord save fails", func(t *testing.T) {
		broker.ExpectedCalls = nil
		broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).
			Return(merr.WrapErrCollectionNotFound("test_collection"))
		assert.Panics(t, func() {
			notifyFunc(&segmentFlushPack{})
		})
	})

	t.Run("stale segment not found", func(t *testing.T) {
		broker.ExpectedCalls = nil
		broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).
			Return(merr.WrapErrSegmentNotFound(0))
		assert.NotPanics(t, func() {
			notifyFunc(&segmentFlushPack{flushed: false})
		})
	})

	// issue https://github.com/milvus-io/milvus/issues/17097
	// meta error, datanode shall not panic, just drop the virtual channel
	t.Run("datacoord found meta error", func(t *testing.T) {
		broker.ExpectedCalls = nil
		broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).
			Return(merr.WrapErrChannelNotFound("channel"))
		assert.NotPanics(t, func() {
			notifyFunc(&segmentFlushPack{})
		})
	})

	t.Run("datacoord call error", func(t *testing.T) {
		broker.ExpectedCalls = nil
		broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).
			Return(errors.New("mock"))
		assert.Panics(t, func() {
			notifyFunc(&segmentFlushPack{})
		})
	})
}

func TestDropVirtualChannelFunc(t *testing.T) {
	meta := NewMetaFactory().GetCollectionMeta(1, "testCollection", schemapb.DataType_Int64)
	broker := broker.NewMockBroker(t)
	broker.EXPECT().DescribeCollection(mock.Anything, mock.Anything, mock.Anything).
		Return(&milvuspb.DescribeCollectionResponse{
			Status: merr.Status(nil),
			Schema: meta.GetSchema(),
		}, nil)
	broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).Return(nil).Maybe()
	broker.EXPECT().DropVirtualChannel(mock.Anything, mock.Anything).Return(nil, nil).Maybe()
	vchanName := "vchan_01"

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cm := storage.NewLocalChunkManager(storage.RootPath(flushTestDir))
	defer cm.RemoveWithPrefix(ctx, cm.RootPath())

	channel := newChannel(vchanName, 1, nil, broker, cm)

	flushingCache := newCache()
	dsService := &dataSyncService{
		collectionID:     1,
		channel:          channel,
		broker:           broker,
		flushingSegCache: flushingCache,
		vchannelName:     vchanName,
	}
	dropFunc := dropVirtualChannelFunc(dsService, retry.Attempts(1))
	t.Run("normal run", func(t *testing.T) {
		channel.addSegment(
			context.TODO(),
			addSegmentReq{
				segType:     datapb.SegmentType_New,
				segID:       2,
				collID:      1,
				partitionID: 10,
				startPos: &msgpb.MsgPosition{
					ChannelName: vchanName,
					MsgID:       []byte{1, 2, 3},
					Timestamp:   10,
				}, endPos: nil,
			})
		assert.NotPanics(t, func() {
			dropFunc([]*segmentFlushPack{
				{
					segmentID:  1,
					insertLogs: map[UniqueID]*datapb.Binlog{1: {LogPath: "/dev/test/id"}},
					statsLogs:  map[UniqueID]*datapb.Binlog{1: {LogPath: "/dev/test/id-stats"}},
					deltaLogs:  []*datapb.Binlog{{LogPath: "/dev/test/del"}},
					pos: &msgpb.MsgPosition{
						ChannelName: vchanName,
						MsgID:       []byte{1, 2, 3},
						Timestamp:   10,
					},
				},
				{
					segmentID:  1,
					insertLogs: map[UniqueID]*datapb.Binlog{1: {LogPath: "/dev/test/idi_2"}},
					statsLogs:  map[UniqueID]*datapb.Binlog{1: {LogPath: "/dev/test/id-stats-2"}},
					deltaLogs:  []*datapb.Binlog{{LogPath: "/dev/test/del-2"}},
					pos: &msgpb.MsgPosition{
						ChannelName: vchanName,
						MsgID:       []byte{1, 2, 3},
						Timestamp:   30,
					},
				},
			})
		})
	})
	t.Run("datacoord_return_error", func(t *testing.T) {
		broker.ExpectedCalls = nil
		broker.EXPECT().DropVirtualChannel(mock.Anything, mock.Anything).
			Return(nil, errors.New("mock"))
		assert.Panics(t, func() {
			dropFunc(nil)
		})
	})

	// issue https://github.com/milvus-io/milvus/issues/17097
	// meta error, datanode shall not panic, just drop the virtual channel
	t.Run("datacoord_return_channel_not_found", func(t *testing.T) {
		broker.ExpectedCalls = nil
		broker.EXPECT().DropVirtualChannel(mock.Anything, mock.Anything).
			Return(nil, merr.WrapErrChannelNotFound("channel"))
		assert.NotPanics(t, func() {
			dropFunc(nil)
		})
	})
}