milvus/internal/datacoord/compaction_trigger_v2_test.go

90 lines
2.6 KiB
Go
Raw Normal View History

package datacoord
import (
"testing"
"github.com/pingcap/log"
"github.com/samber/lo"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"go.uber.org/zap"
"github.com/milvus-io/milvus/internal/proto/datapb"
)
func TestCompactionTriggerManagerSuite(t *testing.T) {
suite.Run(t, new(CompactionTriggerManagerSuite))
}
type CompactionTriggerManagerSuite struct {
suite.Suite
mockAlloc *NMockAllocator
mockPlanContext *MockCompactionPlanContext
testLabel *CompactionGroupLabel
m *CompactionTriggerManager
}
func (s *CompactionTriggerManagerSuite) SetupTest() {
s.mockAlloc = NewNMockAllocator(s.T())
s.mockPlanContext = NewMockCompactionPlanContext(s.T())
s.testLabel = &CompactionGroupLabel{
CollectionID: 1,
PartitionID: 10,
Channel: "ch-1",
}
meta := &meta{segments: &SegmentsInfo{
segments: genSegmentsForMeta(s.testLabel),
}}
s.m = NewCompactionTriggerManager(meta, s.mockAlloc, s.mockPlanContext)
}
func (s *CompactionTriggerManagerSuite) TestNotify() {
viewManager := NewCompactionViewManager(s.m.meta, s.m, s.m.allocator)
collSegs := s.m.meta.GetCompactableSegmentGroupByCollection()
segments, found := collSegs[1]
s.Require().True(found)
levelZeroSegments := lo.Filter(segments, func(info *SegmentInfo, _ int) bool {
return info.GetLevel() == datapb.SegmentLevel_L0
})
latestL0Segments := GetViewsByInfo(levelZeroSegments...)
s.Require().NotEmpty(latestL0Segments)
levelZeroView := viewManager.getChangedLevelZeroViews(1, latestL0Segments)
s.Require().Equal(1, len(levelZeroView))
cView, ok := levelZeroView[0].(*LevelZeroSegmentsView)
s.True(ok)
s.NotNil(cView)
log.Info("view", zap.Any("cView", cView))
s.mockAlloc.EXPECT().allocID(mock.Anything).Return(1, nil)
s.mockPlanContext.EXPECT().execCompactionPlan(mock.Anything, mock.Anything).
Run(func(signal *compactionSignal, plan *datapb.CompactionPlan) {
s.EqualValues(19530, signal.id)
s.True(signal.isGlobal)
s.False(signal.isForce)
s.EqualValues(30000, signal.pos.GetTimestamp())
s.Equal(s.testLabel.CollectionID, signal.collectionID)
s.Equal(s.testLabel.PartitionID, signal.partitionID)
s.NotNil(plan)
s.Equal(s.testLabel.Channel, plan.GetChannel())
s.Equal(datapb.CompactionType_Level0DeleteCompaction, plan.GetType())
expectedSegs := []int64{100, 101, 102}
gotSegs := lo.Map(plan.GetSegmentBinlogs(), func(b *datapb.CompactionSegmentBinlogs, _ int) int64 {
return b.GetSegmentID()
})
s.ElementsMatch(expectedSegs, gotSegs)
log.Info("generated plan", zap.Any("plan", plan))
}).Return(nil).Once()
s.m.Notify(19530, TriggerTypeLevelZeroView, levelZeroView)
}