mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-02 11:59:00 +08:00
112 lines
3.0 KiB
Go
112 lines
3.0 KiB
Go
package master
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
|
|
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/etcdpb"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/schemapb"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
etcdkv "github.com/zilliztech/milvus-distributed/internal/kv/etcd"
|
|
"github.com/zilliztech/milvus-distributed/internal/proto/indexbuilderpb"
|
|
"go.etcd.io/etcd/clientv3"
|
|
)
|
|
|
|
func TestPersistenceScheduler(t *testing.T) {
|
|
//Init environment Params
|
|
Init()
|
|
|
|
ctx := context.Background()
|
|
|
|
//Init client, use Mock instead
|
|
flushClient := &MockWriteNodeClient{}
|
|
buildIndexClient := &MockBuildIndexClient{}
|
|
loadIndexClient := &MockLoadIndexClient{}
|
|
|
|
etcdAddr := Params.EtcdAddress
|
|
cli, err := clientv3.New(clientv3.Config{Endpoints: []string{etcdAddr}})
|
|
assert.Nil(t, err)
|
|
etcdKV := etcdkv.NewEtcdKV(cli, "/etcd/test/root")
|
|
|
|
_, err = cli.Delete(context.TODO(), "/etcd/test/root", clientv3.WithPrefix())
|
|
assert.Nil(t, err)
|
|
|
|
meta, err := NewMetaTable(etcdKV)
|
|
assert.Nil(t, err)
|
|
defer meta.client.Close()
|
|
|
|
err = meta.AddCollection(&etcdpb.CollectionMeta{
|
|
ID: 1,
|
|
Schema: &schemapb.CollectionSchema{
|
|
Name: "testcoll",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: 1},
|
|
{FieldID: 100, DataType: schemapb.DataType_VECTOR_FLOAT, IndexParams: []*commonpb.KeyValuePair{{Key: "k", Value: "v"}}},
|
|
},
|
|
},
|
|
})
|
|
assert.Nil(t, err)
|
|
err = meta.AddSegment(&etcdpb.SegmentMeta{
|
|
SegmentID: 1,
|
|
CollectionID: 1,
|
|
})
|
|
assert.Nil(t, err)
|
|
|
|
//Init scheduler
|
|
indexLoadSch := NewIndexLoadScheduler(ctx, loadIndexClient, meta)
|
|
indexBuildSch := NewIndexBuildScheduler(ctx, buildIndexClient, meta, indexLoadSch)
|
|
cnt := 0
|
|
flushSch := NewFlushScheduler(ctx, flushClient, meta, indexBuildSch, func() (Timestamp, error) {
|
|
cnt++
|
|
return Timestamp(cnt), nil
|
|
})
|
|
|
|
//scheduler start
|
|
err = indexLoadSch.Start()
|
|
assert.Nil(t, err)
|
|
defer indexLoadSch.Close()
|
|
|
|
err = indexBuildSch.Start()
|
|
assert.Nil(t, err)
|
|
defer indexBuildSch.Close()
|
|
|
|
err = flushSch.Start()
|
|
assert.Nil(t, err)
|
|
defer flushSch.Close()
|
|
|
|
//start from flush scheduler
|
|
err = flushSch.Enqueue(UniqueID(1))
|
|
assert.Nil(t, err)
|
|
//wait flush segment request sent to write node
|
|
time.Sleep(100 * time.Millisecond)
|
|
segDes, err := flushClient.DescribeSegment(UniqueID(1))
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, false, segDes.IsClosed)
|
|
|
|
//wait flush to finish
|
|
time.Sleep(3 * time.Second)
|
|
|
|
segDes, err = flushClient.DescribeSegment(UniqueID(1))
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, UniqueID(1), segDes.SegmentID)
|
|
assert.Equal(t, true, segDes.IsClosed)
|
|
|
|
//wait flush segment request sent to build index node
|
|
time.Sleep(100 * time.Microsecond)
|
|
idxDes, err := buildIndexClient.DescribeIndex(UniqueID(1))
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, indexbuilderpb.IndexStatus_INPROGRESS, idxDes.Status)
|
|
|
|
//wait build index to finish
|
|
time.Sleep(3 * time.Second)
|
|
|
|
idxDes, err = buildIndexClient.DescribeIndex(UniqueID(1))
|
|
assert.Nil(t, err)
|
|
assert.Equal(t, indexbuilderpb.IndexStatus_FINISHED, idxDes.Status)
|
|
|
|
}
|