mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-12 13:15:09 +08:00
675 lines
19 KiB
Go
675 lines
19 KiB
Go
package proxy
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"testing"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/mock"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
|
"github.com/milvus-io/milvus/internal/allocator"
|
|
"github.com/milvus-io/milvus/internal/mocks"
|
|
"github.com/milvus-io/milvus/internal/parser/planparserv2"
|
|
"github.com/milvus-io/milvus/internal/proto/internalpb"
|
|
"github.com/milvus-io/milvus/internal/proto/planpb"
|
|
"github.com/milvus-io/milvus/internal/proto/querypb"
|
|
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
|
|
"github.com/milvus-io/milvus/internal/util/streamrpc"
|
|
"github.com/milvus-io/milvus/pkg/common"
|
|
"github.com/milvus-io/milvus/pkg/mq/msgstream"
|
|
"github.com/milvus-io/milvus/pkg/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
)
|
|
|
|
func Test_GetExpr(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_delete",
|
|
Description: "",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.StartOfUserFieldID,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: common.StartOfUserFieldID + 1,
|
|
Name: "non_pk",
|
|
IsPrimaryKey: false,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
},
|
|
}
|
|
t.Run("delelte with complex pk expr", func(t *testing.T) {
|
|
expr := "pk < 4"
|
|
plan, err := planparserv2.CreateRetrievePlan(schema, expr)
|
|
assert.NoError(t, err)
|
|
isSimple, _ := getExpr(plan)
|
|
assert.False(t, isSimple)
|
|
})
|
|
|
|
t.Run("delete with no-pk field expr", func(t *testing.T) {
|
|
expr := "non_pk in [1, 2, 3]"
|
|
plan, err := planparserv2.CreateRetrievePlan(schema, expr)
|
|
assert.NoError(t, err)
|
|
isSimple, _ := getExpr(plan)
|
|
assert.False(t, isSimple)
|
|
})
|
|
|
|
t.Run("delete with simple expr", func(t *testing.T) {
|
|
expr := "pk in [1, 2, 3]"
|
|
plan, err := planparserv2.CreateRetrievePlan(schema, expr)
|
|
assert.NoError(t, err)
|
|
isSimple, _ := getExpr(plan)
|
|
assert.True(t, isSimple)
|
|
})
|
|
}
|
|
|
|
func TestDeleteTask_GetChannels(t *testing.T) {
|
|
collectionID := UniqueID(0)
|
|
collectionName := "col-0"
|
|
channels := []pChan{"mock-chan-0", "mock-chan-1"}
|
|
cache := NewMockCache(t)
|
|
cache.On("GetCollectionID",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(collectionID, nil)
|
|
globalMetaCache = cache
|
|
chMgr := newMockChannelsMgr()
|
|
chMgr.getChannelsFunc = func(collectionID UniqueID) ([]pChan, error) {
|
|
return channels, nil
|
|
}
|
|
dt := deleteTask{
|
|
ctx: context.Background(),
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: collectionName,
|
|
},
|
|
chMgr: chMgr,
|
|
}
|
|
err := dt.setChannels()
|
|
assert.NoError(t, err)
|
|
resChannels := dt.getChannels()
|
|
assert.ElementsMatch(t, channels, resChannels)
|
|
assert.ElementsMatch(t, channels, dt.pChannels)
|
|
|
|
chMgr.getChannelsFunc = func(collectionID UniqueID) ([]pChan, error) {
|
|
return nil, fmt.Errorf("mock err")
|
|
}
|
|
// get channels again, should return task's pChannels, so getChannelsFunc should not invoke again
|
|
resChannels = dt.getChannels()
|
|
assert.ElementsMatch(t, channels, resChannels)
|
|
}
|
|
|
|
func TestDeleteTask_PreExecute(t *testing.T) {
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: "test_delete",
|
|
Description: "",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.StartOfUserFieldID,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: common.StartOfUserFieldID + 1,
|
|
Name: "non_pk",
|
|
IsPrimaryKey: false,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
},
|
|
}
|
|
|
|
t.Run("empty collection name", func(t *testing.T) {
|
|
dt := deleteTask{}
|
|
assert.Error(t, dt.PreExecute(context.Background()))
|
|
})
|
|
|
|
t.Run("fail to get collection id", func(t *testing.T) {
|
|
dt := deleteTask{
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: "foo",
|
|
}}
|
|
cache := NewMockCache(t)
|
|
cache.On("GetCollectionID",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(int64(0), errors.New("mock GetCollectionID err"))
|
|
globalMetaCache = cache
|
|
assert.Error(t, dt.PreExecute(context.Background()))
|
|
})
|
|
|
|
t.Run("fail partition key mode", func(t *testing.T) {
|
|
dt := deleteTask{req: &milvuspb.DeleteRequest{
|
|
CollectionName: "foo",
|
|
DbName: "db_1",
|
|
}}
|
|
cache := NewMockCache(t)
|
|
cache.On("GetCollectionID",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(int64(10000), nil)
|
|
cache.On("GetCollectionSchema",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(nil, errors.New("mock GetCollectionSchema err"))
|
|
|
|
globalMetaCache = cache
|
|
assert.Error(t, dt.PreExecute(context.Background()))
|
|
})
|
|
|
|
t.Run("invalid partition name", func(t *testing.T) {
|
|
dt := deleteTask{req: &milvuspb.DeleteRequest{
|
|
CollectionName: "foo",
|
|
DbName: "db_1",
|
|
PartitionName: "aaa",
|
|
}}
|
|
cache := NewMockCache(t)
|
|
cache.On("GetCollectionID",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(int64(10000), nil)
|
|
cache.On("GetCollectionSchema",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(&schemapb.CollectionSchema{
|
|
Name: "test_delete",
|
|
Description: "",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.StartOfUserFieldID,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPartitionKey: true,
|
|
},
|
|
},
|
|
}, nil)
|
|
|
|
globalMetaCache = cache
|
|
assert.Error(t, dt.PreExecute(context.Background()))
|
|
})
|
|
|
|
t.Run("invalie partition", func(t *testing.T) {
|
|
|
|
dt := deleteTask{
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: "foo",
|
|
DbName: "db_1",
|
|
PartitionName: "aaa",
|
|
Expr: "non_pk in [1, 2, 3]",
|
|
},
|
|
}
|
|
cache := NewMockCache(t)
|
|
cache.On("GetCollectionID",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(int64(10000), nil)
|
|
cache.On("GetCollectionSchema",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(schema, nil)
|
|
cache.On("GetPartitionID",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(int64(0), errors.New("mock GetPartitionID err"))
|
|
|
|
globalMetaCache = cache
|
|
assert.Error(t, dt.PreExecute(context.Background()))
|
|
|
|
dt.req.PartitionName = "aaa"
|
|
assert.Error(t, dt.PreExecute(context.Background()))
|
|
|
|
cache.On("GetPartitionID",
|
|
mock.Anything, // context.Context
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
mock.AnythingOfType("string"),
|
|
).Return(int64(100001), nil)
|
|
assert.Error(t, dt.PreExecute(context.Background()))
|
|
})
|
|
}
|
|
|
|
func TestDeleteTask_Execute(t *testing.T) {
|
|
collectionName := "test_delete"
|
|
collectionID := int64(111)
|
|
partitionName := "default"
|
|
partitionID := int64(222)
|
|
channels := []string{"test_channel"}
|
|
dbName := "test_1"
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: collectionName,
|
|
Description: "",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.StartOfUserFieldID,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: common.StartOfUserFieldID + 1,
|
|
Name: "non_pk",
|
|
IsPrimaryKey: false,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
},
|
|
}
|
|
t.Run("empty expr", func(t *testing.T) {
|
|
dt := deleteTask{}
|
|
assert.Error(t, dt.Execute(context.Background()))
|
|
})
|
|
|
|
t.Run("get channel failed", func(t *testing.T) {
|
|
mockMgr := NewMockChannelsMgr(t)
|
|
dt := deleteTask{
|
|
chMgr: mockMgr,
|
|
req: &milvuspb.DeleteRequest{
|
|
Expr: "pk in [1,2]",
|
|
},
|
|
}
|
|
|
|
mockMgr.EXPECT().getOrCreateDmlStream(mock.Anything).Return(nil, errors.New("mock error"))
|
|
assert.Error(t, dt.Execute(context.Background()))
|
|
})
|
|
|
|
t.Run("create plan failed", func(t *testing.T) {
|
|
mockMgr := NewMockChannelsMgr(t)
|
|
dt := deleteTask{
|
|
chMgr: mockMgr,
|
|
schema: schema,
|
|
req: &milvuspb.DeleteRequest{
|
|
Expr: "????",
|
|
},
|
|
}
|
|
stream := msgstream.NewMockMsgStream(t)
|
|
mockMgr.EXPECT().getOrCreateDmlStream(mock.Anything).Return(stream, nil)
|
|
assert.Error(t, dt.Execute(context.Background()))
|
|
})
|
|
|
|
t.Run("alloc failed", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
mockMgr := NewMockChannelsMgr(t)
|
|
rc := mocks.NewRootCoord(t)
|
|
allocator, err := allocator.NewIDAllocator(ctx, rc, paramtable.GetNodeID())
|
|
assert.NoError(t, err)
|
|
allocator.Close()
|
|
|
|
dt := deleteTask{
|
|
chMgr: mockMgr,
|
|
schema: schema,
|
|
collectionID: collectionID,
|
|
partitionID: partitionID,
|
|
vChannels: channels,
|
|
idAllocator: allocator,
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: collectionName,
|
|
PartitionName: partitionName,
|
|
DbName: dbName,
|
|
Expr: "pk in [1,2]",
|
|
},
|
|
}
|
|
stream := msgstream.NewMockMsgStream(t)
|
|
mockMgr.EXPECT().getOrCreateDmlStream(mock.Anything).Return(stream, nil)
|
|
|
|
assert.Error(t, dt.Execute(context.Background()))
|
|
})
|
|
|
|
t.Run("simple delete failed", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
mockMgr := NewMockChannelsMgr(t)
|
|
rc := mocks.NewRootCoord(t)
|
|
rc.EXPECT().AllocID(mock.Anything, mock.Anything).Return(
|
|
&rootcoordpb.AllocIDResponse{
|
|
Status: merr.Status(nil),
|
|
ID: 0,
|
|
Count: 1,
|
|
}, nil)
|
|
allocator, err := allocator.NewIDAllocator(ctx, rc, paramtable.GetNodeID())
|
|
allocator.Start()
|
|
assert.NoError(t, err)
|
|
|
|
dt := deleteTask{
|
|
chMgr: mockMgr,
|
|
schema: schema,
|
|
collectionID: collectionID,
|
|
partitionID: partitionID,
|
|
vChannels: channels,
|
|
idAllocator: allocator,
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: collectionName,
|
|
PartitionName: partitionName,
|
|
DbName: dbName,
|
|
Expr: "pk in [1,2]",
|
|
},
|
|
}
|
|
stream := msgstream.NewMockMsgStream(t)
|
|
mockMgr.EXPECT().getOrCreateDmlStream(mock.Anything).Return(stream, nil)
|
|
stream.EXPECT().Produce(mock.Anything).Return(errors.New("mock error"))
|
|
assert.Error(t, dt.Execute(context.Background()))
|
|
})
|
|
|
|
t.Run("complex delete query rpc failed", func(t *testing.T) {
|
|
mockMgr := NewMockChannelsMgr(t)
|
|
qn := mocks.NewMockQueryNode(t)
|
|
lb := NewMockLBPolicy(t)
|
|
|
|
dt := deleteTask{
|
|
chMgr: mockMgr,
|
|
schema: schema,
|
|
collectionID: collectionID,
|
|
partitionID: partitionID,
|
|
vChannels: channels,
|
|
lb: lb,
|
|
result: &milvuspb.MutationResult{
|
|
Status: merr.Status(nil),
|
|
IDs: &schemapb.IDs{
|
|
IdField: nil,
|
|
},
|
|
},
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: collectionName,
|
|
PartitionName: partitionName,
|
|
DbName: dbName,
|
|
Expr: "pk < 3",
|
|
},
|
|
}
|
|
stream := msgstream.NewMockMsgStream(t)
|
|
mockMgr.EXPECT().getOrCreateDmlStream(mock.Anything).Return(stream, nil)
|
|
lb.EXPECT().Execute(mock.Anything, mock.Anything).Call.Return(func(ctx context.Context, workload CollectionWorkLoad) error {
|
|
return workload.exec(ctx, 1, qn)
|
|
})
|
|
|
|
qn.EXPECT().QueryStream(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error"))
|
|
assert.Error(t, dt.Execute(context.Background()))
|
|
assert.Equal(t, int64(0), dt.result.DeleteCnt)
|
|
})
|
|
|
|
t.Run("complex delete query failed", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
mockMgr := NewMockChannelsMgr(t)
|
|
rc := mocks.NewRootCoord(t)
|
|
qn := mocks.NewMockQueryNode(t)
|
|
lb := NewMockLBPolicy(t)
|
|
rc.EXPECT().AllocID(mock.Anything, mock.Anything).Return(
|
|
&rootcoordpb.AllocIDResponse{
|
|
Status: merr.Status(nil),
|
|
ID: 0,
|
|
Count: 1,
|
|
}, nil)
|
|
allocator, err := allocator.NewIDAllocator(ctx, rc, paramtable.GetNodeID())
|
|
allocator.Start()
|
|
assert.NoError(t, err)
|
|
|
|
dt := deleteTask{
|
|
chMgr: mockMgr,
|
|
schema: schema,
|
|
collectionID: collectionID,
|
|
partitionID: partitionID,
|
|
vChannels: channels,
|
|
idAllocator: allocator,
|
|
lb: lb,
|
|
result: &milvuspb.MutationResult{
|
|
Status: merr.Status(nil),
|
|
IDs: &schemapb.IDs{
|
|
IdField: nil,
|
|
},
|
|
},
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: collectionName,
|
|
PartitionName: partitionName,
|
|
DbName: dbName,
|
|
Expr: "pk < 3",
|
|
},
|
|
}
|
|
stream := msgstream.NewMockMsgStream(t)
|
|
mockMgr.EXPECT().getOrCreateDmlStream(mock.Anything).Return(stream, nil)
|
|
lb.EXPECT().Execute(mock.Anything, mock.Anything).Call.Return(func(ctx context.Context, workload CollectionWorkLoad) error {
|
|
return workload.exec(ctx, 1, qn)
|
|
})
|
|
|
|
qn.EXPECT().QueryStream(mock.Anything, mock.Anything, mock.Anything).Run(
|
|
func(ctx context.Context, req *querypb.QueryRequest, streamer streamrpc.QueryStreamer) {
|
|
client := streamrpc.NewLocalQueryClient(ctx)
|
|
server := client.CreateServer()
|
|
streamer.SetClient(client)
|
|
|
|
server.Send(&internalpb.RetrieveResults{
|
|
Status: merr.Status(nil),
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: []int64{0, 1, 2},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
|
|
server.Send(&internalpb.RetrieveResults{
|
|
Status: merr.Status(errors.New("mock error")),
|
|
})
|
|
}).Return(nil)
|
|
stream.EXPECT().Produce(mock.Anything).Return(nil)
|
|
|
|
assert.Error(t, dt.Execute(context.Background()))
|
|
// query failed but still delete some data before failed.
|
|
assert.Equal(t, int64(3), dt.result.DeleteCnt)
|
|
})
|
|
|
|
t.Run("complex delete produce failed", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
mockMgr := NewMockChannelsMgr(t)
|
|
rc := mocks.NewRootCoord(t)
|
|
qn := mocks.NewMockQueryNode(t)
|
|
lb := NewMockLBPolicy(t)
|
|
rc.EXPECT().AllocID(mock.Anything, mock.Anything).Return(
|
|
&rootcoordpb.AllocIDResponse{
|
|
Status: merr.Status(nil),
|
|
ID: 0,
|
|
Count: 1,
|
|
}, nil)
|
|
allocator, err := allocator.NewIDAllocator(ctx, rc, paramtable.GetNodeID())
|
|
allocator.Start()
|
|
assert.NoError(t, err)
|
|
|
|
dt := deleteTask{
|
|
chMgr: mockMgr,
|
|
schema: schema,
|
|
collectionID: collectionID,
|
|
partitionID: partitionID,
|
|
vChannels: channels,
|
|
idAllocator: allocator,
|
|
lb: lb,
|
|
result: &milvuspb.MutationResult{
|
|
Status: merr.Status(nil),
|
|
IDs: &schemapb.IDs{
|
|
IdField: nil,
|
|
},
|
|
},
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: collectionName,
|
|
PartitionName: partitionName,
|
|
DbName: dbName,
|
|
Expr: "pk < 3",
|
|
},
|
|
}
|
|
stream := msgstream.NewMockMsgStream(t)
|
|
mockMgr.EXPECT().getOrCreateDmlStream(mock.Anything).Return(stream, nil)
|
|
lb.EXPECT().Execute(mock.Anything, mock.Anything).Call.Return(func(ctx context.Context, workload CollectionWorkLoad) error {
|
|
return workload.exec(ctx, 1, qn)
|
|
})
|
|
|
|
qn.EXPECT().QueryStream(mock.Anything, mock.Anything, mock.Anything).Run(
|
|
func(ctx context.Context, req *querypb.QueryRequest, streamer streamrpc.QueryStreamer) {
|
|
client := streamrpc.NewLocalQueryClient(ctx)
|
|
server := client.CreateServer()
|
|
streamer.SetClient(client)
|
|
|
|
server.Send(&internalpb.RetrieveResults{
|
|
Status: merr.Status(nil),
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: []int64{0, 1, 2},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
server.FinishSend(nil)
|
|
}).Return(nil)
|
|
stream.EXPECT().Produce(mock.Anything).Return(errors.New("mock error"))
|
|
|
|
assert.Error(t, dt.Execute(context.Background()))
|
|
assert.Equal(t, int64(0), dt.result.DeleteCnt)
|
|
})
|
|
|
|
t.Run("complex delete success", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
mockMgr := NewMockChannelsMgr(t)
|
|
rc := mocks.NewRootCoord(t)
|
|
qn := mocks.NewMockQueryNode(t)
|
|
lb := NewMockLBPolicy(t)
|
|
rc.EXPECT().AllocID(mock.Anything, mock.Anything).Return(
|
|
&rootcoordpb.AllocIDResponse{
|
|
Status: merr.Status(nil),
|
|
ID: 0,
|
|
Count: 1,
|
|
}, nil)
|
|
allocator, err := allocator.NewIDAllocator(ctx, rc, paramtable.GetNodeID())
|
|
allocator.Start()
|
|
assert.NoError(t, err)
|
|
|
|
dt := deleteTask{
|
|
chMgr: mockMgr,
|
|
schema: schema,
|
|
collectionID: collectionID,
|
|
partitionID: partitionID,
|
|
vChannels: channels,
|
|
idAllocator: allocator,
|
|
lb: lb,
|
|
result: &milvuspb.MutationResult{
|
|
Status: merr.Status(nil),
|
|
IDs: &schemapb.IDs{
|
|
IdField: nil,
|
|
},
|
|
},
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: collectionName,
|
|
PartitionName: partitionName,
|
|
DbName: dbName,
|
|
Expr: "pk < 3",
|
|
},
|
|
}
|
|
stream := msgstream.NewMockMsgStream(t)
|
|
mockMgr.EXPECT().getOrCreateDmlStream(mock.Anything).Return(stream, nil)
|
|
lb.EXPECT().Execute(mock.Anything, mock.Anything).Call.Return(func(ctx context.Context, workload CollectionWorkLoad) error {
|
|
return workload.exec(ctx, 1, qn)
|
|
})
|
|
|
|
qn.EXPECT().QueryStream(mock.Anything, mock.Anything, mock.Anything).Run(
|
|
func(ctx context.Context, req *querypb.QueryRequest, streamer streamrpc.QueryStreamer) {
|
|
client := streamrpc.NewLocalQueryClient(ctx)
|
|
server := client.CreateServer()
|
|
streamer.SetClient(client)
|
|
|
|
server.Send(&internalpb.RetrieveResults{
|
|
Status: merr.Status(nil),
|
|
Ids: &schemapb.IDs{
|
|
IdField: &schemapb.IDs_IntId{
|
|
IntId: &schemapb.LongArray{
|
|
Data: []int64{0, 1, 2},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
server.FinishSend(nil)
|
|
}).Return(nil)
|
|
stream.EXPECT().Produce(mock.Anything).Return(nil)
|
|
|
|
assert.NoError(t, dt.Execute(context.Background()))
|
|
assert.Equal(t, int64(3), dt.result.DeleteCnt)
|
|
})
|
|
}
|
|
|
|
func TestDeleteTask_SimpleDelete(t *testing.T) {
|
|
collectionName := "test_delete"
|
|
collectionID := int64(111)
|
|
partitionName := "default"
|
|
partitionID := int64(222)
|
|
dbName := "test_1"
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Name: collectionName,
|
|
Description: "",
|
|
AutoID: false,
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: common.StartOfUserFieldID,
|
|
Name: "pk",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: common.StartOfUserFieldID + 1,
|
|
Name: "non_pk",
|
|
IsPrimaryKey: false,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
},
|
|
}
|
|
|
|
task := deleteTask{
|
|
schema: schema,
|
|
collectionID: collectionID,
|
|
partitionID: partitionID,
|
|
req: &milvuspb.DeleteRequest{
|
|
CollectionName: collectionName,
|
|
PartitionName: partitionName,
|
|
DbName: dbName,
|
|
},
|
|
}
|
|
t.Run("get PK failed", func(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
expr := &planpb.Expr_TermExpr{
|
|
TermExpr: &planpb.TermExpr{
|
|
ColumnInfo: &planpb.ColumnInfo{
|
|
DataType: schemapb.DataType_BinaryVector,
|
|
}}}
|
|
stream := msgstream.NewMockMsgStream(t)
|
|
err := task.simpleDelete(ctx, expr, stream)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|