milvus/internal/distributed/proxy/httpserver/handler_v1_test.go
SimFG 26f06dd732
Format the code (#27275)
Signed-off-by: SimFG <bang.fu@zilliz.com>
2023-09-21 09:45:27 +08:00

1495 lines
53 KiB
Go

package httpserver
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"github.com/cockroachdb/errors"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"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-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proxy"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/util"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
const (
DefaultCollectionName = "book"
DefaultErrorMessage = "UNKNOWN ERROR"
ReturnSuccess = 0
ReturnFail = 1
ReturnWrongStatus = 2
ReturnTrue = 3
ReturnFalse = 4
URIPrefix = "/v1"
)
var StatusSuccess = commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Reason: "",
}
var ErrDefault = errors.New(DefaultErrorMessage)
var DefaultShowCollectionsResp = milvuspb.ShowCollectionsResponse{
Status: &StatusSuccess,
CollectionNames: []string{DefaultCollectionName},
}
var DefaultDescCollectionResp = milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(false),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}
var DefaultLoadStateResp = milvuspb.GetLoadStateResponse{
Status: &StatusSuccess,
State: commonpb.LoadState_LoadStateLoaded,
}
var DefaultDescIndexesReqp = milvuspb.DescribeIndexResponse{
Status: &StatusSuccess,
IndexDescriptions: generateIndexes(),
}
var DefaultTrueResp = milvuspb.BoolResponse{
Status: &StatusSuccess,
Value: true,
}
var DefaultFalseResp = milvuspb.BoolResponse{
Status: &StatusSuccess,
Value: false,
}
func initHTTPServer(proxy types.ProxyComponent, needAuth bool) *gin.Engine {
h := NewHandlers(proxy)
ginHandler := gin.Default()
app := ginHandler.Group("/v1", genAuthMiddleWare(needAuth))
NewHandlers(h.proxy).RegisterRoutesToV1(app)
return ginHandler
}
/*
mock authentication
----|----------------|----------------|----------------
| username | password | result
----|----------------|----------------|----------------
|"" |"" |fail
|root |Milvus |success
|root |{not Milvus} |fail
|{not root} |* |success
----|----------------|----------------|----------------
*/
func genAuthMiddleWare(needAuth bool) gin.HandlerFunc {
if needAuth {
return func(c *gin.Context) {
username, password, ok := ParseUsernamePassword(c)
if !ok {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{HTTPReturnCode: Code(merr.ErrNeedAuthenticate), HTTPReturnMessage: merr.ErrNeedAuthenticate.Error()})
} else if username == util.UserRoot && password != util.DefaultRootPassword {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{HTTPReturnCode: Code(merr.ErrNeedAuthenticate), HTTPReturnMessage: merr.ErrNeedAuthenticate.Error()})
} else {
c.Set(ContextUsername, username)
}
}
}
return func(c *gin.Context) {
c.Set(ContextUsername, util.UserRoot)
}
}
func Print(code int32, message string) string {
return fmt.Sprintf("{\"%s\":%d,\"%s\":\"%s\"}", HTTPReturnCode, code, HTTPReturnMessage, message)
}
func PrintErr(err error) string {
return Print(Code(err), err.Error())
}
func TestVectorAuthenticate(t *testing.T) {
paramtable.Init()
mp := mocks.NewMockProxy(t)
mp.EXPECT().
ShowCollections(mock.Anything, mock.Anything).
Return(&DefaultShowCollectionsResp, nil).
Twice()
testEngine := initHTTPServer(mp, true)
t.Run("need authentication", func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/v1/vector/collections", nil)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusUnauthorized)
assert.Equal(t, w.Body.String(), PrintErr(merr.ErrNeedAuthenticate))
})
t.Run("username or password incorrect", func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/v1/vector/collections", nil)
req.SetBasicAuth(util.UserRoot, util.UserRoot)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusUnauthorized)
assert.Equal(t, w.Body.String(), PrintErr(merr.ErrNeedAuthenticate))
})
t.Run("root's password correct", func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/v1/vector/collections", nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), "{\"code\":200,\"data\":[\""+DefaultCollectionName+"\"]}")
})
t.Run("username and password both provided", func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/v1/vector/collections", nil)
req.SetBasicAuth("test", util.UserRoot)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), "{\"code\":200,\"data\":[\""+DefaultCollectionName+"\"]}")
})
}
func TestVectorListCollection(t *testing.T) {
paramtable.Init()
testCases := []testCase{}
mp0 := mocks.NewMockProxy(t)
mp0.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testCases = append(testCases, testCase{
name: "show collections fail",
mp: mp0,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
reason := "cannot create folder"
mp1 := mocks.NewMockProxy(t)
mp1.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&milvuspb.ShowCollectionsResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CannotCreateFolder,
Reason: reason,
},
}, nil).Once()
testCases = append(testCases, testCase{
name: "show collections fail",
mp: mp1,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_CannotCreateFolder), reason),
})
mp := mocks.NewMockProxy(t)
mp.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&DefaultShowCollectionsResp, nil).Once()
testCases = append(testCases, testCase{
name: "list collections success",
mp: mp,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":[\"" + DefaultCollectionName + "\"]}",
})
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
testEngine := initHTTPServer(tt.mp, true)
req := httptest.NewRequest(http.MethodGet, "/v1/vector/collections", nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, tt.exceptCode)
assert.Equal(t, w.Body.String(), tt.expectedBody)
})
}
}
type testCase struct {
name string
mp *mocks.MockProxy
exceptCode int
expectedBody string
}
func TestVectorCollectionsDescribe(t *testing.T) {
paramtable.Init()
testCases := []testCase{}
_, testCases = wrapWithDescribeColl(t, nil, ReturnFail, 1, testCases)
_, testCases = wrapWithDescribeColl(t, nil, ReturnWrongStatus, 1, testCases)
mp2 := mocks.NewMockProxy(t)
mp2, _ = wrapWithDescribeColl(t, mp2, ReturnSuccess, 1, nil)
mp2.EXPECT().GetLoadState(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
mp2, _ = wrapWithDescribeIndex(t, mp2, ReturnSuccess, 1, nil)
testCases = append(testCases, testCase{
name: "get load status fail",
mp: mp2,
exceptCode: http.StatusOK,
expectedBody: "{\"code\":200,\"data\":{\"collectionName\":\"" + DefaultCollectionName + "\",\"description\":\"\",\"enableDynamic\":true,\"fields\":[{\"autoId\":false,\"description\":\"\",\"name\":\"book_id\",\"primaryKey\":true,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"\",\"name\":\"word_count\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"\",\"name\":\"book_intro\",\"primaryKey\":false,\"type\":\"FloatVector(2)\"}],\"indexes\":[{\"fieldName\":\"book_intro\",\"indexName\":\"" + DefaultIndexName + "\",\"metricType\":\"L2\"}],\"load\":\"\",\"shardsNum\":1}}",
})
mp3 := mocks.NewMockProxy(t)
mp3, _ = wrapWithDescribeColl(t, mp3, ReturnSuccess, 1, nil)
mp3.EXPECT().GetLoadState(mock.Anything, mock.Anything).Return(&DefaultLoadStateResp, nil).Once()
mp3, _ = wrapWithDescribeIndex(t, mp3, ReturnFail, 1, nil)
testCases = append(testCases, testCase{
name: "get indexes fail",
mp: mp3,
exceptCode: http.StatusOK,
expectedBody: "{\"code\":200,\"data\":{\"collectionName\":\"" + DefaultCollectionName + "\",\"description\":\"\",\"enableDynamic\":true,\"fields\":[{\"autoId\":false,\"description\":\"\",\"name\":\"book_id\",\"primaryKey\":true,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"\",\"name\":\"word_count\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"\",\"name\":\"book_intro\",\"primaryKey\":false,\"type\":\"FloatVector(2)\"}],\"indexes\":[],\"load\":\"LoadStateLoaded\",\"shardsNum\":1}}",
})
mp4 := mocks.NewMockProxy(t)
mp4, _ = wrapWithDescribeColl(t, mp4, ReturnSuccess, 1, nil)
mp4.EXPECT().GetLoadState(mock.Anything, mock.Anything).Return(&DefaultLoadStateResp, nil).Once()
mp4, _ = wrapWithDescribeIndex(t, mp4, ReturnSuccess, 1, nil)
testCases = append(testCases, testCase{
name: "show collection details success",
mp: mp4,
exceptCode: http.StatusOK,
expectedBody: "{\"code\":200,\"data\":{\"collectionName\":\"" + DefaultCollectionName + "\",\"description\":\"\",\"enableDynamic\":true,\"fields\":[{\"autoId\":false,\"description\":\"\",\"name\":\"book_id\",\"primaryKey\":true,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"\",\"name\":\"word_count\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"\",\"name\":\"book_intro\",\"primaryKey\":false,\"type\":\"FloatVector(2)\"}],\"indexes\":[{\"fieldName\":\"book_intro\",\"indexName\":\"" + DefaultIndexName + "\",\"metricType\":\"L2\"}],\"load\":\"LoadStateLoaded\",\"shardsNum\":1}}",
})
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
testEngine := initHTTPServer(tt.mp, true)
req := httptest.NewRequest(http.MethodGet, "/v1/vector/collections/describe?collectionName="+DefaultCollectionName, nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, tt.exceptCode)
assert.Equal(t, w.Body.String(), tt.expectedBody)
})
}
t.Run("need collectionName", func(t *testing.T) {
testEngine := initHTTPServer(mocks.NewMockProxy(t), true)
req := httptest.NewRequest(http.MethodGet, "/v1/vector/collections/describe?"+DefaultCollectionName, nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, 200)
assert.Equal(t, w.Body.String(), PrintErr(merr.ErrMissingRequiredParameters))
})
}
func TestVectorCreateCollection(t *testing.T) {
paramtable.Init()
testCases := []testCase{}
mp1 := mocks.NewMockProxy(t)
mp1.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testCases = append(testCases, testCase{
name: "create collection fail",
mp: mp1,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
reason := "collection " + DefaultCollectionName + " already exists"
mp2 := mocks.NewMockProxy(t)
mp2.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(&commonpb.Status{
ErrorCode: commonpb.ErrorCode_CannotCreateFile, // 18
Reason: reason,
}, nil).Once()
testCases = append(testCases, testCase{
name: "create collection fail",
mp: mp2,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_CannotCreateFile), reason),
})
mp3 := mocks.NewMockProxy(t)
mp3.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(&StatusSuccess, nil).Once()
mp3.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testCases = append(testCases, testCase{
name: "create index fail",
mp: mp3,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
mp4 := mocks.NewMockProxy(t)
mp4.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(&StatusSuccess, nil).Once()
mp4.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(&StatusSuccess, nil).Once()
mp4.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testCases = append(testCases, testCase{
name: "load collection fail",
mp: mp4,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
mp5 := mocks.NewMockProxy(t)
mp5.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(&StatusSuccess, nil).Once()
mp5.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(&StatusSuccess, nil).Once()
mp5.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&StatusSuccess, nil).Once()
testCases = append(testCases, testCase{
name: "create collection success",
mp: mp5,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":{}}",
})
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
testEngine := initHTTPServer(tt.mp, true)
jsonBody := []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2}`)
bodyReader := bytes.NewReader(jsonBody)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/collections/create", bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, tt.exceptCode)
assert.Equal(t, w.Body.String(), tt.expectedBody)
})
}
}
func TestVectorDropCollection(t *testing.T) {
paramtable.Init()
testCases := []testCase{}
_, testCases = wrapWithHasCollection(t, nil, ReturnFail, 1, testCases)
_, testCases = wrapWithHasCollection(t, nil, ReturnWrongStatus, 1, testCases)
_, testCases = wrapWithHasCollection(t, nil, ReturnFalse, 1, testCases)
mp1 := mocks.NewMockProxy(t)
mp1, _ = wrapWithHasCollection(t, mp1, ReturnTrue, 1, nil)
mp1.EXPECT().DropCollection(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testCases = append(testCases, testCase{
name: "drop collection fail",
mp: mp1,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
reason := "cannot find collection " + DefaultCollectionName
mp2 := mocks.NewMockProxy(t)
mp2, _ = wrapWithHasCollection(t, mp2, ReturnTrue, 1, nil)
mp2.EXPECT().DropCollection(mock.Anything, mock.Anything).Return(&commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists, // 4
Reason: reason,
}, nil).Once()
testCases = append(testCases, testCase{
name: "drop collection fail",
mp: mp2,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_CollectionNotExists), reason),
})
mp3 := mocks.NewMockProxy(t)
mp3, _ = wrapWithHasCollection(t, mp3, ReturnTrue, 1, nil)
mp3.EXPECT().DropCollection(mock.Anything, mock.Anything).Return(&StatusSuccess, nil).Once()
testCases = append(testCases, testCase{
name: "create collection fail",
mp: mp3,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":{}}",
})
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
testEngine := initHTTPServer(tt.mp, true)
jsonBody := []byte(`{"collectionName": "` + DefaultCollectionName + `"}`)
bodyReader := bytes.NewReader(jsonBody)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/collections/drop", bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, tt.exceptCode)
assert.Equal(t, w.Body.String(), tt.expectedBody)
})
}
}
func TestQuery(t *testing.T) {
paramtable.Init()
testCases := []testCase{}
mp2 := mocks.NewMockProxy(t)
mp2, _ = wrapWithDescribeColl(t, mp2, ReturnSuccess, 1, nil)
mp2.EXPECT().Query(mock.Anything, mock.Anything).Return(nil, ErrDefault).Twice()
testCases = append(testCases, testCase{
name: "query fail",
mp: mp2,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
reason := DefaultCollectionName + " name not found"
mp3 := mocks.NewMockProxy(t)
mp3, _ = wrapWithDescribeColl(t, mp3, ReturnSuccess, 1, nil)
mp3.EXPECT().Query(mock.Anything, mock.Anything).Return(&milvuspb.QueryResults{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNameNotFound, // 28
Reason: reason,
},
}, nil).Twice()
testCases = append(testCases, testCase{
name: "query fail",
mp: mp3,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_CollectionNameNotFound), reason),
})
mp4 := mocks.NewMockProxy(t)
mp4, _ = wrapWithDescribeColl(t, mp4, ReturnSuccess, 1, nil)
mp4.EXPECT().Query(mock.Anything, mock.Anything).Return(&milvuspb.QueryResults{
Status: &StatusSuccess,
FieldsData: generateFieldData(),
CollectionName: DefaultCollectionName,
OutputFields: []string{FieldBookID, FieldWordCount, FieldBookIntro},
}, nil).Twice()
testCases = append(testCases, testCase{
name: "query success",
mp: mp4,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":[{\"book_id\":1,\"book_intro\":[0.1,0.11],\"word_count\":1000},{\"book_id\":2,\"book_intro\":[0.2,0.22],\"word_count\":2000},{\"book_id\":3,\"book_intro\":[0.3,0.33],\"word_count\":3000}]}",
})
for _, tt := range testCases {
reqs := []*http.Request{genQueryRequest(), genGetRequest()}
t.Run(tt.name, func(t *testing.T) {
testEngine := initHTTPServer(tt.mp, true)
for _, req := range reqs {
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, tt.exceptCode)
assert.Equal(t, w.Body.String(), tt.expectedBody)
resp := map[string]interface{}{}
err := json.Unmarshal(w.Body.Bytes(), &resp)
assert.Equal(t, err, nil)
if resp[HTTPReturnCode] == float64(200) {
data := resp[HTTPReturnData].([]interface{})
rows := generateQueryResult64(false)
for i, row := range data {
assert.Equal(t, compareRow64(row.(map[string]interface{}), rows[i]), true)
}
}
}
})
}
}
func genQueryRequest() *http.Request {
jsonBody := []byte(`{"collectionName": "` + DefaultCollectionName + `" , "filter": "book_id in [1,2,3]"}`)
bodyReader := bytes.NewReader(jsonBody)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/query", bodyReader)
return req
}
func genGetRequest() *http.Request {
jsonBody := []byte(`{"collectionName": "` + DefaultCollectionName + `" , "id": [1,2,3]}`)
bodyReader := bytes.NewReader(jsonBody)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/get", bodyReader)
return req
}
func TestDelete(t *testing.T) {
paramtable.Init()
testCases := []testCase{}
_, testCases = wrapWithDescribeColl(t, nil, ReturnFail, 1, testCases)
_, testCases = wrapWithDescribeColl(t, nil, ReturnWrongStatus, 1, testCases)
mp2 := mocks.NewMockProxy(t)
mp2, _ = wrapWithDescribeColl(t, mp2, ReturnSuccess, 1, nil)
mp2.EXPECT().Delete(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testCases = append(testCases, testCase{
name: "delete fail",
mp: mp2,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
reason := DefaultCollectionName + " name not found"
mp3 := mocks.NewMockProxy(t)
mp3, _ = wrapWithDescribeColl(t, mp3, ReturnSuccess, 1, nil)
mp3.EXPECT().Delete(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNameNotFound, // 28
Reason: reason,
},
}, nil).Once()
testCases = append(testCases, testCase{
name: "delete fail",
mp: mp3,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_CollectionNameNotFound), reason),
})
mp4 := mocks.NewMockProxy(t)
mp4, _ = wrapWithDescribeColl(t, mp4, ReturnSuccess, 1, nil)
mp4.EXPECT().Delete(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{
Status: &StatusSuccess,
}, nil).Once()
testCases = append(testCases, testCase{
name: "delete success",
mp: mp4,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":{}}",
})
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
testEngine := initHTTPServer(tt.mp, true)
jsonBody := []byte(`{"collectionName": "` + DefaultCollectionName + `" , "id": [1,2,3]}`)
bodyReader := bytes.NewReader(jsonBody)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/delete", bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, tt.exceptCode)
assert.Equal(t, w.Body.String(), tt.expectedBody)
resp := map[string]interface{}{}
err := json.Unmarshal(w.Body.Bytes(), &resp)
assert.Equal(t, err, nil)
})
}
}
func TestInsert(t *testing.T) {
paramtable.Init()
testCases := []testCase{}
_, testCases = wrapWithDescribeColl(t, nil, ReturnFail, 1, testCases)
_, testCases = wrapWithDescribeColl(t, nil, ReturnWrongStatus, 1, testCases)
mp2 := mocks.NewMockProxy(t)
mp2, _ = wrapWithDescribeColl(t, mp2, ReturnSuccess, 1, nil)
mp2.EXPECT().Insert(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testCases = append(testCases, testCase{
name: "insert fail",
mp: mp2,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
reason := DefaultCollectionName + " name not found"
mp3 := mocks.NewMockProxy(t)
mp3, _ = wrapWithDescribeColl(t, mp3, ReturnSuccess, 1, nil)
mp3.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNameNotFound, // 28
Reason: reason,
},
}, nil).Once()
testCases = append(testCases, testCase{
name: "insert fail",
mp: mp3,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_CollectionNameNotFound), reason),
})
mp4 := mocks.NewMockProxy(t)
mp4, _ = wrapWithDescribeColl(t, mp4, ReturnSuccess, 1, nil)
mp4.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{
Status: &StatusSuccess,
}, nil).Once()
testCases = append(testCases, testCase{
name: "id type invalid",
mp: mp4,
exceptCode: 200,
expectedBody: PrintErr(merr.ErrCheckPrimaryKey),
})
mp5 := mocks.NewMockProxy(t)
mp5, _ = wrapWithDescribeColl(t, mp5, ReturnSuccess, 1, nil)
mp5.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{
Status: &StatusSuccess,
IDs: getIntIds(),
InsertCnt: 3,
}, nil).Once()
testCases = append(testCases, testCase{
name: "insert success",
mp: mp5,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":{\"insertCount\":3,\"insertIds\":[1,2,3]}}",
})
mp6 := mocks.NewMockProxy(t)
mp6, _ = wrapWithDescribeColl(t, mp6, ReturnSuccess, 1, nil)
mp6.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{
Status: &StatusSuccess,
IDs: getStrIds(),
InsertCnt: 3,
}, nil).Once()
testCases = append(testCases, testCase{
name: "insert success",
mp: mp6,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":{\"insertCount\":3,\"insertIds\":[\"1\",\"2\",\"3\"]}}",
})
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
testEngine := initHTTPServer(tt.mp, true)
rows := generateSearchResult()
data, _ := json.Marshal(map[string]interface{}{
HTTPCollectionName: DefaultCollectionName,
HTTPReturnData: rows[0],
})
bodyReader := bytes.NewReader(data)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/insert", bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, tt.exceptCode)
assert.Equal(t, w.Body.String(), tt.expectedBody)
resp := map[string]interface{}{}
err := json.Unmarshal(w.Body.Bytes(), &resp)
assert.Equal(t, err, nil)
})
}
t.Run("wrong request body", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp, _ = wrapWithDescribeColl(t, mp, ReturnSuccess, 1, nil)
testEngine := initHTTPServer(mp, true)
bodyReader := bytes.NewReader([]byte(`{"collectionName": "` + DefaultCollectionName + `", "data": {}}`))
req := httptest.NewRequest(http.MethodPost, "/v1/vector/insert", bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, 200)
assert.Equal(t, w.Body.String(), PrintErr(merr.ErrInvalidInsertData))
resp := map[string]interface{}{}
err := json.Unmarshal(w.Body.Bytes(), &resp)
assert.Equal(t, err, nil)
})
}
func TestInsertForDataType(t *testing.T) {
paramtable.Init()
schemas := map[string]*schemapb.CollectionSchema{
"[success]kinds of data type": newCollectionSchema(generateCollectionSchema(false)),
"[success]use binary vector": newCollectionSchema(generateCollectionSchema(true)),
"[success]with dynamic field": withDynamicField(newCollectionSchema(generateCollectionSchema(false))),
}
for name, schema := range schemas {
t.Run(name, func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: schema,
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{
Status: &StatusSuccess,
IDs: getIntIds(),
InsertCnt: 3,
}, nil).Once()
testEngine := initHTTPServer(mp, true)
rows := newSearchResult(generateSearchResult())
data, _ := json.Marshal(map[string]interface{}{
HTTPCollectionName: DefaultCollectionName,
HTTPReturnData: rows,
})
bodyReader := bytes.NewReader(data)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/insert", bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, 200)
assert.Equal(t, w.Body.String(), "{\"code\":200,\"data\":{\"insertCount\":3,\"insertIds\":[1,2,3]}}")
})
}
schemas = map[string]*schemapb.CollectionSchema{
"with unsupport field type": withUnsupportField(newCollectionSchema(generateCollectionSchema(false))),
}
for name, schema := range schemas {
t.Run(name, func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: schema,
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
testEngine := initHTTPServer(mp, true)
rows := newSearchResult(generateSearchResult())
data, _ := json.Marshal(map[string]interface{}{
HTTPCollectionName: DefaultCollectionName,
HTTPReturnData: rows,
})
bodyReader := bytes.NewReader(data)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/insert", bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, 200)
assert.Equal(t, w.Body.String(), PrintErr(merr.ErrInvalidInsertData))
})
}
}
func getIntIds() *schemapb.IDs {
ids := schemapb.IDs{
IdField: &schemapb.IDs_IntId{
IntId: &schemapb.LongArray{
Data: []int64{1, 2, 3},
},
},
}
return &ids
}
func getStrIds() *schemapb.IDs {
ids := schemapb.IDs{
IdField: &schemapb.IDs_StrId{
StrId: &schemapb.StringArray{
Data: []string{"1", "2", "3"},
},
},
}
return &ids
}
func TestSearch(t *testing.T) {
paramtable.Init()
testCases := []testCase{}
mp2 := mocks.NewMockProxy(t)
mp2.EXPECT().Search(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testCases = append(testCases, testCase{
name: "search fail",
mp: mp2,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
})
reason := DefaultCollectionName + " name not found"
mp3 := mocks.NewMockProxy(t)
mp3.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNameNotFound, // 28
Reason: reason,
},
}, nil).Once()
testCases = append(testCases, testCase{
name: "search fail",
mp: mp3,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_CollectionNameNotFound), reason),
})
mp4 := mocks.NewMockProxy(t)
mp4.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{
Status: &StatusSuccess,
Results: &schemapb.SearchResultData{
FieldsData: generateFieldData(),
Scores: []float32{0.01, 0.04, 0.09},
TopK: 3,
},
}, nil).Once()
testCases = append(testCases, testCase{
name: "search success",
mp: mp4,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":[{\"book_id\":1,\"book_intro\":[0.1,0.11],\"distance\":0.01,\"word_count\":1000},{\"book_id\":2,\"book_intro\":[0.2,0.22],\"distance\":0.04,\"word_count\":2000},{\"book_id\":3,\"book_intro\":[0.3,0.33],\"distance\":0.09,\"word_count\":3000}]}",
})
for _, tt := range testCases {
t.Run(tt.name, func(t *testing.T) {
testEngine := initHTTPServer(tt.mp, true)
rows := []float32{0.0, 0.0}
data, _ := json.Marshal(map[string]interface{}{
HTTPCollectionName: DefaultCollectionName,
"vector": rows,
})
bodyReader := bytes.NewReader(data)
req := httptest.NewRequest(http.MethodPost, "/v1/vector/search", bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, tt.exceptCode)
assert.Equal(t, w.Body.String(), tt.expectedBody)
resp := map[string]interface{}{}
err := json.Unmarshal(w.Body.Bytes(), &resp)
assert.Equal(t, err, nil)
if resp[HTTPReturnCode] == float64(200) {
data := resp[HTTPReturnData].([]interface{})
rows := generateQueryResult64(true)
for i, row := range data {
assert.Equal(t, compareRow64(row.(map[string]interface{}), rows[i]), true)
}
}
})
}
}
type ReturnType int
func wrapWithDescribeColl(t *testing.T, mp *mocks.MockProxy, returnType ReturnType, times int, testCases []testCase) (*mocks.MockProxy, []testCase) {
if mp == nil {
mp = mocks.NewMockProxy(t)
}
var call *mocks.MockProxy_DescribeCollection_Call
var testcase testCase
switch returnType {
case ReturnSuccess:
call = mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&DefaultDescCollectionResp, nil)
testcase = testCase{
name: "[share] describe coll success",
mp: mp,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":{}}",
}
case ReturnFail:
call = mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(nil, ErrDefault)
testcase = testCase{
name: "[share] describe coll fail",
mp: mp,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
}
case ReturnWrongStatus:
call = mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNotExists,
Reason: "can't find collection: " + DefaultCollectionName,
},
}, nil)
testcase = testCase{
name: "[share] collection not found",
mp: mp,
exceptCode: 200,
expectedBody: "{\"code\":4,\"message\":\"can't find collection: " + DefaultCollectionName + "\"}",
}
}
if times == 2 {
call.Twice()
} else {
call.Once()
}
if testCases != nil {
testCases = append(testCases, testcase)
}
return mp, testCases
}
func wrapWithHasCollection(t *testing.T, mp *mocks.MockProxy, returnType ReturnType, times int, testCases []testCase) (*mocks.MockProxy, []testCase) {
if mp == nil {
mp = mocks.NewMockProxy(t)
}
var call *mocks.MockProxy_HasCollection_Call
var testcase testCase
switch returnType {
case ReturnTrue:
call = mp.EXPECT().HasCollection(mock.Anything, mock.Anything).Return(&DefaultTrueResp, nil)
testcase = testCase{
name: "[share] collection already exists",
mp: mp,
exceptCode: 200,
expectedBody: "{\"code\":200,\"message\":\"collection " + DefaultCollectionName + " already exist.\"}",
}
case ReturnFalse:
call = mp.EXPECT().HasCollection(mock.Anything, mock.Anything).Return(&DefaultFalseResp, nil)
testcase = testCase{
name: "[share] collection not found",
mp: mp,
exceptCode: 200,
expectedBody: PrintErr(merr.ErrCollectionNotFound),
}
case ReturnFail:
call = mp.EXPECT().HasCollection(mock.Anything, mock.Anything).Return(nil, ErrDefault)
testcase = testCase{
name: "[share] check collection fail",
mp: mp,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
}
case ReturnWrongStatus:
reason := "can't find collection: " + DefaultCollectionName
call = mp.EXPECT().HasCollection(mock.Anything, mock.Anything).Return(&milvuspb.BoolResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError, // 1
Reason: reason,
},
}, nil)
testcase = testCase{
name: "[share] unexpected error",
mp: mp,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_UnexpectedError), reason),
}
}
if times == 2 {
call.Twice()
} else {
call.Once()
}
if testCases != nil {
testCases = append(testCases, testcase)
}
return mp, testCases
}
func TestHttpRequestFormat(t *testing.T) {
paramtable.Init()
errStrs := []string{
PrintErr(merr.ErrIncorrectParameterFormat),
PrintErr(merr.ErrMissingRequiredParameters),
PrintErr(merr.ErrMissingRequiredParameters),
PrintErr(merr.ErrMissingRequiredParameters),
}
requestJsons := [][]byte{
[]byte(`{"collectionName": {"` + DefaultCollectionName + `", "dimension": 2}`),
[]byte(`{"collName": "` + DefaultCollectionName + `", "dimension": 2}`),
[]byte(`{"collName": "` + DefaultCollectionName + `", "dim": 2}`),
[]byte(`{"collectionName": "` + DefaultCollectionName + `"}`),
}
paths := [][]string{
{
URIPrefix + VectorCollectionsCreatePath,
URIPrefix + VectorCollectionsDropPath,
URIPrefix + VectorGetPath,
URIPrefix + VectorSearchPath,
URIPrefix + VectorQueryPath,
URIPrefix + VectorInsertPath,
URIPrefix + VectorDeletePath,
}, {
URIPrefix + VectorCollectionsDropPath,
URIPrefix + VectorGetPath,
URIPrefix + VectorSearchPath,
URIPrefix + VectorQueryPath,
URIPrefix + VectorInsertPath,
URIPrefix + VectorDeletePath,
}, {
URIPrefix + VectorCollectionsCreatePath,
}, {
URIPrefix + VectorGetPath,
URIPrefix + VectorSearchPath,
URIPrefix + VectorQueryPath,
URIPrefix + VectorInsertPath,
URIPrefix + VectorDeletePath,
},
}
for i, pathArr := range paths {
for _, path := range pathArr {
t.Run("request parameters wrong", func(t *testing.T) {
testEngine := initHTTPServer(mocks.NewMockProxy(t), true)
bodyReader := bytes.NewReader(requestJsons[i])
req := httptest.NewRequest(http.MethodPost, path, bodyReader)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, 200)
assert.Equal(t, w.Body.String(), errStrs[i])
})
}
}
}
func TestAuthorization(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
errorStr := Print(Code(merr.ErrServiceUnavailable), "internal: Milvus Proxy is not ready yet. please wait: service unavailable")
jsons := map[string][]byte{
errorStr: []byte(`{"collectionName": "` + DefaultCollectionName + `", "vector": [0.1, 0.2], "filter": "id in [2]", "id": [2], "dimension": 2, "data":[{"book_id":1,"book_intro":[0.1,0.11],"distance":0.01,"word_count":1000},{"book_id":2,"book_intro":[0.2,0.22],"distance":0.04,"word_count":2000},{"book_id":3,"book_intro":[0.3,0.33],"distance":0.09,"word_count":3000}]}`),
}
paths := map[string][]string{
errorStr: {
URIPrefix + VectorGetPath,
URIPrefix + VectorInsertPath,
URIPrefix + VectorDeletePath,
},
}
for res, pathArr := range paths {
for _, path := range pathArr {
t.Run("proxy is not ready", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServer(mp, true)
bodyReader := bytes.NewReader(jsons[res])
req := httptest.NewRequest(http.MethodPost, path, bodyReader)
req.Header.Set("authorization", "Bearer test:test")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusForbidden)
assert.Equal(t, w.Body.String(), res)
})
}
}
paths = map[string][]string{
errorStr: {
URIPrefix + VectorCollectionsCreatePath,
},
}
for res, pathArr := range paths {
for _, path := range pathArr {
t.Run("proxy is not ready", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServer(mp, true)
bodyReader := bytes.NewReader(jsons[res])
req := httptest.NewRequest(http.MethodPost, path, bodyReader)
req.Header.Set("authorization", "Bearer test:test")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusForbidden)
assert.Equal(t, w.Body.String(), res)
})
}
}
paths = map[string][]string{
errorStr: {
URIPrefix + VectorCollectionsDropPath,
},
}
for res, pathArr := range paths {
for _, path := range pathArr {
t.Run("proxy is not ready", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServer(mp, true)
bodyReader := bytes.NewReader(jsons[res])
req := httptest.NewRequest(http.MethodPost, path, bodyReader)
req.Header.Set("authorization", "Bearer test:test")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusForbidden)
assert.Equal(t, w.Body.String(), res)
})
}
}
paths = map[string][]string{
errorStr: {
URIPrefix + VectorCollectionsPath,
URIPrefix + VectorCollectionsDescribePath + "?collectionName=" + DefaultCollectionName,
},
}
for res, pathArr := range paths {
for _, path := range pathArr {
t.Run("proxy is not ready", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServer(mp, true)
req := httptest.NewRequest(http.MethodGet, path, nil)
req.Header.Set("authorization", "Bearer test:test")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusForbidden)
assert.Equal(t, w.Body.String(), res)
})
}
}
paths = map[string][]string{
errorStr: {
URIPrefix + VectorQueryPath,
URIPrefix + VectorSearchPath,
},
}
for res, pathArr := range paths {
for _, path := range pathArr {
t.Run("proxy is not ready", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServer(mp, true)
bodyReader := bytes.NewReader(jsons[res])
req := httptest.NewRequest(http.MethodPost, path, bodyReader)
req.Header.Set("authorization", "Bearer test:test")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusForbidden)
assert.Equal(t, w.Body.String(), res)
})
}
}
}
func TestDatabaseNotFound(t *testing.T) {
paramtable.Init()
t.Run("list database fail", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(nil, ErrDefault).Once()
testEngine := initHTTPServer(mp, true)
req := httptest.NewRequest(http.MethodGet, URIPrefix+VectorCollectionsPath+"?dbName=test", nil)
req.Header.Set("authorization", "Bearer root:Milvus")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), PrintErr(ErrDefault))
})
t.Run("list database without success code", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(&milvuspb.ListDatabasesResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_UnexpectedError,
Reason: "",
},
}, nil).Once()
testEngine := initHTTPServer(mp, true)
req := httptest.NewRequest(http.MethodGet, URIPrefix+VectorCollectionsPath+"?dbName=test", nil)
req.Header.Set("authorization", "Bearer root:Milvus")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), Print(int32(commonpb.ErrorCode_UnexpectedError), ""))
})
t.Run("list database success", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(&milvuspb.ListDatabasesResponse{
Status: &StatusSuccess,
DbNames: []string{"default", "test"},
}, nil).Once()
mp.EXPECT().
ShowCollections(mock.Anything, mock.Anything).
Return(&milvuspb.ShowCollectionsResponse{
Status: &StatusSuccess,
CollectionNames: nil,
}, nil).Once()
testEngine := initHTTPServer(mp, true)
req := httptest.NewRequest(http.MethodGet, URIPrefix+VectorCollectionsPath+"?dbName=test", nil)
req.Header.Set("authorization", "Bearer root:Milvus")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), "{\"code\":200,\"data\":[]}")
})
errorStr := PrintErr(merr.ErrDatabaseNotFound)
paths := map[string][]string{
errorStr: {
URIPrefix + VectorCollectionsPath + "?dbName=test",
URIPrefix + VectorCollectionsDescribePath + "?dbName=test&collectionName=" + DefaultCollectionName,
},
}
for res, pathArr := range paths {
for _, path := range pathArr {
t.Run("GET dbName", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(&milvuspb.ListDatabasesResponse{
Status: &StatusSuccess,
DbNames: []string{"default"},
}, nil).Once()
testEngine := initHTTPServer(mp, true)
req := httptest.NewRequest(http.MethodGet, path, nil)
req.Header.Set("authorization", "Bearer root:Milvus")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), res)
})
}
}
requestBody := `{"dbName": "test", "collectionName": "` + DefaultCollectionName + `", "vector": [0.1, 0.2], "filter": "id in [2]", "id": [2], "dimension": 2, "data":[{"book_id":1,"book_intro":[0.1,0.11],"distance":0.01,"word_count":1000},{"book_id":2,"book_intro":[0.2,0.22],"distance":0.04,"word_count":2000},{"book_id":3,"book_intro":[0.3,0.33],"distance":0.09,"word_count":3000}]}`
paths = map[string][]string{
requestBody: {
URIPrefix + VectorCollectionsCreatePath,
URIPrefix + VectorCollectionsDropPath,
URIPrefix + VectorInsertPath,
URIPrefix + VectorDeletePath,
URIPrefix + VectorQueryPath,
URIPrefix + VectorGetPath,
URIPrefix + VectorSearchPath,
},
}
for request, pathArr := range paths {
for _, path := range pathArr {
t.Run("POST dbName", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(&milvuspb.ListDatabasesResponse{
Status: &StatusSuccess,
DbNames: []string{"default"},
}, nil).Once()
testEngine := initHTTPServer(mp, true)
bodyReader := bytes.NewReader([]byte(request))
req := httptest.NewRequest(http.MethodPost, path, bodyReader)
req.Header.Set("authorization", "Bearer root:Milvus")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), errorStr)
})
}
}
}
func wrapWithDescribeIndex(t *testing.T, mp *mocks.MockProxy, returnType int, times int, testCases []testCase) (*mocks.MockProxy, []testCase) {
if mp == nil {
mp = mocks.NewMockProxy(t)
}
var call *mocks.MockProxy_DescribeIndex_Call
var testcase testCase
switch returnType {
case ReturnSuccess:
call = mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&DefaultDescIndexesReqp, nil)
testcase = testCase{
name: "[share] describe coll success",
mp: mp,
exceptCode: 200,
expectedBody: "{\"code\":200,\"data\":{}}",
}
case ReturnFail:
call = mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(nil, ErrDefault)
testcase = testCase{
name: "[share] describe index fail",
mp: mp,
exceptCode: 200,
expectedBody: PrintErr(ErrDefault),
}
case ReturnWrongStatus:
reason := "index is not exists"
call = mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&milvuspb.DescribeIndexResponse{Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_IndexNotExist,
Reason: reason,
}}, nil)
testcase = testCase{
name: "[share] describe index fail",
mp: mp,
exceptCode: 200,
expectedBody: Print(int32(commonpb.ErrorCode_IndexNotExist), reason),
}
}
if times == 2 {
call.Twice()
} else {
call.Once()
}
if testCases != nil {
testCases = append(testCases, testcase)
}
return mp, testCases
}
func getCollectionSchema(collectionName string) *schemapb.CollectionSchema {
sch := &schemapb.CollectionSchema{
Name: collectionName,
AutoID: false,
}
sch.Fields = getFieldSchema()
return sch
}
func getFieldSchema() []*schemapb.FieldSchema {
fields := []*schemapb.FieldSchema{
{
FieldID: 0,
Name: "RowID",
Description: "RowID field",
DataType: schemapb.DataType_Int64,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "f0_tk1",
Value: "f0_tv1",
},
},
},
{
FieldID: 1,
Name: "Timestamp",
Description: "Timestamp field",
DataType: schemapb.DataType_Int64,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "f1_tk1",
Value: "f1_tv1",
},
},
},
{
FieldID: 100,
Name: "float_vector_field",
Description: "field 100",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{
Key: "dim",
Value: "2",
},
},
IndexParams: []*commonpb.KeyValuePair{
{
Key: "indexkey",
Value: "indexvalue",
},
},
},
{
FieldID: 101,
Name: "int64_field",
Description: "field 106",
DataType: schemapb.DataType_Int64,
TypeParams: []*commonpb.KeyValuePair{},
IndexParams: []*commonpb.KeyValuePair{},
IsPrimaryKey: true,
},
}
return fields
}
func Test_Handles_VectorCollectionsDescribe(t *testing.T) {
paramtable.Init()
mp := mocks.NewMockProxy(t)
h := NewHandlers(mp)
testEngine := gin.New()
app := testEngine.Group("/", func(c *gin.Context) {
username, _, ok := ParseUsernamePassword(c)
if ok {
c.Set(ContextUsername, username)
}
})
h.RegisterRoutesToV1(app)
t.Run("hasn't authenticate", func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, "/vector/collections/describe?collectionName=book", nil)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusUnauthorized)
})
t.Run("auth fail", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
req := httptest.NewRequest(http.MethodGet, "/vector/collections/describe?collectionName=book", nil)
req.SetBasicAuth("test", util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusForbidden)
assert.Equal(t, w.Body.String(), Print(Code(merr.ErrServiceUnavailable), "internal: Milvus Proxy is not ready yet. please wait: service unavailable"))
})
t.Run("describe collection fail with error", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "false")
mp.EXPECT().
DescribeCollection(mock.Anything, mock.Anything).
Return(nil, ErrDefault).
Once()
req := httptest.NewRequest(http.MethodGet, "/vector/collections/describe?collectionName=book", nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), PrintErr(ErrDefault))
})
t.Run("describe collection fail with status code", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "false")
mp.EXPECT().
DescribeCollection(mock.Anything, mock.Anything).
Return(&milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
}, nil).
Once()
req := httptest.NewRequest(http.MethodGet, "/vector/collections/describe?collectionName=book", nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), "{\"code\":1,\"message\":\"\"}")
})
t.Run("get load state and describe index fail with error", func(t *testing.T) {
mp.EXPECT().
DescribeCollection(mock.Anything, mock.Anything).
Return(&milvuspb.DescribeCollectionResponse{
Schema: getCollectionSchema("collectionName"),
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil).
Once()
mp.EXPECT().
GetLoadState(mock.Anything, mock.Anything).
Return(nil, errors.New("error")).
Once()
mp.EXPECT().
DescribeIndex(mock.Anything, mock.Anything).
Return(nil, errors.New("error")).
Once()
req := httptest.NewRequest(http.MethodGet, "/vector/collections/describe?collectionName=book", nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), "{\"code\":200,\"data\":{\"collectionName\":\"\",\"description\":\"\",\"enableDynamic\":false,\"fields\":[{\"autoId\":false,\"description\":\"RowID field\",\"name\":\"RowID\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"Timestamp field\",\"name\":\"Timestamp\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"field 100\",\"name\":\"float_vector_field\",\"primaryKey\":false,\"type\":\"FloatVector(2)\"},{\"autoId\":false,\"description\":\"field 106\",\"name\":\"int64_field\",\"primaryKey\":true,\"type\":\"Int64\"}],\"indexes\":[],\"load\":\"\",\"shardsNum\":0}}")
})
t.Run("get load state and describe index fail with status code", func(t *testing.T) {
mp.EXPECT().
DescribeCollection(mock.Anything, mock.Anything).
Return(&milvuspb.DescribeCollectionResponse{
Schema: getCollectionSchema("collectionName"),
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil).
Once()
mp.EXPECT().
GetLoadState(mock.Anything, mock.Anything).
Return(&milvuspb.GetLoadStateResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
}, nil).
Once()
mp.EXPECT().
DescribeIndex(mock.Anything, mock.Anything).
Return(&milvuspb.DescribeIndexResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
}, nil).
Once()
req := httptest.NewRequest(http.MethodGet, "/vector/collections/describe?collectionName=book", nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), "{\"code\":200,\"data\":{\"collectionName\":\"\",\"description\":\"\",\"enableDynamic\":false,\"fields\":[{\"autoId\":false,\"description\":\"RowID field\",\"name\":\"RowID\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"Timestamp field\",\"name\":\"Timestamp\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"field 100\",\"name\":\"float_vector_field\",\"primaryKey\":false,\"type\":\"FloatVector(2)\"},{\"autoId\":false,\"description\":\"field 106\",\"name\":\"int64_field\",\"primaryKey\":true,\"type\":\"Int64\"}],\"indexes\":[],\"load\":\"\",\"shardsNum\":0}}")
})
t.Run("ok", func(t *testing.T) {
mp.EXPECT().
DescribeCollection(mock.Anything, mock.Anything).
Return(&milvuspb.DescribeCollectionResponse{
Schema: getCollectionSchema("collectionName"),
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil).
Once()
mp.EXPECT().
GetLoadState(mock.Anything, mock.Anything).
Return(&milvuspb.GetLoadStateResponse{
State: commonpb.LoadState_LoadStateLoaded,
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil).
Once()
mp.EXPECT().
DescribeIndex(mock.Anything, mock.Anything).
Return(&milvuspb.DescribeIndexResponse{
IndexDescriptions: []*milvuspb.IndexDescription{
{
IndexName: "in",
FieldName: "fn",
Params: []*commonpb.KeyValuePair{
{
Key: "metric_type",
Value: "L2",
},
},
},
},
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
}, nil).
Once()
req := httptest.NewRequest(http.MethodGet, "/vector/collections/describe?collectionName=book", nil)
req.SetBasicAuth(util.UserRoot, util.DefaultRootPassword)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, w.Code, http.StatusOK)
assert.Equal(t, w.Body.String(), "{\"code\":200,\"data\":{\"collectionName\":\"\",\"description\":\"\",\"enableDynamic\":false,\"fields\":[{\"autoId\":false,\"description\":\"RowID field\",\"name\":\"RowID\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"Timestamp field\",\"name\":\"Timestamp\",\"primaryKey\":false,\"type\":\"Int64\"},{\"autoId\":false,\"description\":\"field 100\",\"name\":\"float_vector_field\",\"primaryKey\":false,\"type\":\"FloatVector(2)\"},{\"autoId\":false,\"description\":\"field 106\",\"name\":\"int64_field\",\"primaryKey\":true,\"type\":\"Int64\"}],\"indexes\":[],\"load\":\"LoadStateLoaded\",\"shardsNum\":0}}")
})
}