2023-06-25 17:20:43 +08:00
|
|
|
package proxy
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
2024-03-28 07:13:10 +08:00
|
|
|
"strings"
|
2023-06-25 17:20:43 +08:00
|
|
|
"testing"
|
|
|
|
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/stretchr/testify/mock"
|
2024-03-28 07:13:10 +08:00
|
|
|
"google.golang.org/grpc/metadata"
|
2023-06-25 17:20:43 +08:00
|
|
|
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
2024-03-28 07:13:10 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util"
|
|
|
|
"github.com/milvus-io/milvus/pkg/util/crypto"
|
2023-09-21 09:45:27 +08:00
|
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
2023-06-25 17:20:43 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
func TestCreateDatabaseTask(t *testing.T) {
|
|
|
|
paramtable.Init()
|
|
|
|
rc := NewRootCoordMock()
|
2023-09-26 09:57:25 +08:00
|
|
|
defer rc.Close()
|
2023-06-25 17:20:43 +08:00
|
|
|
|
|
|
|
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)
|
|
|
|
|
2023-10-20 14:26:09 +08:00
|
|
|
task.Base = nil
|
2023-06-25 17:20:43 +08:00
|
|
|
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()
|
2023-09-26 09:57:25 +08:00
|
|
|
defer rc.Close()
|
2023-06-25 17:20:43 +08:00
|
|
|
|
|
|
|
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)
|
|
|
|
|
2023-10-20 14:26:09 +08:00
|
|
|
task.Base = nil
|
2023-06-25 17:20:43 +08:00
|
|
|
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()
|
2023-09-26 09:57:25 +08:00
|
|
|
defer rc.Close()
|
2023-06-25 17:20:43 +08:00
|
|
|
|
2024-03-28 07:13:10 +08:00
|
|
|
ctx := GetContext(context.Background(), "root:123456")
|
2023-06-25 17:20:43 +08:00
|
|
|
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())
|
2024-03-28 07:13:10 +08:00
|
|
|
|
2024-04-29 19:09:34 +08:00
|
|
|
taskCtx := AppendUserInfoForRPC(ctx)
|
|
|
|
md, ok := metadata.FromOutgoingContext(taskCtx)
|
2024-03-28 07:13:10 +08:00
|
|
|
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])
|
2023-06-25 17:20:43 +08:00
|
|
|
})
|
|
|
|
}
|