milvus/internal/datanode/collection_replica_test.go

package datanode

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
)

func initTestReplicaMeta(t *testing.T, replica collectionReplica, collectionName string, collectionID UniqueID, segmentID UniqueID) {
	// GOOSE TODO remove
	Factory := &MetaFactory{}
	collectionMeta := Factory.CollectionMetaFactory(collectionID, collectionName)

	var err = replica.addCollection(collectionMeta.ID, collectionMeta.Schema)
	require.NoError(t, err)

	collection, err := replica.getCollectionByName(collectionName)
	require.NoError(t, err)
	assert.Equal(t, collection.Name(), collectionName)
	assert.Equal(t, collection.ID(), collectionID)
	assert.Equal(t, replica.getCollectionNum(), 1)

}

func TestReplica_Collection(t *testing.T) {
	Factory := &MetaFactory{}
	collMetaMock := Factory.CollectionMetaFactory(0, "collection0")

	t.Run("Test add collection", func(t *testing.T) {

		replica := newReplica()
		assert.False(t, replica.hasCollection(0))
		num := replica.getCollectionNum()
		assert.Equal(t, 0, num)

		err := replica.addCollection(0, collMetaMock.GetSchema())
		assert.NoError(t, err)

		assert.True(t, replica.hasCollection(0))
		num = replica.getCollectionNum()
		assert.Equal(t, 1, num)

		coll, err := replica.getCollectionByID(0)
		assert.NoError(t, err)
		assert.NotNil(t, coll)
		assert.Equal(t, UniqueID(0), coll.ID())
		assert.Equal(t, "collection0", coll.Name())
		assert.Equal(t, collMetaMock.GetSchema(), coll.Schema())

		coll, err = replica.getCollectionByName("collection0")
		assert.NoError(t, err)
		assert.NotNil(t, coll)
		assert.Equal(t, UniqueID(0), coll.ID())
		assert.Equal(t, "collection0", coll.Name())
		assert.Equal(t, collMetaMock.GetSchema(), coll.Schema())

		collID, err := replica.getCollectionIDByName("collection0")
		assert.NoError(t, err)
		assert.Equal(t, UniqueID(0), collID)

	})

	t.Run("Test remove collection", func(t *testing.T) {
		replica := newReplica()
		err := replica.addCollection(0, collMetaMock.GetSchema())
		require.NoError(t, err)

		numsBefore := replica.getCollectionNum()
		coll, err := replica.getCollectionByID(0)
		require.NotNil(t, coll)
		require.NoError(t, err)

		err = replica.removeCollection(0)
		assert.NoError(t, err)
		numsAfter := replica.getCollectionNum()
		assert.Equal(t, 1, numsBefore-numsAfter)

		coll, err = replica.getCollectionByID(0)
		assert.Nil(t, coll)
		assert.Error(t, err)
		err = replica.removeCollection(999999999)
		assert.Error(t, err)
	})

	t.Run("Test errors", func(t *testing.T) {
		replica := newReplica()
		require.False(t, replica.hasCollection(0))
		require.Equal(t, 0, replica.getCollectionNum())

		coll, err := replica.getCollectionByName("Name-not-exist")
		assert.Error(t, err)
		assert.Nil(t, coll)

		coll, err = replica.getCollectionByID(0)
		assert.Error(t, err)
		assert.Nil(t, coll)

		collID, err := replica.getCollectionIDByName("Name-not-exist")
		assert.Error(t, err)
		assert.Zero(t, collID)

		err = replica.removeCollection(0)
		assert.Error(t, err)
	})

}

func TestReplica_Segment(t *testing.T) {
	t.Run("Test segment", func(t *testing.T) {
		replica := newReplica()
		assert.False(t, replica.hasSegment(0))

		err := replica.addSegment(0, 1, 2, make([]*internalpb2.MsgPosition, 0))
		assert.NoError(t, err)
		assert.True(t, replica.hasSegment(0))

		seg, err := replica.getSegmentByID(0)
		assert.NoError(t, err)
		assert.NotNil(t, seg)
		assert.Equal(t, UniqueID(1), seg.collectionID)
		assert.Equal(t, UniqueID(2), seg.partitionID)

		assert.Equal(t, int64(0), seg.numRows)

		err = replica.updateStatistics(0, 100)
		assert.NoError(t, err)
		assert.Equal(t, int64(100), seg.numRows)

		update, err := replica.getSegmentStatisticsUpdates(0)
		assert.NoError(t, err)
		assert.Equal(t, UniqueID(0), update.SegmentID)
		assert.Equal(t, int64(100), update.NumRows)
		assert.True(t, update.IsNewSegment)
	})

	t.Run("Test errors", func(t *testing.T) {
		replica := newReplica()
		require.False(t, replica.hasSegment(0))

		seg, err := replica.getSegmentByID(0)
		assert.Error(t, err)
		assert.Nil(t, seg)

		err = replica.removeSegment(0)
		assert.Error(t, err)

		err = replica.updateStatistics(0, 0)
		assert.Error(t, err)

		update, err := replica.getSegmentStatisticsUpdates(0)
		assert.Error(t, err)
		assert.Nil(t, update)
	})
}