test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-12-01 03:18:29 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Fix unstable search task ut (#28101)

Browse Source 
Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
This commit is contained in:
yihao.dai 2023-11-02 12:52:15 +08:00 committed by GitHub
parent d39ffba5d1
commit 7bd44bd671
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 33 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
internal/proxy/task_search_test.go
View File

@ -28,6 +28,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
@ -1938,32 +1939,45 @@ func TestSearchTask_Requery(t *testing.T) {
t.Run("Test normal", func(t *testing.T) {
schema := constructCollectionSchema(pkField, vecField, dim, collection)
qn := mocks.NewMockQueryNodeClient(t)
qn.EXPECT().Query(mock.Anything, mock.Anything).
Return(&internalpb.RetrieveResults{
Ids: &schemapb.IDs{
qn.EXPECT().Query(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, request *querypb.QueryRequest, option ...grpc.CallOption) (*internalpb.RetrieveResults, error) {
idFieldData := &schemapb.FieldData{
Type: schemapb.DataType_Int64,
FieldName: pkField,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: ids,
},
},
},
},
}
idField := &schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: ids,
},
},
},
FieldsData: []*schemapb.FieldData{
{
Type: schemapb.DataType_Int64,
FieldName: pkField,
Field: &schemapb.FieldData_Scalars{
Scalars: &schemapb.ScalarField{
Data: &schemapb.ScalarField_LongData{
LongData: &schemapb.LongArray{
Data: ids,
},
},
},
}
if request.GetReq().GetOutputFieldsId()[0] == 100 {
return &internalpb.RetrieveResults{
Ids: idField,
FieldsData: []*schemapb.FieldData{
idFieldData,
newFloatVectorFieldData(vecField, rows, dim),
},
}, nil
}
return &internalpb.RetrieveResults{
Ids: idField,
FieldsData: []*schemapb.FieldData{
newFloatVectorFieldData(vecField, rows, dim),
idFieldData,
},
newFloatVectorFieldData(vecField, rows, dim),
},
}, nil)
}, nil
})
lb := NewMockLBPolicy(t)
lb.EXPECT().Execute(mock.Anything, mock.Anything).Run(func(ctx context.Context, workload CollectionWorkLoad) {