milvus/internal/proxy/task_database_test.go

166 lines
4.1 KiB
Go
Raw Normal View History

package proxy
import (
"context"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"google.golang.org/grpc/metadata"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/pkg/util"
"github.com/milvus-io/milvus/pkg/util/crypto"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
func TestCreateDatabaseTask(t *testing.T) {
paramtable.Init()
rc := NewRootCoordMock()
defer rc.Close()
ctx := context.Background()
task := &createDatabaseTask{
Condition: NewTaskCondition(ctx),
CreateDatabaseRequest: &milvuspb.CreateDatabaseRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_CreateDatabase,
MsgID: 100,
Timestamp: 100,
},
DbName: "db",
},
ctx: ctx,
rootCoord: rc,
result: nil,
}
t.Run("ok", func(t *testing.T) {
err := task.PreExecute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.MsgType_CreateDatabase, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
assert.Equal(t, "db", task.GetDbName())
err = task.Execute(ctx)
assert.NoError(t, err)
task.Base = nil
err = task.OnEnqueue()
assert.NoError(t, err)
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
assert.Equal(t, UniqueID(0), task.ID())
})
t.Run("pre execute fail", func(t *testing.T) {
task.DbName = "#0xc0de"
err := task.PreExecute(ctx)
assert.Error(t, err)
})
}
func TestDropDatabaseTask(t *testing.T) {
paramtable.Init()
rc := NewRootCoordMock()
defer rc.Close()
ctx := context.Background()
task := &dropDatabaseTask{
Condition: NewTaskCondition(ctx),
DropDatabaseRequest: &milvuspb.DropDatabaseRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_DropDatabase,
MsgID: 100,
Timestamp: 100,
},
DbName: "db",
},
ctx: ctx,
rootCoord: rc,
result: nil,
}
cache := NewMockCache(t)
cache.On("RemoveDatabase",
mock.Anything, // context.Context
mock.AnythingOfType("string"),
).Maybe()
globalMetaCache = cache
t.Run("ok", func(t *testing.T) {
err := task.PreExecute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.MsgType_DropDatabase, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
assert.Equal(t, "db", task.GetDbName())
err = task.Execute(ctx)
assert.NoError(t, err)
task.Base = nil
err = task.OnEnqueue()
assert.NoError(t, err)
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
assert.Equal(t, UniqueID(0), task.ID())
})
t.Run("pre execute fail", func(t *testing.T) {
task.DbName = "#0xc0de"
err := task.PreExecute(ctx)
assert.Error(t, err)
})
}
func TestListDatabaseTask(t *testing.T) {
paramtable.Init()
rc := NewRootCoordMock()
defer rc.Close()
ctx := GetContext(context.Background(), "root:123456")
task := &listDatabaseTask{
Condition: NewTaskCondition(ctx),
ListDatabasesRequest: &milvuspb.ListDatabasesRequest{
Base: &commonpb.MsgBase{
MsgType: commonpb.MsgType_ListDatabases,
MsgID: 100,
Timestamp: 100,
},
},
ctx: ctx,
rootCoord: rc,
result: nil,
}
t.Run("ok", func(t *testing.T) {
err := task.PreExecute(ctx)
assert.NoError(t, err)
assert.Equal(t, commonpb.MsgType_ListDatabases, task.Type())
assert.Equal(t, UniqueID(100), task.ID())
assert.Equal(t, Timestamp(100), task.BeginTs())
assert.Equal(t, Timestamp(100), task.EndTs())
err = task.Execute(ctx)
assert.NoError(t, err)
assert.NotNil(t, task.result)
err = task.OnEnqueue()
assert.NoError(t, err)
assert.Equal(t, paramtable.GetNodeID(), task.GetBase().GetSourceID())
assert.Equal(t, UniqueID(0), task.ID())
taskCtx := AppendUserInfoForRPC(ctx)
md, ok := metadata.FromOutgoingContext(taskCtx)
assert.True(t, ok)
authorization, ok := md[strings.ToLower(util.HeaderAuthorize)]
assert.True(t, ok)
expectAuth := crypto.Base64Encode("root:root")
assert.Equal(t, expectAuth, authorization[0])
})
}