wehub/milvus-io--milvus
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
milvus-io--milvus/internal/distributed/proxy/httpserver/handler_v2_test.go
T
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Details
Publish Builder / amazonlinux2023 (push) Failing after 1s
Details
Build and test / UT for Go (push) Has been skipped
Details
Publish KRTE Images / KRTE (push) Failing after 1s
Details
Build and test / Integration Test (push) Has been skipped
Details
Build and test / Upload Code Coverage (push) Has been skipped
Details
Publish Builder / rockylinux9 (push) Failing after 1s
Details
Publish Builder / ubuntu22.04 (push) Failing after 0s
Details
Publish Builder / ubuntu24.04 (push) Failing after 0s
Details
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Details
Publish Test Images / PyTest (push) Failing after 0s
Details
Build and test / UT for Cpp (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

5313 lines
219 KiB
Go
Raw Permalink Blame History

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package httpserver
import (
"bytes"
"context"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
mhttp "github.com/milvus-io/milvus/internal/http"
"github.com/milvus-io/milvus/internal/json"
"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/internal/util/indexparamcheck"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
"github.com/milvus-io/milvus/pkg/v3/util"
"github.com/milvus-io/milvus/pkg/v3/util/crypto"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
const (
DefaultPartitionName = "_default"
)
type rawTestCase struct {
path string
errMsg string
errCode int32
}
type requestBodyTestCase struct {
path string
requestBody []byte
errMsg string
errCode int32
}
type DefaultReq struct {
DbName string `json:"dbName"`
}
func (DefaultReq) GetBase() *commonpb.MsgBase {
return &commonpb.MsgBase{}
}
func (req *DefaultReq) GetDbName() string { return req.DbName }
func init() {
paramtable.Init()
streaming.SetupNoopWALForTest()
}
func sendReqAndVerify(t *testing.T, testEngine *gin.Engine, testName, method string, testcase requestBodyTestCase) {
t.Run(testName, func(t *testing.T) {
req := httptest.NewRequest(method, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Contains(t, returnBody.Message, testcase.errMsg)
}
})
}
func TestTraceLogRequestFieldRedactsRESTSnapshotExternalSpec(t *testing.T) {
externalSpec := `{"extfs":{"cloud_provider":"aws","access_key_id":"AKIAEXAMPLE","access_key_value":"SUPERSECRET","region":"us-west-2"}}`
testCases := []struct {
name string
req any
}{
{
name: "restore external snapshot",
req: &RestoreExternalSnapshotReq{
DbName: "db",
TargetCollectionName: "restored",
SnapshotMetadataURI: "s3://bucket/export-root/snapshots/100/metadata/1.json",
ExternalSpec: externalSpec,
},
},
{
name: "export snapshot",
req: &ExportSnapshotReq{
DbName: "db",
CollectionName: "source",
Name: "snapshot",
TargetS3Path: "s3://bucket/export-root",
ExternalSpec: externalSpec,
},
},
}
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
field := getTraceLogRequestFieldWithoutSensitiveInfo(testCase.req)
request := fmt.Sprint(field.Interface)
assert.NotContains(t, request, "AKIAEXAMPLE")
assert.NotContains(t, request, "SUPERSECRET")
assert.Contains(t, request, "***")
})
}
}
func TestHTTPWrapper(t *testing.T) {
postTestCases := []requestBodyTestCase{}
postTestCasesTrace := []requestBodyTestCase{}
ginHandler := gin.Default()
app := ginHandler.Group("", genAuthMiddleWare(false))
path := "/wrapper/post"
app.POST(path, wrapperPost(func() any { return &DefaultReq{} }, func(ctx context.Context, c *gin.Context, req any, dbName string) (interface{}, error) {
return nil, nil
}))
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{}`),
})
path = "/wrapper/post/param"
app.POST(path, wrapperPost(func() any { return &CollectionNameReq{} }, func(ctx context.Context, c *gin.Context, req any, dbName string) (interface{}, error) {
return nil, nil
}))
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `"}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{}`),
errMsg: "missing required parameters, error: Key: 'CollectionNameReq.CollectionName' Error:Field validation for 'CollectionName' failed on the 'required' tag",
errCode: 1802, // ErrMissingRequiredParameters
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(``),
errMsg: "can only accept json format request, the request body should be nil, however {} is valid",
errCode: 1801, // ErrIncorrectParameterFormat
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "book", "dbName"}`),
errMsg: "can only accept json format request, error: invalid character '}' after object key",
errCode: 1801, // ErrIncorrectParameterFormat
})
path = "/wrapper/post/trace"
app.POST(path, wrapperPost(func() any { return &DefaultReq{} }, wrapperTraceLog(func(ctx context.Context, c *gin.Context, req any, dbName string) (interface{}, error) {
return nil, nil
})))
postTestCasesTrace = append(postTestCasesTrace, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `"}`),
})
path = "/wrapper/post/trace/wrong"
app.POST(path, wrapperPost(func() any { return &DefaultReq{} }, wrapperTraceLog(func(ctx context.Context, c *gin.Context, req any, dbName string) (interface{}, error) {
return nil, merr.ErrCollectionNotFound
})))
postTestCasesTrace = append(postTestCasesTrace, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `"}`),
})
path = "/wrapper/post/trace/call"
app.POST(path, wrapperPost(func() any { return &DefaultReq{} }, wrapperTraceLog(func(ctx context.Context, c *gin.Context, req any, dbName string) (interface{}, error) {
return wrapperProxy(ctx, c, req, false, false, "", func(reqctx context.Context, req any) (any, error) {
return nil, nil
})
})))
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `"}`),
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
if testcase.errCode != 0 {
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
fmt.Println(w.Body.String())
})
}
for _, i := range []string{"1", "2", "3"} {
paramtable.Get().Save(proxy.Params.CommonCfg.TraceLogMode.Key, i)
for _, testcase := range postTestCasesTrace {
t.Run("post"+testcase.path+"["+i+"]", func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
if testcase.errCode != 0 {
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
assert.Equal(t, testcase.errCode, returnBody.Code)
}
fmt.Println(w.Body.String())
})
}
}
}
func TestGrpcWrapper(t *testing.T) {
getTestCases := []rawTestCase{}
getTestCasesNeedAuth := []rawTestCase{}
needAuthPrefix := "/auth"
ginHandler := gin.Default()
app := ginHandler.Group("")
appNeedAuth := ginHandler.Group(needAuthPrefix, genAuthMiddleWare(true))
path := "/wrapper/grpc/-0"
handle := func(reqctx context.Context, req any) (any, error) {
return nil, nil
}
app.GET(path, func(c *gin.Context) {
ctx := proxy.NewContextWithMetadata(c, "", DefaultDbName)
wrapperProxy(ctx, c, &DefaultReq{}, false, false, "", handle)
})
appNeedAuth.GET(path, func(c *gin.Context) {
username, _ := c.Get(ContextUsername)
ctx := proxy.NewContextWithMetadata(c, username.(string), DefaultDbName)
wrapperProxy(ctx, c, &milvuspb.DescribeCollectionRequest{}, true, false, "", handle)
})
getTestCases = append(getTestCases, rawTestCase{
path: path,
})
getTestCasesNeedAuth = append(getTestCasesNeedAuth, rawTestCase{
path: needAuthPrefix + path,
})
path = "/wrapper/grpc/01"
handle = func(reqctx context.Context, req any) (any, error) {
return nil, merr.ErrNeedAuthenticate // 1800
}
app.GET(path, func(c *gin.Context) {
ctx := proxy.NewContextWithMetadata(c, "", DefaultDbName)
wrapperProxy(ctx, c, &DefaultReq{}, false, false, "", handle)
})
appNeedAuth.GET(path, func(c *gin.Context) {
username, _ := c.Get(ContextUsername)
ctx := proxy.NewContextWithMetadata(c, username.(string), DefaultDbName)
wrapperProxy(ctx, c, &milvuspb.DescribeCollectionRequest{}, true, false, "", handle)
})
getTestCases = append(getTestCases, rawTestCase{
path: path,
errCode: 65535,
})
getTestCasesNeedAuth = append(getTestCasesNeedAuth, rawTestCase{
path: needAuthPrefix + path,
})
path = "/wrapper/grpc/00"
handle = func(reqctx context.Context, req any) (any, error) {
return &milvuspb.BoolResponse{
Status: commonSuccessStatus,
}, nil
}
app.GET(path, func(c *gin.Context) {
ctx := proxy.NewContextWithMetadata(c, "", DefaultDbName)
wrapperProxy(ctx, c, &DefaultReq{}, false, false, "", handle)
})
appNeedAuth.GET(path, func(c *gin.Context) {
username, _ := c.Get(ContextUsername)
ctx := proxy.NewContextWithMetadata(c, username.(string), DefaultDbName)
wrapperProxy(ctx, c, &milvuspb.DescribeCollectionRequest{}, true, false, "", handle)
})
getTestCases = append(getTestCases, rawTestCase{
path: path,
})
getTestCasesNeedAuth = append(getTestCasesNeedAuth, rawTestCase{
path: needAuthPrefix + path,
})
path = "/wrapper/grpc/10"
handle = func(reqctx context.Context, req any) (any, error) {
return &milvuspb.BoolResponse{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNameNotFound, // 28
Reason: "",
},
}, nil
}
app.GET(path, func(c *gin.Context) {
ctx := proxy.NewContextWithMetadata(c, "", DefaultDbName)
wrapperProxy(ctx, c, &DefaultReq{}, false, false, "", handle)
})
appNeedAuth.GET(path, func(c *gin.Context) {
username, _ := c.Get(ContextUsername)
ctx := proxy.NewContextWithMetadata(c, username.(string), DefaultDbName)
wrapperProxy(ctx, c, &milvuspb.DescribeCollectionRequest{}, true, false, "", handle)
})
getTestCases = append(getTestCases, rawTestCase{
path: path,
errCode: 65535,
})
getTestCasesNeedAuth = append(getTestCasesNeedAuth, rawTestCase{
path: needAuthPrefix + path,
})
for _, testcase := range getTestCases {
t.Run("get"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, testcase.path, nil)
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
if testcase.errCode != 0 {
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
fmt.Println(w.Body.String())
})
}
for _, testcase := range getTestCasesNeedAuth {
t.Run("get"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, testcase.path, nil)
req.SetBasicAuth(util.UserRoot, getDefaultRootPassword())
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
if testcase.errCode != 0 {
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
fmt.Println(w.Body.String())
})
}
path = "/wrapper/grpc/auth"
app.GET(path, func(c *gin.Context) {
wrapperProxy(context.Background(), c, &milvuspb.DescribeCollectionRequest{}, true, false, "", handle)
})
appNeedAuth.GET(path, func(c *gin.Context) {
ctx := proxy.NewContextWithMetadata(c, "test", DefaultDbName)
wrapperProxy(ctx, c, &milvuspb.LoadCollectionRequest{}, true, false, "", handle)
})
t.Run("check authorization", func(t *testing.T) {
req := httptest.NewRequest(http.MethodGet, path, nil)
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusUnauthorized, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, int32(1800), returnBody.Code)
assert.Equal(t, "user hasn't authenticated", returnBody.Message)
fmt.Println(w.Body.String())
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
req = httptest.NewRequest(http.MethodGet, needAuthPrefix+path, nil)
req.SetBasicAuth("test", getDefaultRootPassword())
w = httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusForbidden, w.Code)
err = json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, int32(65535), returnBody.Code)
assert.Equal(t, "rpc error: code = PermissionDenied desc = PrivilegeLoad: permission deny to test in the `default` database", returnBody.Message)
fmt.Println(w.Body.String())
})
}
type headerTestCase struct {
path string
headers map[string]string
status int
}
func TestTimeout(t *testing.T) {
headerTestCases := []headerTestCase{}
ginHandler := gin.Default()
app := ginHandler.Group("")
path := "/middleware/timeout/5"
app.POST(path, timeoutMiddleware(func(c *gin.Context) {
time.Sleep(5 * time.Second)
}))
headerTestCases = append(headerTestCases, headerTestCase{
path: path, // wait 5s
})
headerTestCases = append(headerTestCases, headerTestCase{
path: path, // timeout 3s
headers: map[string]string{mhttp.HTTPHeaderRequestTimeout: "3"},
status: http.StatusRequestTimeout,
})
path = "/middleware/timeout/31"
app.POST(path, timeoutMiddleware(func(c *gin.Context) {
time.Sleep(31 * time.Second)
}))
headerTestCases = append(headerTestCases, headerTestCase{
path: path, // timeout 30s
status: http.StatusRequestTimeout,
})
headerTestCases = append(headerTestCases, headerTestCase{
path: path, // wait 32s
headers: map[string]string{mhttp.HTTPHeaderRequestTimeout: "32"},
})
for _, testcase := range headerTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, nil)
for key, value := range testcase.headers {
req.Header.Set(key, value)
}
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
if testcase.status == 0 {
assert.Equal(t, http.StatusOK, w.Code)
} else {
assert.Equal(t, testcase.status, w.Code)
}
fmt.Println(w.Body.String())
})
}
}
func TestTimeoutResponseRecorderIsolation(t *testing.T) {
realResponse := httptest.NewRecorder()
realCtx, _ := gin.CreateTestContext(realResponse)
buffer := &bytes.Buffer{}
recorder := newTimeoutResponseRecorder(buffer)
recorder.Header().Set("X-Test", "value")
recorder.WriteHeader(http.StatusAccepted)
recorder.WriteHeaderNow()
recorder.Flush()
n, err := recorder.WriteString("body")
assert.NoError(t, err)
assert.Equal(t, 4, n)
assert.Empty(t, realResponse.Body.String())
assert.Empty(t, realResponse.Header().Get("X-Test"))
assert.Equal(t, http.StatusAccepted, recorder.Status())
assert.Equal(t, 4, recorder.Size())
assert.True(t, recorder.Written())
assert.NoError(t, recorder.CommitTo(realCtx.Writer))
assert.Equal(t, http.StatusAccepted, realResponse.Code)
assert.Equal(t, "value", realResponse.Header().Get("X-Test"))
assert.Equal(t, "body", realResponse.Body.String())
recorder.CloseForTimeout()
n, err = recorder.Write([]byte("late"))
assert.Error(t, err)
assert.Equal(t, 0, n)
}
func TestTimeoutMiddlewareCommitsBufferedResponsesAndMetadata(t *testing.T) {
testCases := []struct {
name string
status int
code int32
write func(c *gin.Context, code int32, message string)
}{
{
name: "return",
status: http.StatusCreated,
code: 11,
write: func(c *gin.Context, code int32, message string) {
HTTPReturn(c, http.StatusCreated, gin.H{HTTPReturnCode: code, HTTPReturnMessage: message})
},
},
{
name: "stream",
status: http.StatusAccepted,
code: 12,
write: func(c *gin.Context, code int32, message string) {
HTTPReturnStream(c, http.StatusAccepted, gin.H{HTTPReturnCode: code, HTTPReturnMessage: message})
},
},
{
name: "abort",
status: http.StatusNonAuthoritativeInfo,
code: 13,
write: func(c *gin.Context, code int32, message string) {
HTTPAbortReturn(c, http.StatusNonAuthoritativeInfo, gin.H{HTTPReturnCode: code, HTTPReturnMessage: message})
},
},
}
for _, testcase := range testCases {
t.Run(testcase.name, func(t *testing.T) {
ginHandler := gin.New()
originalCtx := make(chan *gin.Context, 1)
path := "/middleware/timeout/commit/" + testcase.name
ginHandler.Use(func(c *gin.Context) {
c.Set(ContextUsername, "root")
originalCtx <- c
c.Next()
})
ginHandler.POST(path, timeoutMiddleware(func(c *gin.Context) {
username, ok := c.Get(ContextUsername)
assert.True(t, ok)
assert.Equal(t, "root", username)
c.Header("X-Test", testcase.name)
c.Set(ContextRequest, "request-"+testcase.name)
c.Set(ContextResponse, "response-"+testcase.name)
c.Set("traceID", "trace-"+testcase.name)
testcase.write(c, testcase.code, testcase.name)
}))
req := httptest.NewRequest(http.MethodPost, path, nil)
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, testcase.status, w.Code)
assert.Equal(t, testcase.name, w.Header().Get("X-Test"))
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, testcase.code, returnBody.Code)
assert.Equal(t, testcase.name, returnBody.Message)
ctx := <-originalCtx
value, ok := ctx.Get(HTTPReturnCode)
assert.True(t, ok)
assert.Equal(t, testcase.code, value)
value, ok = ctx.Get(HTTPReturnMessage)
assert.True(t, ok)
assert.Equal(t, testcase.name, value)
value, ok = ctx.Get(ContextRequest)
assert.True(t, ok)
assert.Equal(t, "request-"+testcase.name, value)
value, ok = ctx.Get(ContextResponse)
assert.True(t, ok)
assert.Equal(t, "response-"+testcase.name, value)
value, ok = ctx.Get("traceID")
assert.True(t, ok)
assert.Equal(t, "trace-"+testcase.name, value)
})
}
}
func TestTimeoutMiddlewarePropagatesAbortToOriginalContext(t *testing.T) {
ginHandler := gin.New()
path := "/middleware/timeout/abort-propagation"
nextCalled := false
ginHandler.POST(path, timeoutMiddleware(func(c *gin.Context) {
HTTPAbortReturn(c, http.StatusOK, gin.H{HTTPReturnCode: int32(100), HTTPReturnMessage: "aborted"})
}), func(c *gin.Context) {
nextCalled = true
HTTPReturn(c, http.StatusOK, gin.H{HTTPReturnCode: int32(200), HTTPReturnMessage: "next"})
})
req := httptest.NewRequest(http.MethodPost, path, nil)
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.False(t, nextCalled)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(100), returnBody.Code)
assert.Equal(t, "aborted", returnBody.Message)
}
func TestTimeoutMiddlewareLateHandlerWritesUseCopiedContext(t *testing.T) {
paramtable.Get().Save(paramtable.Get().HTTPCfg.RequestTimeoutMs.Key, "10")
t.Cleanup(func() {
paramtable.Get().Reset(paramtable.Get().HTTPCfg.RequestTimeoutMs.Key)
})
ginHandler := gin.New()
path := "/middleware/timeout/late-write"
originalCtx := make(chan *gin.Context, 1)
handlerCtx := make(chan *gin.Context, 1)
lateWriteDone := make(chan struct{}, 1)
ginHandler.Use(func(c *gin.Context) {
originalCtx <- c
c.Next()
})
ginHandler.POST(path, timeoutMiddleware(func(c *gin.Context) {
handlerCtx <- c
<-c.Request.Context().Done()
HTTPAbortReturn(c, http.StatusOK, gin.H{HTTPReturnCode: int32(12345), HTTPReturnMessage: "late write"})
lateWriteDone <- struct{}{}
}))
req := httptest.NewRequest(http.MethodPost, path, nil)
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusRequestTimeout, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, merr.TimeoutCode, returnBody.Code)
assert.Equal(t, "request timeout", returnBody.Message)
assert.NotContains(t, w.Body.String(), "late write")
ctx := <-originalCtx
copiedCtx := <-handlerCtx
assert.False(t, ctx == copiedCtx)
select {
case <-lateWriteDone:
case <-time.After(time.Second):
t.Fatal("late handler write did not finish")
}
value, ok := ctx.Get(HTTPReturnCode)
assert.True(t, ok)
assert.Equal(t, merr.TimeoutCode, value)
value, ok = ctx.Get(HTTPReturnMessage)
assert.True(t, ok)
assert.Equal(t, "request timeout", value)
}
func TestRestfulSizeMiddlewarePreservesRequestContextCancel(t *testing.T) {
ginHandler := gin.New()
app := ginHandler.Group("")
path := "/middleware/restful-size/context-cancel"
requestCtx, cancel := context.WithCancel(context.Background())
defer cancel()
observed := make(chan error, 1)
app.POST(path, restfulSizeMiddleware(timeoutMiddleware(func(c *gin.Context) {
ctx := c.Request.Context()
cancel()
select {
case <-ctx.Done():
observed <- ctx.Err()
case <-time.After(100 * time.Millisecond):
observed <- errors.New("request context was not canceled")
}
}), false))
req := httptest.NewRequest(http.MethodPost, path, nil).WithContext(requestCtx)
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.True(t, errors.Is(<-observed, context.Canceled))
}
func TestTimeoutMiddlewarePassesDeadline(t *testing.T) {
ginHandler := gin.Default()
app := ginHandler.Group("")
path := "/middleware/timeout/deadline"
app.POST(path, timeoutMiddleware(func(c *gin.Context) {
deadline, ok := c.Request.Context().Deadline()
assert.True(t, ok)
assert.LessOrEqual(t, time.Until(deadline), 3*time.Second)
assert.Greater(t, time.Until(deadline), time.Second)
}))
req := httptest.NewRequest(http.MethodPost, path, nil)
req.Header.Set(mhttp.HTTPHeaderRequestTimeout, "3")
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
}
func TestTimeoutMiddlewareRejectsInvalidRequestTimeout(t *testing.T) {
for _, requestTimeout := range []string{"3.5", "abc"} {
t.Run(requestTimeout, func(t *testing.T) {
ginHandler := gin.New()
path := "/middleware/timeout/invalid"
called := make(chan struct{}, 1)
ginHandler.POST(path, timeoutMiddleware(func(c *gin.Context) {
called <- struct{}{}
HTTPReturn(c, http.StatusOK, gin.H{HTTPReturnCode: merr.Code(nil)})
}))
req := httptest.NewRequest(http.MethodPost, path, nil)
req.Header.Set(mhttp.HTTPHeaderRequestTimeout, requestTimeout)
w := httptest.NewRecorder()
ginHandler.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, merr.Code(merr.ErrParameterInvalid), returnBody.Code)
assert.Contains(t, returnBody.Message, mhttp.HTTPHeaderRequestTimeout)
assert.Contains(t, returnBody.Message, requestTimeout)
select {
case <-called:
t.Fatal("handler was invoked for invalid Request-Timeout")
default:
}
})
}
}
func TestDocInDocOutCreateCollection(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Times(1)
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(CollectionCategory, CreateAction)
const baseRequestBody = `{
"collectionName": "doc_in_doc_out_demo",
"schema": {
"autoId": false,
"enableDynamicField": false,
"fields": [
{
"fieldName": "my_id",
"dataType": "Int64",
"isPrimary": true
},
{
"fieldName": "document_content",
"dataType": "VarChar",
"elementTypeParams": {
"max_length": "9000"
}
},
{
"fieldName": "sparse_vector_1",
"dataType": "SparseFloatVector"
}
],
"functions": %s
}
}`
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(fmt.Sprintf(baseRequestBody, `[
{
"name": "bm25_fn_1",
"type": "BM25",
"inputFieldNames": ["document_content"],
"outputFieldNames": ["sparse_vector_1"]
}
]`)),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(fmt.Sprintf(baseRequestBody, `[
{
"name": "bm25_fn_1",
"type": "BM25_",
"inputFieldNames": ["document_content"],
"outputFieldNames": ["sparse_vector_1"]
}
]`)),
errMsg: "Unsupported function type: BM25_",
errCode: 1100,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(fmt.Sprintf(baseRequestBody, `[
{
"name": "bm25_fn_1",
"inputFieldNames": ["document_content"],
"outputFieldNames": ["sparse_vector_1"]
}
]`)),
errMsg: "Unsupported function type:", // unprovided function type is empty string
errCode: 1100,
})
for _, testcase := range postTestCases {
sendReqAndVerify(t, testEngine, "post"+testcase.path, http.MethodPost, testcase)
}
}
func TestDocInDocOutCreateCollectionQuickDisallowFunction(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(CollectionCategory, CreateAction)
const baseRequestBody = `{
"collectionName": "doc_in_doc_out_demo",
"dimension": 2,
"idType": "Varchar",
"schema": {
"autoId": false,
"enableDynamicField": false,
"functions": [
{
"name": "bm25_fn_1",
"type": "BM25",
"inputFieldNames": ["document_content"],
"outputFieldNames": ["sparse_vector_1"]
}
]
}
}`
testcase := requestBodyTestCase{
path: path,
requestBody: []byte(baseRequestBody),
errMsg: "functions are not supported for quickly create collection",
errCode: 1100,
}
sendReqAndVerify(t, testEngine, "post"+testcase.path, http.MethodPost, testcase)
}
func TestDocInDocOutDescribeCollection(t *testing.T) {
paramtable.Init()
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateDocInDocOutCollectionSchema(schemapb.DataType_Int64),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().GetLoadState(mock.Anything, mock.Anything).Return(&DefaultLoadStateResp, nil).Once()
mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&DefaultDescIndexesReqp, nil).Once()
mp.EXPECT().ListAliases(mock.Anything, mock.Anything).Return(&milvuspb.ListAliasesResponse{
Status: &StatusSuccess,
Aliases: []string{DefaultAliasName},
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
testcase := requestBodyTestCase{
path: versionalV2(CollectionCategory, DescribeAction),
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `"}`),
}
sendReqAndVerify(t, testEngine, testcase.path, http.MethodPost, testcase)
}
func TestDocInDocOutInsert(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServerV2(mp, false)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateDocInDocOutCollectionSchema(schemapb.DataType_Int64),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once()
testcase := requestBodyTestCase{
path: versionalV2(EntityCategory, InsertAction),
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "varchar_field": "some text"}]}`),
}
sendReqAndVerify(t, testEngine, testcase.path, http.MethodPost, testcase)
}
func TestDocInDocOutInsertInvalid(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServerV2(mp, false)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateDocInDocOutCollectionSchema(schemapb.DataType_Int64),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
// function output field data is now passed through to proxy for validation
insertErr := errors.Wrap(merr.ErrInvalidInsertData, "not allowed to provide data for BM25 function output field")
mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{
Status: merr.Status(insertErr),
}, nil).Once()
testcase := requestBodyTestCase{
path: versionalV2(EntityCategory, InsertAction),
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": {"1": 0.1}, "varchar_field": "some text"}]}`),
errCode: 1804,
errMsg: "not allowed to provide data for BM25 function output field",
}
sendReqAndVerify(t, testEngine, testcase.path, http.MethodPost, testcase)
}
func TestSearchWithRerank(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServerV2(mp, false)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, true, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{
TopK: int64(3),
OutputFields: []string{FieldWordCount},
FieldsData: generateFieldData(),
Ids: generateIDs(schemapb.DataType_Int64, 3),
Scores: DefaultScores,
}}, nil).Once()
testcase := requestBodyTestCase{
path: versionalV2(EntityCategory, SearchAction),
requestBody: []byte(`{
"collectionName": "book",
"data": [[0.1, 0.2]],
"limit": 4,
"outputFields": ["word_count"],
"functionScore": {
"functions": [{
"name": "testRank",
"description": "",
"type": "Rerank",
"inputFieldNames": ["FieldWordCount"],
"params": {"name": "decay"}
}]}}`),
}
sendReqAndVerify(t, testEngine, testcase.path, http.MethodPost, testcase)
}
func TestHybridSearchWithRerank(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServerV2(mp, false)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, true, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().HybridSearch(mock.Anything, mock.MatchedBy(func(req *milvuspb.HybridSearchRequest) bool {
return assert.ObjectsAreEqual([]string{"part_a"}, req.GetPartitionNames())
})).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{
TopK: int64(3),
OutputFields: []string{FieldWordCount},
FieldsData: generateFieldData(),
Ids: generateIDs(schemapb.DataType_Int64, 3),
Scores: DefaultScores,
}}, nil).Once()
queryTestCases := requestBodyTestCase{
path: versionalV2(EntityCategory, HybridSearchAction),
requestBody: []byte(`{"collectionName": "hello_milvus", "partitionNames": ["part_a"], "search": [{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3}, {"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3}], "functionScore": {"functions": [{"name": "testRank", "type": "Rerank", "inputFieldNames": ["FieldWordCount"], "params": {"name": "decay"}}]}}`),
}
sendReqAndVerify(t, testEngine, queryTestCases.path, http.MethodPost, queryTestCases)
}
func TestDocInDocOutSearch(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServerV2(mp, false)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateDocInDocOutCollectionSchema(schemapb.DataType_Int64),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{
TopK: int64(3),
OutputFields: []string{FieldWordCount},
FieldsData: generateFieldData(),
Ids: generateIDs(schemapb.DataType_Int64, 3),
Scores: DefaultScores,
}}, nil).Once()
testcase := requestBodyTestCase{
path: versionalV2(EntityCategory, SearchAction),
requestBody: []byte(`{"collectionName": "book", "data": ["query data"], "limit": 4, "outputFields": ["word_count"]}`),
}
sendReqAndVerify(t, testEngine, testcase.path, http.MethodPost, testcase)
}
func TestCreateIndex(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Twice()
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(IndexCategory, CreateAction)
// the previous format
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2", "params": {"index_type": "L2", "nlist": 10}}]}`),
})
// the current format
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2", "indexType": "L2", "params":{"nlist": 10}}]}`),
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
fmt.Println(w.Body.String())
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
}
func TestCompact(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().ManualCompaction(mock.Anything, mock.Anything).Return(&milvuspb.ManualCompactionResponse{CompactionID: 1}, nil).Once()
mp.EXPECT().ManualCompaction(mock.Anything, mock.Anything).Return(
&milvuspb.ManualCompactionResponse{
Status: &commonpb.Status{
Code: 1100,
Reason: "mock",
},
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(CollectionCategory, CompactAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"invalid_name"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().GetCompactionState(mock.Anything, mock.Anything).Return(&milvuspb.GetCompactionStateResponse{}, nil).Once()
mp.EXPECT().GetCompactionState(mock.Anything, mock.Anything).Return(
&milvuspb.GetCompactionStateResponse{
Status: &commonpb.Status{
Code: 1100,
Reason: "mock",
},
}, nil).Once()
path = versionalV2(CollectionCategory, CompactionStateAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"invalid_name"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
fmt.Println(w.Body.String())
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
}
func TestFlush(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().Flush(mock.Anything, mock.Anything).Return(&milvuspb.FlushResponse{
Status: merr.Success(),
CollSegIDs: map[string]*schemapb.LongArray{
"test": {Data: []int64{1}},
},
CollFlushTs: map[string]uint64{
"test": 100,
},
}, nil).Once()
mp.EXPECT().GetFlushState(mock.Anything, mock.MatchedBy(func(req *milvuspb.GetFlushStateRequest) bool {
return req.GetCollectionName() == "test" && req.GetFlushTs() == 100 && assert.ObjectsAreEqual([]int64{1}, req.GetSegmentIDs())
})).Return(&milvuspb.GetFlushStateResponse{
Status: merr.Success(),
Flushed: true,
}, nil).Once()
mp.EXPECT().Flush(mock.Anything, mock.Anything).Return(
&milvuspb.FlushResponse{
Status: &commonpb.Status{
Code: 1100,
Reason: "mock",
},
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(CollectionCategory, FlushAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"invalid_name"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
fmt.Println(w.Body.String())
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
}
func TestWaitForFlush(t *testing.T) {
t.Run("missing flush timestamp", func(t *testing.T) {
h := &HandlersV2{proxy: mocks.NewMockProxy(t)}
err := h.waitForFlush(context.Background(), "", "test", &milvuspb.FlushResponse{
CollSegIDs: map[string]*schemapb.LongArray{
"test": {Data: []int64{1}},
},
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "failed to get flush timestamp")
})
t.Run("get flush state rpc error", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetFlushState(mock.Anything, mock.Anything).Return(nil, errors.New("mock rpc")).Once()
h := &HandlersV2{proxy: mp}
err := h.waitForFlush(context.Background(), "", "test", &milvuspb.FlushResponse{
CollSegIDs: map[string]*schemapb.LongArray{
"test": {Data: []int64{1}},
},
CollFlushTs: map[string]uint64{
"test": 100,
},
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "mock rpc")
})
t.Run("context canceled before flush completes", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetFlushState(mock.Anything, mock.Anything).Return(&milvuspb.GetFlushStateResponse{
Status: merr.Success(),
}, nil).Once()
ctx, cancel := context.WithCancel(context.Background())
cancel()
h := &HandlersV2{proxy: mp}
err := h.waitForFlush(ctx, "", "test", &milvuspb.FlushResponse{
CollSegIDs: map[string]*schemapb.LongArray{
"test": {Data: []int64{1}},
},
CollFlushTs: map[string]uint64{
"test": 100,
},
})
assert.ErrorIs(t, err, context.Canceled)
})
}
func TestDatabase(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().CreateDatabase(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().CreateDatabase(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(DataBaseCategory, CreateAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"invalid_name"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().DropDatabase(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().DropDatabase(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
path = versionalV2(DataBaseCategory, DropAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"mock"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(&milvuspb.ListDatabasesResponse{DbNames: []string{"a", "b", "c"}, DbIds: []int64{100, 101, 102}}, nil).Once()
mp.EXPECT().ListDatabases(mock.Anything, mock.Anything).Return(&milvuspb.ListDatabasesResponse{
Status: &commonpb.Status{
Code: 1100,
Reason: "mock",
},
}, nil).Once()
path = versionalV2(DataBaseCategory, ListAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"mock"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().DescribeDatabase(mock.Anything, mock.Anything).Return(&milvuspb.DescribeDatabaseResponse{DbName: "test", DbID: 100}, nil).Once()
mp.EXPECT().DescribeDatabase(mock.Anything, mock.Anything).Return(&milvuspb.DescribeDatabaseResponse{
Status: &commonpb.Status{
Code: 1100,
Reason: "mock",
},
}, nil).Once()
path = versionalV2(DataBaseCategory, DescribeAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"mock"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
path = versionalV2(DataBaseCategory, AlterAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"mock"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
path = versionalV2(DataBaseCategory, DropPropertiesAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"mock"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
path = versionalV2(DataBaseCategory, AlterPropertiesAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"test"}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"dbName":"mock"}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
fmt.Println(w.Body.String())
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
}
func TestColletcionProperties(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().AlterCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterCollection(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(CollectionCategory, AlterPropertiesAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"test", "properties":{"mmap": true}}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"mock", "properties":{"mmap": true}}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().AlterCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterCollection(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
path = versionalV2(CollectionCategory, DropPropertiesAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"test", "propertyKeys":["mmap"]}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"mock", "propertyKeys":["mmap"]}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
fmt.Println(w.Body.String())
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
}
func TestIndexProperties(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().AlterIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterIndex(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(IndexCategory, AlterPropertiesAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"test", "indexName":"test", "properties":{"mmap": true}}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"mock", "indexName":"test", "properties":{"mmap": true}}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
mp.EXPECT().AlterIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterIndex(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
path = versionalV2(IndexCategory, DropPropertiesAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"test","indexName":"test", "propertyKeys":["test"]}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"mock","indexName":"test", "propertyKeys":["test"]}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
fmt.Println(w.Body.String())
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
}
func TestDescribeIndexWithWarmup(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&milvuspb.DescribeIndexResponse{
Status: &StatusSuccess,
IndexDescriptions: []*milvuspb.IndexDescription{
{
IndexName: DefaultIndexName,
FieldName: FieldBookIntro,
Params: []*commonpb.KeyValuePair{
{Key: common.MetricTypeKey, Value: DefaultMetricType},
{Key: common.IndexTypeKey, Value: "IVF_FLAT"},
{Key: common.WarmupKey, Value: "sync"},
},
State: commonpb.IndexState_Finished,
},
},
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(IndexCategory, DescribeAction),
bytes.NewReader([]byte(`{"collectionName": "`+DefaultCollectionName+`", "indexName": "`+DefaultIndexName+`"}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp map[string]any
err := json.Unmarshal(w.Body.Bytes(), &resp)
assert.Nil(t, err)
assert.Equal(t, float64(0), resp["code"])
data := resp["data"].([]any)
assert.Equal(t, 1, len(data))
indexInfo := data[0].(map[string]any)
assert.Equal(t, "sync", indexInfo["warmup"])
}
func TestCollectionFieldProperties(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().AlterCollectionField(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterCollectionField(mock.Anything, mock.Anything).Return(
&commonpb.Status{
Code: 1100,
Reason: "mock",
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(CollectionFieldCategory, AlterPropertiesAction)
// success
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"test", "fieldName":"test", "fieldParams":{"max_length": 100}}`),
})
// mock fail
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName":"mock", "fieldName":"test", "fieldParams":{"max_length": 100}}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
fmt.Println(w.Body.String())
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
}
func TestCreateCollection(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
postTestCases := []requestBodyTestCase{}
mp := mocks.NewMockProxy(t)
mp.EXPECT().CreateCollection(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) {
hasNegativeTTL := false
for _, prop := range req.GetProperties() {
if prop.GetKey() == common.CollectionTTLConfigKey && prop.GetValue() == "-100" {
hasNegativeTTL = true
break
}
}
if hasNegativeTTL {
for _, prop := range req.GetProperties() {
if prop.GetKey() == "ttl.seconds" && prop.GetValue() == "-100" {
return merr.Status(merr.WrapErrParameterInvalidMsg("unexpected ttl.seconds property forwarded")), nil
}
}
return merr.Status(merr.WrapErrParameterInvalidMsg("collection TTL is out of range, expect [-1, 3155760000], got -100")), nil
}
return commonSuccessStatus, nil
}).Times(17)
mp.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Times(8)
mp.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Times(7)
mp.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonErrorStatus, nil).Twice()
mp.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(commonErrorStatus, nil).Twice()
testEngine := initHTTPServerV2(mp, false)
path := versionalV2(CollectionCategory, CreateAction)
// quickly create collection
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `"}`),
errMsg: "dimension is required for quickly create collection(default metric type: COSINE): invalid parameter[expected=collectionName & dimension][actual=collectionName]",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": "256", "enableDynamicField": "false", "shardsNum": "2", "consistencyLevel": "Strong", "ttlSeconds": "3600"}}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": "256", "enableDynamicField": false, "shardsNum": "2", "consistencyLevel": "Strong", "ttlSeconds": "3600"}}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": 256, "enableDynamicField": false, "shardsNum": 2, "consistencyLevel": "Strong", "ttlSeconds": 3600}}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": 256, "enableDynamicField": false, "shardsNum": 2, "consistencyLevel": "unknown", "ttlSeconds": 3600}}`),
errMsg: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded: invalid parameter[expected=Strong, Session, Bounded, Eventually, Customized][actual=unknown]",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": 256}, "properties": {"collection.ttl.seconds": -100}}`),
errMsg: "collection TTL is out of range, expect [-1, 3155760000], got -100: invalid parameter",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": 256}, "properties": {"ttl.seconds": -100}}`),
errMsg: "collection TTL is out of range, expect [-1, 3155760000], got -100: invalid parameter",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": 256, "ttlSeconds": 3600}, "properties": {"collection.ttl.seconds": 3600}}`),
errMsg: "collection TTL is specified multiple times: invalid parameter",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": 256}, "properties": {"collection.ttl.seconds": 3600, "ttl.seconds": 3600}}`),
errMsg: "collection TTL is specified multiple times: invalid parameter",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "unknown"}`),
errMsg: "idType can only be [Int64, VarChar], default: Int64: invalid parameter[expected=Int64, Varchar][actual=unknown]",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "metricType": "L2"}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "isPartitionKey": false, "elementTypeParams": {}},
{"fieldName": "partition_field", "dataType": "VarChar", "isPartitionKey": true, "elementTypeParams": {"max_length": 256}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}, "params": {"partitionsNum": "32"}}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}}`),
errMsg: "invalid parameter, data type int64 is invalid(case sensitive).",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "null_fid", "dataType": "Int64", "nullable": true},
{"fieldName": "default_fid_bool", "dataType": "Bool", "defaultValue": true},
{"fieldName": "default_fid_int8", "dataType": "Int8", "defaultValue": 10},
{"fieldName": "default_fid_int16", "dataType": "Int16", "defaultValue": 10},
{"fieldName": "default_fid_int32", "dataType": "Int32", "defaultValue": 10},
{"fieldName": "default_fid_int64", "dataType": "Int64", "defaultValue": 10},
{"fieldName": "default_fid_float32", "dataType": "Float", "defaultValue": 10},
{"fieldName": "default_fid_double", "dataType": "Double", "defaultValue": 10},
{"fieldName": "default_fid_varchar", "dataType": "VarChar", "defaultValue": "a"},
{"fieldName": "word_count", "dataType": "Array", "elementDataType": "Int64", "elementTypeParams": {"max_capacity": 2}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}}`),
})
// dim should not be specified for SparseFloatVector field
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "isPartitionKey": false, "elementTypeParams": {}},
{"fieldName": "partition_field", "dataType": "VarChar", "isPartitionKey": true, "elementTypeParams": {"max_length": 256}},
{"fieldName": "book_intro", "dataType": "SparseFloatVector", "elementTypeParams": {}}
]
}, "params": {"partitionsNum": "32"}}`),
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Array", "elementDataType": "int64", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}}`),
errMsg: "invalid parameter, element data type int64 is invalid(case sensitive).",
errCode: 1100, // ErrParameterInvalid
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64","isClusteringKey":true, "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_xxx", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "missing required parameters, error: `book_xxx` hasn't defined in schema",
errCode: 1802,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "isPartitionKey": true, "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}, "params": {"partitionKeyIsolation": "true"}}`),
})
// Struct sub-vector index via collection_create indexParams: the
// structName[subName] form is registered in fieldNames so users can
// create the sub-vector index alongside top-level indexes in one call.
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
],
"structFields": [{
"fieldName": "my_struct",
"fields": [
{"fieldName": "sub_vec", "dataType": "ArrayOfVector", "elementDataType": "FloatVector", "elementTypeParams": {"dim": 2, "max_capacity": 4}}
]
}]
}, "indexParams": [
{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"},
{"fieldName": "my_struct[sub_vec]", "indexName": "sub_vec_idx", "metricType": "L2"}
]}`),
})
// Struct sub-field qualified name that does not exist in schema is rejected.
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
],
"structFields": [{
"fieldName": "my_struct",
"fields": [
{"fieldName": "sub_vec", "dataType": "ArrayOfVector", "elementDataType": "FloatVector", "elementTypeParams": {"dim": 2, "max_capacity": 4}}
]
}]
}, "indexParams": [
{"fieldName": "my_struct[nonexistent]", "indexName": "bad", "metricType": "L2"}
]}`),
errMsg: "missing required parameters, error: `my_struct[nonexistent]` hasn't defined in schema",
errCode: 1802,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "metricType": "L2"}`),
errMsg: "",
errCode: 65535,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "",
errCode: 65535,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "metricType": "L2"}`),
errMsg: "",
errCode: 65535,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "SparseFloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "",
errCode: 65535,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "default_fid", "dataType": "Bool", "defaultValue":10, "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "SparseFloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "convert defaultValue fail, err:cannot use \"10\"(type: float64) as bool default value: invalid parameter",
errCode: 1100,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "default_fid", "dataType": "VarChar", "defaultValue":true, "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "SparseFloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "convert defaultValue fail, err:cannot use \"true\"(type: bool) as string default value: invalid parameter",
errCode: 1100,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "default_fid", "dataType": "Int8", "defaultValue":"10", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "SparseFloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "convert defaultValue fail, err:cannot use \"\"10\"(type: string) as int default value: invalid parameter",
errCode: 1100,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "default_fid", "dataType": "Int64", "defaultValue":"10", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "SparseFloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "convert defaultValue fail, err:cannot use \"10\"(type: string) as long default value: invalid parameter",
errCode: 1100,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "default_fid", "dataType": "Float", "defaultValue":"10", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "SparseFloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "convert defaultValue fail, err:cannot use \"10\"(type: string) as float default value: invalid parameter",
errCode: 1100,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "word_count", "dataType": "Int64", "elementTypeParams": {}},
{"fieldName": "default_fid", "dataType": "Double", "defaultValue":"10", "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "SparseFloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}]}`),
errMsg: "convert defaultValue fail, err:cannot use \"10\"(type: string) as float default value: invalid parameter",
errCode: 1100,
})
postTestCases = append(postTestCases, requestBodyTestCase{
path: path,
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `", "dimension": 2, "idType": "Varchar",` +
`"params": {"max_length": 256, "enableDynamicField": 100, "shardsNum": 2, "consistencyLevel": "unknown", "ttlSeconds": 3600}}`),
errMsg: "parse enableDynamicField fail, err:strconv.ParseBool: parsing \"100\": invalid syntax",
errCode: 65535,
})
for _, testcase := range postTestCases {
t.Run("post"+testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
fmt.Println(w.Body.String())
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
// collection-level warmup params: assert Properties contains the 4 warmup key-value pairs
t.Run("warmup properties in CreateCollectionRequest", func(t *testing.T) {
mp2 := mocks.NewMockProxy(t)
mp2.EXPECT().CreateCollection(mock.Anything, mock.MatchedBy(func(req *milvuspb.CreateCollectionRequest) bool {
warmupProps := make(map[string]string)
for _, kv := range req.Properties {
warmupProps[kv.Key] = kv.Value
}
return warmupProps[common.WarmupScalarFieldKey] == "sync" &&
warmupProps[common.WarmupScalarIndexKey] == "sync" &&
warmupProps[common.WarmupVectorFieldKey] == "sync" &&
warmupProps[common.WarmupVectorIndexKey] == "disable"
})).Return(commonSuccessStatus, nil).Once()
mp2.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp2.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
testEngine2 := initHTTPServerV2(mp2, false)
reqBody := []byte(`{"collectionName": "` + DefaultCollectionName + `", "schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}},
{"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}
]
}, "indexParams": [{"fieldName": "book_intro", "indexName": "book_intro_vector", "metricType": "L2"}],
"params": {"warmup.scalarField": "sync", "warmup.scalarIndex": "sync", "warmup.vectorField": "sync", "warmup.vectorIndex": "disable"}}`)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(reqBody))
w := httptest.NewRecorder()
testEngine2.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, int32(0), returnBody.Code)
})
}
func TestBinaryMetricTypesForQuickCreate(t *testing.T) {
params := paramtable.Get()
key := params.AutoIndexConfig.BinaryIndexParams.Key
testcases := []struct {
name string
config string
expected []string
}{
{
name: "bin ivf flat uses ivf binary metrics",
config: `{"nlist": 1024, "index_type": "BIN_IVF_FLAT", "metric_type": "HAMMING"}`,
expected: indexparamcheck.BinIvfMetrics,
},
{
name: "missing index type falls back to ivf binary metrics",
config: `{"metric_type": "HAMMING"}`,
expected: indexparamcheck.BinIvfMetrics,
},
{
name: "unknown index type falls back to ivf binary metrics",
config: `{"index_type": "UNKNOWN", "metric_type": "HAMMING"}`,
expected: indexparamcheck.BinIvfMetrics,
},
{
name: "bin flat uses id map binary metrics",
config: `{"index_type": "BIN_FLAT", "metric_type": "HAMMING"}`,
expected: indexparamcheck.BinIDMapMetrics,
},
}
for _, testcase := range testcases {
t.Run(testcase.name, func(t *testing.T) {
params.Save(key, testcase.config)
t.Cleanup(func() {
params.Reset(key)
})
assert.ElementsMatch(t, testcase.expected, binaryMetricTypesForQuickCreate())
})
}
}
func TestCreateCollectionQuickVectorFieldType(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
path := versionalV2(CollectionCategory, CreateAction)
t.Run("reject invalid vector field type", func(t *testing.T) {
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(`{
"collectionName": "book",
"dimension": 4,
"vectorFieldType": "InvalidVectorType"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(1100), returnBody.Code)
assert.Equal(t, "vectorFieldType can only be [FloatVector, BinaryVector, Float16Vector, BFloat16Vector, SparseFloatVector], default: FloatVector: invalid parameter[expected=FloatVector, BinaryVector, Float16Vector, BFloat16Vector, SparseFloatVector][actual=InvalidVectorType]", returnBody.Message)
assert.Nil(t, proxy.createReq)
assert.Empty(t, proxy.createIndexReqs)
})
t.Run("reject sparse vector dimension", func(t *testing.T) {
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(`{
"collectionName": "book",
"dimension": 4,
"vectorFieldType": "SparseFloatVector"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(1100), returnBody.Code)
assert.Contains(t, returnBody.Message, "dimension should not be specified for SparseFloatVector quick create")
assert.Nil(t, proxy.createReq)
assert.Empty(t, proxy.createIndexReqs)
})
invalidMetricCases := []struct {
name string
body string
expectedMsg string
}{
{
name: "reject binary vector invalid metric",
body: `{
"collectionName": "book",
"dimension": 8,
"metricType": "COSINE",
"vectorFieldType": "BinaryVector"
}`,
expectedMsg: "binary vector index does not support metric type: COSINE",
},
{
name: "reject binary vector substructure metric",
body: `{
"collectionName": "book",
"dimension": 8,
"metricType": "SUBSTRUCTURE",
"vectorFieldType": "BinaryVector"
}`,
expectedMsg: "binary vector index does not support metric type: SUBSTRUCTURE",
},
{
name: "reject binary vector superstructure metric",
body: `{
"collectionName": "book",
"dimension": 8,
"metricType": "SUPERSTRUCTURE",
"vectorFieldType": "BinaryVector"
}`,
expectedMsg: "binary vector index does not support metric type: SUPERSTRUCTURE",
},
{
name: "reject binary vector mhjaccard metric",
body: `{
"collectionName": "book",
"dimension": 8,
"metricType": "MHJACCARD",
"vectorFieldType": "BinaryVector"
}`,
expectedMsg: "binary vector index does not support metric type: MHJACCARD",
},
{
name: "reject sparse float vector invalid metric",
body: `{
"collectionName": "book",
"metricType": "COSINE",
"vectorFieldType": "SparseFloatVector"
}`,
expectedMsg: "only IP&BM25 is the supported metric type for sparse index",
},
{
name: "reject sparse float vector bm25 metric",
body: `{
"collectionName": "book",
"metricType": "BM25",
"vectorFieldType": "SparseFloatVector"
}`,
expectedMsg: "only BM25 Function output field support BM25 metric type",
},
}
for _, testcase := range invalidMetricCases {
t.Run(testcase.name, func(t *testing.T) {
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(testcase.body)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(1100), returnBody.Code)
assert.Contains(t, returnBody.Message, testcase.expectedMsg)
assert.Nil(t, proxy.createReq)
assert.Empty(t, proxy.createIndexReqs)
})
}
t.Run("binary vector metric follows configured binary auto index", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().AutoIndexConfig.BinaryIndexParams.Key, `{"index_type": "BIN_FLAT", "metric_type": "HAMMING"}`)
defer paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.BinaryIndexParams.Key)
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(`{
"collectionName": "book",
"dimension": 8,
"metricType": "SUBSTRUCTURE",
"vectorFieldType": "BinaryVector"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(0), returnBody.Code)
require.NotNil(t, proxy.createReq)
require.Len(t, proxy.createIndexReqs, 1)
require.Len(t, proxy.createIndexReqs[0].GetExtraParams(), 1)
assert.Equal(t, common.MetricTypeKey, proxy.createIndexReqs[0].GetExtraParams()[0].GetKey())
assert.Equal(t, "SUBSTRUCTURE", proxy.createIndexReqs[0].GetExtraParams()[0].GetValue())
})
testcases := []struct {
name string
body string
expectedType schemapb.DataType
expectedDimParam bool
expectedMetric string
}{
{
name: "binary vector default metric",
body: `{
"collectionName": "book",
"dimension": 8,
"vectorFieldType": "BinaryVector"
}`,
expectedType: schemapb.DataType_BinaryVector,
expectedDimParam: true,
expectedMetric: paramtable.BinaryVectorDefaultMetricType,
},
{
name: "float16 vector",
body: `{
"collectionName": "book",
"dimension": 4,
"vectorFieldType": "Float16Vector"
}`,
expectedType: schemapb.DataType_Float16Vector,
expectedDimParam: true,
expectedMetric: DefaultMetricType,
},
{
name: "bfloat16 vector",
body: `{
"collectionName": "book",
"dimension": 4,
"vectorFieldType": "BFloat16Vector"
}`,
expectedType: schemapb.DataType_BFloat16Vector,
expectedDimParam: true,
expectedMetric: DefaultMetricType,
},
{
name: "sparse float vector default metric",
body: `{
"collectionName": "book",
"vectorFieldType": "SparseFloatVector"
}`,
expectedType: schemapb.DataType_SparseFloatVector,
expectedDimParam: false,
expectedMetric: paramtable.SparseFloatVectorDefaultMetricType,
},
}
for _, testcase := range testcases {
t.Run(testcase.name, func(t *testing.T) {
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(testcase.body)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(0), returnBody.Code)
require.NotNil(t, proxy.createReq)
collSchema := &schemapb.CollectionSchema{}
require.NoError(t, proto.Unmarshal(proxy.createReq.GetSchema(), collSchema))
require.Len(t, collSchema.GetFields(), 2)
vectorField := collSchema.GetFields()[1]
assert.Equal(t, testcase.expectedType, vectorField.GetDataType())
dimFound := false
for _, typeParam := range vectorField.GetTypeParams() {
if typeParam.GetKey() == Dim {
dimFound = true
}
}
assert.Equal(t, testcase.expectedDimParam, dimFound)
require.Len(t, proxy.createIndexReqs, 1)
indexReq := proxy.createIndexReqs[0]
assert.Equal(t, DefaultVectorFieldName, indexReq.GetFieldName())
assert.Equal(t, DefaultVectorFieldName, indexReq.GetIndexName())
require.Len(t, indexReq.GetExtraParams(), 1)
assert.Equal(t, common.MetricTypeKey, indexReq.GetExtraParams()[0].GetKey())
assert.Equal(t, testcase.expectedMetric, indexReq.GetExtraParams()[0].GetValue())
})
}
}
func TestCreateCollectionTopLevelConsistencyLevel(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
path := versionalV2(CollectionCategory, CreateAction)
t.Run("reject invalid top-level consistency level", func(t *testing.T) {
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(`{
"collectionName": "book",
"dimension": 4,
"consistencyLevel": "Invalid"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(1100), returnBody.Code)
assert.Contains(t, returnBody.Message, "consistencyLevel can only be")
assert.Nil(t, proxy.createReq)
assert.Empty(t, proxy.createIndexReqs)
})
t.Run("accept valid top-level consistency level", func(t *testing.T) {
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(`{
"collectionName": "book",
"dimension": 4,
"consistencyLevel": "Strong"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(0), returnBody.Code)
require.NotNil(t, proxy.createReq)
assert.Equal(t, commonpb.ConsistencyLevel_Strong, proxy.createReq.GetConsistencyLevel())
require.Len(t, proxy.createIndexReqs, 1)
})
t.Run("reject conflicting consistency levels", func(t *testing.T) {
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(`{
"collectionName": "book",
"dimension": 4,
"consistencyLevel": "Strong",
"params": {"consistencyLevel": "Bounded"}
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, int32(1100), returnBody.Code)
assert.Contains(t, returnBody.Message, "top-level consistencyLevel conflicts with params.consistencyLevel")
assert.Nil(t, proxy.createReq)
assert.Empty(t, proxy.createIndexReqs)
})
}
func TestCreateCollectionStructArrayDuplicateName(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
testEngine := initHTTPServerV2(mocks.NewMockProxy(t), false)
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionCategory, CreateAction), bytes.NewReader([]byte(`{
"collectionName": "test",
"schema": {
"fields": [
{"fieldName": "book_id", "dataType": "Int64", "isPrimary": true, "elementTypeParams": {}}
],
"structFields": [{
"fieldName": "book_id",
"fields": [
{"fieldName": "sub_int", "dataType": "Array", "elementDataType": "Int32"}
]
}]
}
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, merr.Code(merr.ErrParameterInvalid), returnBody.Code)
assert.Contains(t, returnBody.Message, "duplicated field name: book_id")
}
func versionalV2(category string, action string) string {
return "/v2/vectordb" + category + action
}
func initHTTPServerV2(proxy types.ProxyComponent, needAuth bool) *gin.Engine {
h := NewHandlersV2(proxy)
ginHandler := gin.Default()
appV2 := ginHandler.Group("/v2/vectordb", genAuthMiddleWare(needAuth))
h.RegisterRoutesToV2(appV2)
return ginHandler
}
type externalCollectionRESTProxy struct {
mockProxyComponent
createReq *milvuspb.CreateCollectionRequest
createIndexReqs []*milvuspb.CreateIndexRequest
describeResp *milvuspb.DescribeCollectionResponse
refreshReq *milvuspb.RefreshExternalCollectionRequest
listReq *milvuspb.ListRefreshExternalCollectionJobsRequest
progressReq *milvuspb.GetRefreshExternalCollectionProgressRequest
restoreReq *milvuspb.RestoreExternalSnapshotRequest
exportReq *milvuspb.ExportSnapshotRequest
getRestoreSnapshotStateReq *milvuspb.GetRestoreSnapshotStateRequest
listRestoreSnapshotJobsReq *milvuspb.ListRestoreSnapshotJobsRequest
}
func (m *externalCollectionRESTProxy) CreateCollection(ctx context.Context, request *milvuspb.CreateCollectionRequest) (*commonpb.Status, error) {
m.createReq = request
return commonSuccessStatus, nil
}
func (m *externalCollectionRESTProxy) CreateIndex(ctx context.Context, request *milvuspb.CreateIndexRequest) (*commonpb.Status, error) {
m.createIndexReqs = append(m.createIndexReqs, request)
return commonSuccessStatus, nil
}
func (m *externalCollectionRESTProxy) DescribeCollection(ctx context.Context, request *milvuspb.DescribeCollectionRequest) (*milvuspb.DescribeCollectionResponse, error) {
if m.describeResp != nil {
return m.describeResp, nil
}
return m.mockProxyComponent.DescribeCollection(ctx, request)
}
func (m *externalCollectionRESTProxy) GetLoadState(ctx context.Context, request *milvuspb.GetLoadStateRequest) (*milvuspb.GetLoadStateResponse, error) {
return &DefaultLoadStateResp, nil
}
func (m *externalCollectionRESTProxy) DescribeIndex(ctx context.Context, request *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error) {
return &DefaultDescIndexesReqp, nil
}
func (m *externalCollectionRESTProxy) ListAliases(ctx context.Context, request *milvuspb.ListAliasesRequest) (*milvuspb.ListAliasesResponse, error) {
return &milvuspb.ListAliasesResponse{
Status: &StatusSuccess,
Aliases: []string{DefaultAliasName},
}, nil
}
func (m *externalCollectionRESTProxy) RefreshExternalCollection(ctx context.Context, request *milvuspb.RefreshExternalCollectionRequest) (*milvuspb.RefreshExternalCollectionResponse, error) {
m.refreshReq = request
return &milvuspb.RefreshExternalCollectionResponse{
Status: commonSuccessStatus,
JobId: 1001,
}, nil
}
func (m *externalCollectionRESTProxy) GetRefreshExternalCollectionProgress(ctx context.Context, request *milvuspb.GetRefreshExternalCollectionProgressRequest) (*milvuspb.GetRefreshExternalCollectionProgressResponse, error) {
m.progressReq = request
return &milvuspb.GetRefreshExternalCollectionProgressResponse{
Status: commonSuccessStatus,
JobInfo: &milvuspb.RefreshExternalCollectionJobInfo{
JobId: request.GetJobId(),
CollectionName: "external_books",
State: milvuspb.RefreshExternalCollectionState_RefreshInProgress,
Progress: 42,
ExternalSource: "s3://bucket/books",
ExternalSpec: `{"format":"parquet"}`,
StartTime: 10,
},
}, nil
}
func (m *externalCollectionRESTProxy) ListRefreshExternalCollectionJobs(ctx context.Context, request *milvuspb.ListRefreshExternalCollectionJobsRequest) (*milvuspb.ListRefreshExternalCollectionJobsResponse, error) {
m.listReq = request
return &milvuspb.ListRefreshExternalCollectionJobsResponse{
Status: commonSuccessStatus,
Jobs: []*milvuspb.RefreshExternalCollectionJobInfo{
{
JobId: 1001,
CollectionName: request.GetCollectionName(),
State: milvuspb.RefreshExternalCollectionState_RefreshCompleted,
Progress: 100,
ExternalSource: "s3://bucket/books",
ExternalSpec: `{"format":"parquet"}`,
StartTime: 10,
EndTime: 20,
},
},
}, nil
}
func (m *externalCollectionRESTProxy) RestoreExternalSnapshot(ctx context.Context, request *milvuspb.RestoreExternalSnapshotRequest) (*milvuspb.RestoreExternalSnapshotResponse, error) {
m.restoreReq = request
return &milvuspb.RestoreExternalSnapshotResponse{
Status: merr.Success(),
JobId: 2001,
}, nil
}
func (m *externalCollectionRESTProxy) ExportSnapshot(ctx context.Context, request *milvuspb.ExportSnapshotRequest) (*milvuspb.ExportSnapshotResponse, error) {
m.exportReq = request
return &milvuspb.ExportSnapshotResponse{
Status: merr.Success(),
SnapshotMetadataUri: request.GetTargetS3Path() + "/snapshots/100/metadata/1.json",
}, nil
}
func (m *externalCollectionRESTProxy) GetRestoreSnapshotState(ctx context.Context, request *milvuspb.GetRestoreSnapshotStateRequest) (*milvuspb.GetRestoreSnapshotStateResponse, error) {
m.getRestoreSnapshotStateReq = request
return &milvuspb.GetRestoreSnapshotStateResponse{
Status: merr.Success(),
Info: &milvuspb.RestoreSnapshotInfo{
JobId: request.GetJobId(),
SnapshotName: "snapshot_1",
DbName: "default",
CollectionName: "restored_collection",
State: milvuspb.RestoreSnapshotState_RestoreSnapshotCompleted,
Progress: 100,
},
}, nil
}
func (m *externalCollectionRESTProxy) ListRestoreSnapshotJobs(ctx context.Context, request *milvuspb.ListRestoreSnapshotJobsRequest) (*milvuspb.ListRestoreSnapshotJobsResponse, error) {
m.listRestoreSnapshotJobsReq = request
return &milvuspb.ListRestoreSnapshotJobsResponse{
Status: merr.Success(),
Jobs: []*milvuspb.RestoreSnapshotInfo{{
JobId: 2001,
SnapshotName: "snapshot_1",
DbName: request.GetDbName(),
CollectionName: request.GetCollectionName(),
State: milvuspb.RestoreSnapshotState_RestoreSnapshotPending,
Progress: 10,
}},
}, nil
}
func TestFieldSchemaGetProtoWithExternalField(t *testing.T) {
field := &FieldSchema{
FieldName: "title",
DataType: "VarChar",
ExternalField: "book_title",
ElementTypeParams: map[string]interface{}{
common.MaxLengthKey: 256,
},
}
got, err := field.GetProto(context.Background())
assert.NoError(t, err)
assert.Equal(t, "book_title", got.GetExternalField())
assert.Equal(t, schemapb.DataType_VarChar, got.GetDataType())
}
func TestCreateExternalCollectionRESTV2(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
testcases := []struct {
name string
body string
source string
spec string
}{
{
name: "schema level external config",
body: `{
"dbName": "default",
"collectionName": "external_books",
"schema": {
"externalSource": "s3://bucket/books",
"externalSpec": "{\"format\":\"parquet\"}",
"fields": [
{"fieldName": "book_intro", "dataType": "FloatVector", "externalField": "embedding", "elementTypeParams": {"dim": 2}},
{"fieldName": "title", "dataType": "VarChar", "externalField": "book_title", "elementTypeParams": {"max_length": 256}}
]
}
}`,
source: "s3://bucket/books",
spec: `{"format":"parquet"}`,
},
}
for _, testcase := range testcases {
t.Run(testcase.name, func(t *testing.T) {
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionCategory, CreateAction), bytes.NewReader([]byte(testcase.body)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.NotNil(t, proxy.createReq)
collSchema := &schemapb.CollectionSchema{}
assert.NoError(t, proto.Unmarshal(proxy.createReq.GetSchema(), collSchema))
assert.Equal(t, testcase.source, collSchema.GetExternalSource())
assert.Equal(t, testcase.spec, collSchema.GetExternalSpec())
assert.Equal(t, "embedding", collSchema.GetFields()[0].GetExternalField())
assert.Equal(t, "book_title", collSchema.GetFields()[1].GetExternalField())
})
}
}
func TestCreateExternalCollectionTopLevelConfigRejectedRESTV2(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionCategory, CreateAction), bytes.NewReader([]byte(`{
"collectionName": "external_books",
"externalSource": "s3://bucket/books",
"externalSpec": "{\"format\":\"parquet\"}",
"schema": {
"fields": [
{"fieldName": "book_intro", "dataType": "FloatVector", "externalField": "embedding", "elementTypeParams": {"dim": 2}},
{"fieldName": "title", "dataType": "VarChar", "externalField": "book_title", "elementTypeParams": {"max_length": 256}}
]
}
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Nil(t, proxy.createReq)
returnBody := &ReturnErrMsg{}
assert.NoError(t, json.Unmarshal(w.Body.Bytes(), returnBody))
assert.NotZero(t, returnBody.Code)
assert.Contains(t, returnBody.Message, "schema.externalSource/schema.externalSpec")
}
func TestCreateExternalCollectionQuickCreateRejectedRESTV2(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionCategory, CreateAction), bytes.NewReader([]byte(`{
"collectionName": "external_books",
"dimension": 2,
"schema": {
"externalSource": "s3://bucket/books",
"externalSpec": "{\"format\":\"parquet\"}"
}
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Nil(t, proxy.createReq)
returnBody := &ReturnErrMsg{}
assert.NoError(t, json.Unmarshal(w.Body.Bytes(), returnBody))
assert.NotZero(t, returnBody.Code)
assert.Contains(t, returnBody.Message, "external collection is not supported by quick create")
}
func TestDescribeExternalCollectionRESTV2(t *testing.T) {
paramtable.Init()
schema := generateCollectionSchema(schemapb.DataType_Int64, false, false)
schema.ExternalSource = "s3://bucket/books"
schema.ExternalSpec = `{"format":"parquet"}`
schema.Fields[1].ExternalField = "word_count_col"
schema.Fields[2].ExternalField = "embedding"
proxy := &externalCollectionRESTProxy{
describeResp: &milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: schema,
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
},
}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionCategory, DescribeAction), bytes.NewReader([]byte(`{"collectionName": "`+DefaultCollectionName+`"}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp map[string]interface{}
assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
assert.Equal(t, float64(0), resp[HTTPReturnCode])
data := resp[HTTPReturnData].(map[string]interface{})
assert.Equal(t, "s3://bucket/books", data["externalSource"])
assert.Equal(t, `{"format":"parquet"}`, data["externalSpec"])
fields := data["fields"].([]interface{})
wordCount := fields[1].(map[string]interface{})
vector := fields[2].(map[string]interface{})
assert.Equal(t, "word_count_col", wordCount["externalField"])
assert.Equal(t, "embedding", vector["externalField"])
}
func TestExternalCollectionJobRoutesRESTV2(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(ExternalCollectionJobCategory, RefreshAction), bytes.NewReader([]byte(`{
"dbName": "default",
"collectionName": "external_books",
"externalSource": "s3://bucket/books_v2",
"externalSpec": "{\"format\":\"parquet\"}"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, "external_books", proxy.refreshReq.GetCollectionName())
assert.Equal(t, "s3://bucket/books_v2", proxy.refreshReq.GetExternalSource())
req = httptest.NewRequest(http.MethodPost, versionalV2(ExternalCollectionJobCategory, DescribeAction), bytes.NewReader([]byte(`{"jobId": 1001}`)))
w = httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, int64(1001), proxy.progressReq.GetJobId())
assert.Contains(t, w.Body.String(), `"progress":42`)
assert.Contains(t, w.Body.String(), `"externalSpec":"{\"format\":\"parquet\"}"`)
req = httptest.NewRequest(http.MethodPost, versionalV2(ExternalCollectionJobCategory, ListAction), bytes.NewReader([]byte(`{"collectionName": "external_books"}`)))
w = httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, "external_books", proxy.listReq.GetCollectionName())
assert.Contains(t, w.Body.String(), `"state":"RefreshCompleted"`)
assert.Contains(t, w.Body.String(), `"externalSpec":"{\"format\":\"parquet\"}"`)
}
func TestRestoreExternalSnapshotRESTV2(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(SnapshotJobCategory, RestoreExternalAction), bytes.NewReader([]byte(`{
"dbName": "default",
"targetCollectionName": "restored_books",
"snapshotMetadataURI": "s3://bucket/files/snapshots/meta.json",
"externalSpec": "{\"extfs\":{\"cloud_provider\":\"aws\",\"region\":\"us-west-2\",\"use_iam\":\"true\"}}"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, "default", proxy.restoreReq.GetDbName())
assert.Equal(t, "restored_books", proxy.restoreReq.GetTargetCollectionName())
assert.Equal(t, "s3://bucket/files/snapshots/meta.json", proxy.restoreReq.GetSnapshotMetadataUri())
assert.Equal(t, `{"extfs":{"cloud_provider":"aws","region":"us-west-2","use_iam":"true"}}`, proxy.restoreReq.GetExternalSpec())
assert.Contains(t, w.Body.String(), `"jobId":2001`)
}
func TestExportSnapshotRESTV2(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(SnapshotJobCategory, ExportAction), bytes.NewReader([]byte(`{
"dbName": "default",
"collectionName": "source_books",
"snapshotName": "snapshot_1",
"targetS3Path": "s3://foreign-bucket/export-root",
"externalSpec": "{\"extfs\":{\"cloud_provider\":\"aws\",\"region\":\"us-west-2\",\"use_iam\":\"true\"}}"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, "default", proxy.exportReq.GetDbName())
assert.Equal(t, "source_books", proxy.exportReq.GetCollectionName())
assert.Equal(t, "snapshot_1", proxy.exportReq.GetName())
assert.Equal(t, "s3://foreign-bucket/export-root", proxy.exportReq.GetTargetS3Path())
assert.Equal(t, `{"extfs":{"cloud_provider":"aws","region":"us-west-2","use_iam":"true"}}`, proxy.exportReq.GetExternalSpec())
assert.Contains(t, w.Body.String(), `"snapshotMetadataURI":"s3://foreign-bucket/export-root/snapshots/100/metadata/1.json"`)
}
func TestRestoreSnapshotJobRESTV2(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
proxy := &externalCollectionRESTProxy{}
testEngine := initHTTPServerV2(proxy, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(SnapshotJobCategory, DescribeAction), bytes.NewReader([]byte(`{"jobId":"2001"}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, int64(2001), proxy.getRestoreSnapshotStateReq.GetJobId())
assert.Contains(t, w.Body.String(), `"collectionName":"restored_collection"`)
assert.Contains(t, w.Body.String(), `"state":"RestoreSnapshotCompleted"`)
req = httptest.NewRequest(http.MethodPost, versionalV2(SnapshotJobCategory, ListAction), bytes.NewReader([]byte(`{
"dbName": "default",
"collectionName": "restored_collection"
}`)))
w = httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, "default", proxy.listRestoreSnapshotJobsReq.GetDbName())
assert.Equal(t, "restored_collection", proxy.listRestoreSnapshotJobsReq.GetCollectionName())
assert.Contains(t, w.Body.String(), `"records":[`)
assert.Contains(t, w.Body.String(), `"state":"RestoreSnapshotPending"`)
}
func TestRestoreSnapshotJobRESTV2InvalidJobID(t *testing.T) {
paramtable.Init()
testEngine := initHTTPServerV2(&externalCollectionRESTProxy{}, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(SnapshotJobCategory, DescribeAction), bytes.NewReader([]byte(`{"jobId":"bad"}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
assert.NoError(t, json.Unmarshal(w.Body.Bytes(), returnBody))
assert.NotZero(t, returnBody.Code)
assert.Contains(t, returnBody.Message, "int64 jobId")
}
func TestRestoreExternalSnapshotRESTV2RequiresMetadataURI(t *testing.T) {
paramtable.Init()
testEngine := initHTTPServerV2(&externalCollectionRESTProxy{}, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(SnapshotJobCategory, RestoreExternalAction), bytes.NewReader([]byte(`{
"targetCollectionName": "restored_books"
}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
assert.NoError(t, json.Unmarshal(w.Body.Bytes(), returnBody))
assert.NotZero(t, returnBody.Code)
assert.Contains(t, returnBody.Message, "missing required parameters")
}
func TestExternalCollectionJobDescribeRequiresJobIDRESTV2(t *testing.T) {
paramtable.Init()
testEngine := initHTTPServerV2(&externalCollectionRESTProxy{}, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(ExternalCollectionJobCategory, DescribeAction), bytes.NewReader([]byte(`{}`)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
assert.NoError(t, json.Unmarshal(w.Body.Bytes(), returnBody))
assert.NotZero(t, returnBody.Code)
assert.Contains(t, returnBody.Message, "missing required parameters")
}
/**
| path| ListDatabases | ShowCollections | HasCollection | DescribeCollection | GetLoadState | DescribeIndex | GetCollectionStatistics | GetLoadingProgress |
|collections | | 1 | | | | | | |
|has?coll | | | 1 | | | | | |
|desc?coll | | | | 1 | 1 | 1 | | |
|stats?coll | | | | | | | 1 | |
|loadState?coll| | | | | 1 | | | 1 |
|collections | | 1 | | | | | | |
|has/coll/ | | | 1 | | | | | |
|has/coll/default/| | | 1 | | | | | |
|has/coll/db/ | 1 | | | | | | | |
|desc/coll/ | | | | 1 | 1 | 1 | | |
|stats/coll/ | | | | | | | 1 | |
|loadState/coll| | | | | 1 | | | 1 |
| path| ShowPartitions | HasPartition | GetPartitionStatistics |
|partitions?coll | 1 | | |
|has?coll&part | | 1 | |
|stats?coll&part | | | 1 |
|partitions/coll | 1 | | |
|has/coll/part | | 1 | |
|stats/coll/part | | | 1 |
| path| ListCredUsers | SelectUser |
|users | 1 | |
|desc?user | | 1 |
|users | 1 | |
|desc/user | | 1 |
| path| SelectRole | SelectGrant |
|roles | 1 | |
|desc?role | | 1 |
|roles | 1 | |
|desc/role | | 1 |
| path| DescribeCollection | DescribeIndex |
|indexes | 0 | 1 |
|desc?index | | 1 |
|indexes | 0 | 1 |
|desc/index | | 1 |
| path| ListAliases | DescribeAlias |
|aliases | 1 | |
|desc?alias | | 1 |
|aliases | 1 | |
|desc/alias | | 1 |
*/
func TestMethodGet(t *testing.T) {
paramtable.Init()
mp := mocks.NewMockProxy(t)
mp.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&milvuspb.ShowCollectionsResponse{
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().ShowCollections(mock.Anything, mock.Anything).Return(&milvuspb.ShowCollectionsResponse{
Status: &StatusSuccess,
CollectionNames: []string{DefaultCollectionName},
}, nil).Once()
mp.EXPECT().HasCollection(mock.Anything, mock.Anything).Return(&milvuspb.BoolResponse{
Status: &StatusSuccess,
Value: true,
}, nil).Once()
mp.EXPECT().HasCollection(mock.Anything, mock.Anything).Return(&milvuspb.BoolResponse{Status: commonErrorStatus}, nil).Once()
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Twice()
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{Status: commonErrorStatus}, nil).Once()
mp.EXPECT().GetLoadState(mock.Anything, mock.Anything).Return(&milvuspb.GetLoadStateResponse{Status: commonErrorStatus}, nil).Once()
mp.EXPECT().GetLoadState(mock.Anything, mock.Anything).Return(&DefaultLoadStateResp, nil).Times(4)
mp.EXPECT().GetLoadState(mock.Anything, mock.Anything).Return(&milvuspb.GetLoadStateResponse{
Status: &StatusSuccess,
State: commonpb.LoadState_LoadStateNotExist,
}, nil).Once()
mp.EXPECT().GetLoadState(mock.Anything, mock.Anything).Return(&milvuspb.GetLoadStateResponse{
Status: &StatusSuccess,
State: commonpb.LoadState_LoadStateNotLoad,
}, nil).Once()
mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&milvuspb.DescribeIndexResponse{Status: commonErrorStatus}, nil).Once()
mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&DefaultDescIndexesReqp, nil).Times(3)
mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(nil, merr.WrapErrIndexNotFoundForCollection(DefaultCollectionName)).Once()
mp.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&milvuspb.DescribeIndexResponse{
Status: merr.Status(merr.WrapErrIndexNotFoundForCollection(DefaultCollectionName)),
}, nil).Once()
mp.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(&milvuspb.GetCollectionStatisticsResponse{
Status: commonSuccessStatus,
Stats: []*commonpb.KeyValuePair{
{Key: "row_count", Value: "0"},
},
}, nil).Once()
mp.EXPECT().GetCollectionStatistics(mock.Anything, mock.Anything).Return(&milvuspb.GetCollectionStatisticsResponse{
Status: commonSuccessStatus,
Stats: []*commonpb.KeyValuePair{
{Key: "row_count", Value: "abc"},
},
}, nil).Once()
mp.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).Return(&milvuspb.GetLoadingProgressResponse{
Status: commonSuccessStatus,
Progress: int64(77),
}, nil).Once()
mp.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).Return(&milvuspb.GetLoadingProgressResponse{
Status: commonSuccessStatus,
Progress: int64(100),
}, nil).Once()
mp.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).Return(&milvuspb.GetLoadingProgressResponse{Status: commonErrorStatus}, nil).Once()
mp.EXPECT().ShowPartitions(mock.Anything, mock.Anything).Return(&milvuspb.ShowPartitionsResponse{
Status: &StatusSuccess,
PartitionNames: []string{DefaultPartitionName},
}, nil).Once()
mp.EXPECT().HasPartition(mock.Anything, mock.Anything).Return(&milvuspb.BoolResponse{
Status: &StatusSuccess,
Value: true,
}, nil).Once()
mp.EXPECT().GetPartitionStatistics(mock.Anything, mock.Anything).Return(&milvuspb.GetPartitionStatisticsResponse{
Status: commonSuccessStatus,
Stats: []*commonpb.KeyValuePair{
{Key: "row_count", Value: "0"},
},
}, nil).Once()
mp.EXPECT().ListCredUsers(mock.Anything, mock.Anything).Return(&milvuspb.ListCredUsersResponse{
Status: &StatusSuccess,
Usernames: []string{util.UserRoot},
}, nil).Once()
mp.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(&milvuspb.SelectUserResponse{
Status: &StatusSuccess,
Results: []*milvuspb.UserResult{
{User: &milvuspb.UserEntity{Name: util.UserRoot}, Roles: []*milvuspb.RoleEntity{
{Name: util.RoleAdmin},
}},
},
}, nil).Once()
mp.EXPECT().SelectRole(mock.Anything, mock.Anything).Return(&milvuspb.SelectRoleResponse{
Status: &StatusSuccess,
Results: []*milvuspb.RoleResult{
{Role: &milvuspb.RoleEntity{Name: util.RoleAdmin}},
},
}, nil).Once()
mp.EXPECT().SelectRole(mock.Anything, mock.Anything).Return(&milvuspb.SelectRoleResponse{
Status: &StatusSuccess,
Results: []*milvuspb.RoleResult{
{Role: &milvuspb.RoleEntity{Name: util.RoleAdmin, Description: "admin role"}},
},
}, nil).Once()
mp.EXPECT().SelectGrant(mock.Anything, mock.Anything).Return(&milvuspb.SelectGrantResponse{
Status: &StatusSuccess,
Entities: []*milvuspb.GrantEntity{
{
Role: &milvuspb.RoleEntity{Name: util.RoleAdmin},
Object: &milvuspb.ObjectEntity{Name: "global"},
ObjectName: "",
DbName: util.DefaultDBName,
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: util.UserRoot},
Privilege: &milvuspb.PrivilegeEntity{Name: "*"},
},
},
},
}, nil).Once()
mp.EXPECT().ListAliases(mock.Anything, mock.Anything).Return(&milvuspb.ListAliasesResponse{Status: commonErrorStatus}, nil).Once()
mp.EXPECT().ListAliases(mock.Anything, mock.Anything).Return(&milvuspb.ListAliasesResponse{
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().ListAliases(mock.Anything, mock.Anything).Return(&milvuspb.ListAliasesResponse{
Status: &StatusSuccess,
Aliases: []string{DefaultAliasName},
}, nil).Once()
mp.EXPECT().DescribeAlias(mock.Anything, mock.Anything).Return(&milvuspb.DescribeAliasResponse{
Status: &StatusSuccess,
Alias: DefaultAliasName,
}, nil).Once()
mp.EXPECT().ListPrivilegeGroups(mock.Anything, mock.Anything).Return(&milvuspb.ListPrivilegeGroupsResponse{
Status: &StatusSuccess,
PrivilegeGroups: []*milvuspb.PrivilegeGroupInfo{{GroupName: "group1", Privileges: []*milvuspb.PrivilegeEntity{{Name: "*"}}}},
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []rawTestCase{}
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, HasAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, HasAction),
errMsg: "",
errCode: 65535,
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, DescribeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, DescribeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, DescribeAction),
errMsg: "",
errCode: 65535,
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, StatsAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, StatsAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, LoadStateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, LoadStateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, LoadStateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, LoadStateAction),
errCode: 100,
errMsg: "collection not found[collection=book]",
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, LoadStateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PartitionCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PartitionCategory, HasAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PartitionCategory, StatsAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(UserCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(UserCategory, DescribeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(RoleCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(RoleCategory, DescribeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(IndexCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(IndexCategory, DescribeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(IndexCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(IndexCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(AliasCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(AliasCategory, DescribeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PrivilegeGroupCategory, ListAction),
})
for _, testcase := range queryTestCases {
t.Run(testcase.path, func(t *testing.T) {
bodyReader := bytes.NewReader([]byte(`{` +
`"collectionName": "` + DefaultCollectionName + `",` +
`"partitionName": "` + DefaultPartitionName + `",` +
`"indexName": "` + DefaultIndexName + `",` +
`"userName": "` + util.UserRoot + `",` +
`"roleName": "` + util.RoleAdmin + `",` +
`"aliasName": "` + DefaultAliasName + `",` +
`"privilegeGroupName": "pg"` +
`}`))
req := httptest.NewRequest(http.MethodPost, testcase.path, bodyReader)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
})
}
}
var commonSuccessStatus = &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Code: merr.Code(nil),
Reason: "",
}
var commonErrorStatus = &commonpb.Status{
ErrorCode: commonpb.ErrorCode_CollectionNameNotFound, // 28
Reason: "",
}
func TestRoleDescriptionV2(t *testing.T) {
paramtable.Init()
mp := mocks.NewMockProxy(t)
mp.EXPECT().CreateRole(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.CreateRoleRequest) (*commonpb.Status, error) {
assert.Equal(t, "test_role", req.GetEntity().GetName())
assert.Equal(t, "role description", req.GetEntity().GetDescription())
return commonSuccessStatus, nil
}).Once()
mp.EXPECT().AlterRole(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.AlterRoleRequest) (*commonpb.Status, error) {
assert.Equal(t, "test_role", req.GetRoleName())
assert.Equal(t, "updated description", req.GetDescription())
return commonSuccessStatus, nil
}).Once()
mp.EXPECT().SelectRole(mock.Anything, mock.Anything).Return(&milvuspb.SelectRoleResponse{
Status: &StatusSuccess,
Results: []*milvuspb.RoleResult{
{Role: &milvuspb.RoleEntity{Name: "test_role", Description: "updated description"}},
},
}, nil).Once()
mp.EXPECT().SelectGrant(mock.Anything, mock.Anything).Return(&milvuspb.SelectGrantResponse{
Status: &StatusSuccess,
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
post := func(path, body string) *httptest.ResponseRecorder {
req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader([]byte(body)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
return w
}
w := post(versionalV2(RoleCategory, CreateAction), `{"roleName": "test_role", "description": "role description"}`)
assert.Equal(t, http.StatusOK, w.Code)
w = post(versionalV2(RoleCategory, AlterAction), `{"roleName": "test_role", "description": "updated description"}`)
assert.Equal(t, http.StatusOK, w.Code)
w = post(versionalV2(RoleCategory, DescribeAction), `{"roleName": "test_role"}`)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := map[string]any{}
assert.Nil(t, json.Unmarshal(w.Body.Bytes(), &returnBody))
assert.EqualValues(t, 0, returnBody[HTTPReturnCode])
assert.Equal(t, "updated description", returnBody[HTTPReturnDescription])
}
func TestMethodDelete(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().DropCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().DropPartition(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().DeleteCredential(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().DropRole(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().DropIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().DropAlias(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().DropPrivilegeGroup(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []rawTestCase{}
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, DropAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PartitionCategory, DropAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(UserCategory, DropAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(RoleCategory, DropAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(IndexCategory, DropAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(AliasCategory, DropAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PrivilegeGroupCategory, DropAction),
})
for _, testcase := range queryTestCases {
t.Run(testcase.path, func(t *testing.T) {
bodyReader := bytes.NewReader([]byte(`{"collectionName": "` + DefaultCollectionName + `", "partitionName": "` + DefaultPartitionName +
`", "userName": "` + util.UserRoot + `", "roleName": "` + util.RoleAdmin + `", "indexName": "` + DefaultIndexName + `", "aliasName": "` + DefaultAliasName + `", "privilegeGroupName": "pg"}`))
req := httptest.NewRequest(http.MethodPost, testcase.path, bodyReader)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
fmt.Println(w.Body.String())
})
}
}
func TestMethodPost(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().CreateCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().RenameCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Times(3)
mp.EXPECT().ReleaseCollection(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().CreatePartition(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().LoadPartitions(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().ReleasePartitions(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().CreateCredential(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().UpdateCredential(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().OperateUserRole(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Twice()
mp.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Times(2)
mp.EXPECT().OperatePrivilegeV2(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Times(2)
mp.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Twice()
mp.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(commonErrorStatus, nil).Once()
mp.EXPECT().CreateAlias(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().AlterAlias(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().CreatePrivilegeGroup(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
mp.EXPECT().OperatePrivilegeGroup(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Twice()
mp.EXPECT().ImportV2(mock.Anything, mock.Anything).Return(&internalpb.ImportResponse{
Status: commonSuccessStatus, JobID: "1234567890",
}, nil).Once()
mp.EXPECT().ListImports(mock.Anything, mock.Anything).Return(&internalpb.ListImportsResponse{
Status: &StatusSuccess,
JobIDs: []string{"1", "2", "3", "4"},
States: []internalpb.ImportJobState{
internalpb.ImportJobState_Pending,
internalpb.ImportJobState_Importing,
internalpb.ImportJobState_Failed,
internalpb.ImportJobState_Completed,
},
Reasons: []string{"", "", "mock reason", ""},
Progresses: []int64{0, 30, 0, 100},
CollectionNames: []string{"AAA", "BBB", "CCC", "DDD"},
}, nil).Once()
mp.EXPECT().GetImportProgress(mock.Anything, mock.Anything).Return(&internalpb.GetImportProgressResponse{
Status: &StatusSuccess,
State: internalpb.ImportJobState_Completed,
Reason: "",
Progress: 100,
}, nil).Twice()
mp.EXPECT().GetSegmentsInfo(mock.Anything, mock.Anything).Return(&internalpb.GetSegmentsInfoResponse{
Status: &StatusSuccess,
SegmentInfos: []*internalpb.SegmentInfo{
{
SegmentID: 3,
CollectionID: 1,
PartitionID: 2,
NumRows: 1000,
VChannel: "ch-1",
State: commonpb.SegmentState_Flushed,
Level: commonpb.SegmentLevel_L1,
IsSorted: true,
InsertLogs: []*internalpb.FieldBinlog{
{
FieldID: 0,
LogIDs: []int64{1, 5, 9},
},
{
FieldID: 1,
LogIDs: []int64{2, 6, 10},
},
{
FieldID: 100,
LogIDs: []int64{3, 7, 11},
},
{
FieldID: 101,
LogIDs: []int64{4, 8, 12},
},
},
DeltaLogs: []*internalpb.FieldBinlog{
{
FieldID: 100,
LogIDs: []int64{13, 14, 15},
},
},
StatsLogs: []*internalpb.FieldBinlog{
{
FieldID: 100,
LogIDs: []int64{16},
},
},
},
},
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []rawTestCase{}
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, CreateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, RenameAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, LoadAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, RefreshLoadAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(CollectionCategory, ReleaseAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PartitionCategory, CreateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PartitionCategory, LoadAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PartitionCategory, ReleaseAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(UserCategory, CreateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(UserCategory, UpdatePasswordAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(UserCategory, GrantRoleAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(UserCategory, RevokeRoleAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(RoleCategory, CreateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(RoleCategory, GrantPrivilegeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(RoleCategory, RevokePrivilegeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(RoleCategory, GrantPrivilegeActionV2),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(RoleCategory, RevokePrivilegeActionV2),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(IndexCategory, CreateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(IndexCategory, CreateAction),
errMsg: "",
errCode: 65535,
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(AliasCategory, CreateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(AliasCategory, AlterAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(ImportJobCategory, CreateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(ImportJobCategory, ListAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(ImportJobCategory, GetProgressAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(ImportJobCategory, DescribeAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PrivilegeGroupCategory, CreateAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PrivilegeGroupCategory, AddPrivilegesToGroupAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(PrivilegeGroupCategory, RemovePrivilegesFromGroupAction),
})
queryTestCases = append(queryTestCases, rawTestCase{
path: versionalV2(SegmentCategory, DescribeAction),
})
for _, testcase := range queryTestCases {
t.Run(testcase.path, func(t *testing.T) {
bodyReader := bytes.NewReader([]byte(`{` +
`"collectionName": "` + DefaultCollectionName + `", "newCollectionName": "test", "newDbName": "",` +
`"partitionName": "` + DefaultPartitionName + `", "partitionNames": ["` + DefaultPartitionName + `"],` +
`"schema": {"fields": [{"fieldName": "book_id", "dataType": "Int64", "elementTypeParams": {}}, {"fieldName": "book_intro", "dataType": "FloatVector", "elementTypeParams": {"dim": 2}}]},` +
`"indexParams": [{"indexName": "` + DefaultIndexName + `", "fieldName": "book_intro", "metricType": "L2", "params": {"nlist": 30, "index_type": "IVF_FLAT"}}],` +
`"userName": "` + util.UserRoot + `", "password": "Milvus", "newPassword": "milvus", "roleName": "` + util.RoleAdmin + `",` +
`"roleName": "` + util.RoleAdmin + `", "objectType": "Global", "objectName": "*", "privilege": "*",` +
`"privilegeGroupName": "pg", "privileges": ["create", "drop"],` +
`"aliasName": "` + DefaultAliasName + `",` +
`"jobId": "1234567890",` +
`"files": [["book.json"]]` +
`}`))
req := httptest.NewRequest(http.MethodPost, testcase.path, bodyReader)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
if testcase.errCode != 0 {
assert.Equal(t, testcase.errMsg, returnBody.Message)
}
fmt.Println(w.Body.String())
})
}
}
func TestUserDescriptionV2(t *testing.T) {
paramtable.Init()
description := "用户描述"
mp := mocks.NewMockProxy(t)
mp.EXPECT().
CreateCredential(mock.Anything, mock.MatchedBy(func(req *milvuspb.CreateCredentialRequest) bool {
return req.GetUsername() == util.UserRoot &&
req.GetPassword() == crypto.Base64Encode("Milvus") &&
req.Description != nil &&
req.GetDescription() == description
})).
Return(commonSuccessStatus, nil).
Once()
mp.EXPECT().
UpdateCredential(mock.Anything, mock.MatchedBy(func(req *milvuspb.UpdateCredentialRequest) bool {
return req.GetUsername() == util.UserRoot &&
req.GetOldPassword() == "" &&
req.GetNewPassword() == "" &&
req.Description != nil &&
req.GetDescription() == description
})).
Return(commonSuccessStatus, nil).
Once()
mp.EXPECT().
SelectUser(mock.Anything, mock.Anything).
Return(&milvuspb.SelectUserResponse{
Status: &StatusSuccess,
Results: []*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{Name: util.UserRoot},
Roles: []*milvuspb.RoleEntity{{Name: util.RoleAdmin}},
Description: description,
},
},
}, nil).
Once()
testEngine := initHTTPServerV2(mp, false)
for _, testcase := range []struct {
path string
body string
}{
{
path: versionalV2(UserCategory, CreateAction),
body: `{"userName":"` + util.UserRoot + `","password":"Milvus","description":"` + description + `"}`,
},
{
path: versionalV2(UserCategory, UpdatePasswordAction),
body: `{"userName":"` + util.UserRoot + `","description":"` + description + `"}`,
},
{
path: versionalV2(UserCategory, DescribeAction),
body: `{"userName":"` + util.UserRoot + `"}`,
},
} {
t.Run(testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader([]byte(testcase.body)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
resp := map[string]interface{}{}
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
assert.Equal(t, float64(0), resp[HTTPReturnCode])
if testcase.path == versionalV2(UserCategory, DescribeAction) {
assert.Equal(t, description, resp[HTTPReturnDescription])
assert.Equal(t, []interface{}{util.RoleAdmin}, resp[HTTPReturnData])
}
})
}
}
func TestCreateImportJobPreservesRBACRoleContext(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
paramtable.Get().Save(proxy.Params.CommonCfg.RootShouldBindRole.Key, "true")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.RootShouldBindRole.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().ImportV2(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, req *internalpb.ImportRequest) (*internalpb.ImportResponse, error) {
roles, ok := ctx.Value(proxy.RBACRoleContextKey).([]string)
assert.True(t, ok)
assert.ElementsMatch(t, []string{util.RoleAdmin, util.RolePublic}, roles)
return &internalpb.ImportResponse{
Status: commonSuccessStatus,
JobID: "1234567890",
}, nil
}).Once()
testEngine := initHTTPServerV2(mp, true)
bodyReader := bytes.NewReader([]byte(`{"collectionName": "` + DefaultCollectionName + `", "files": [["book.json"]]}`))
req := httptest.NewRequest(http.MethodPost, versionalV2(ImportJobCategory, CreateAction), bodyReader)
req.SetBasicAuth(util.UserRoot, getDefaultRootPassword())
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.NoError(t, err)
assert.Equal(t, merr.Code(nil), returnBody.Code)
}
func validateTestCases(t *testing.T, testEngine *gin.Engine, queryTestCases []requestBodyTestCase, allowInt64 bool) {
for i, testcase := range queryTestCases {
t.Run(testcase.path, func(t *testing.T) {
bodyReader := bytes.NewReader(testcase.requestBody)
req := httptest.NewRequest(http.MethodPost, versionalV2(EntityCategory, testcase.path), bodyReader)
if allowInt64 {
req.Header.Set(HTTPHeaderAllowInt64, "true")
}
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code, "case %d: ", i, string(testcase.requestBody))
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err, "case %d: ", i)
assert.Equal(t, testcase.errCode, returnBody.Code, "case: %d, request body: %s ", i, string(testcase.requestBody))
if testcase.errCode != 0 {
assert.Contains(t, returnBody.Message, testcase.errMsg, "case: %d, request body: %s", i, string(testcase.requestBody))
}
fmt.Println(w.Body.String())
})
}
}
func TestDML(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(2)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(3)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(6)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{Status: commonErrorStatus}, nil).Times(4)
mp.EXPECT().Query(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
if matchCountRule(req.OutputFields) {
for _, pair := range req.QueryParams {
if pair.GetKey() == proxy.LimitKey {
return nil, errors.New("mock error")
}
}
}
return &milvuspb.QueryResults{Status: commonSuccessStatus, OutputFields: []string{}, FieldsData: []*schemapb.FieldData{}}, nil
}).Times(4)
mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once()
mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{}}}}}, nil).Once()
mp.EXPECT().Upsert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, UpsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once()
mp.EXPECT().Upsert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, UpsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_StrId{StrId: &schemapb.StringArray{Data: []string{}}}}}, nil).Once()
mp.EXPECT().Delete(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus}, nil).Once()
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, true, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().Upsert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, UpsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []requestBodyTestCase{}
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: QueryAction,
requestBody: []byte(`{"collectionName": "book", "filter": "book_id in [2, 4, 6, 8]", "outputFields": ["book_id", "word_count", "book_intro"], "offset": 1}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: GetAction,
requestBody: []byte(`{"collectionName": "book", "outputFields": ["book_id", "word_count", "book_intro"]}`),
errMsg: "missing required parameters, error: Key: 'CollectionIDReq.ID' Error:Field validation for 'ID' failed on the 'required' tag",
errCode: 1802, // ErrMissingRequiredParameters
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: QueryAction,
requestBody: []byte(`{"collectionName": "book", "filter": "book_id in [2, 4, 6, 8]"}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: QueryAction,
requestBody: []byte(`{"collectionName": "book", "filter": "", "outputFields": ["count(*)"], "limit": 10}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: QueryAction,
requestBody: []byte(`{"collectionName": "book", "filter": "", "outputFields": ["book_id", "word_count", "book_intro"], "limit": 10, "consistencyLevel": "AAA"}`),
errMsg: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded",
errCode: 1100, // ErrParameterInvalid
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: GetAction,
requestBody: []byte(`{"collectionName": "book", "id": [2, 4, 6, 8], "outputFields": ["book_id", "word_count", "book_intro"], "consistencyLevel": "AAA"}`),
errMsg: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded",
errCode: 1100, // ErrParameterInvalid
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: InsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: InsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: UpsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: UpsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: DeleteAction,
requestBody: []byte(`{"collectionName": "book", "filter": "book_id in [0]"}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: DeleteAction,
requestBody: []byte(`{"collectionName": "book", "id" : [0]}`),
errMsg: "missing required parameters, error: Key: 'CollectionFilterReq.Filter' Error:Field validation for 'Filter' failed on the 'required' tag",
errCode: 1802, // ErrMissingRequiredParameters
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: GetAction,
requestBody: []byte(`{"collectionName": "book", "id": [2, 4, 6, 8, 0], "outputFields": ["book_id", "word_count", "book_intro"]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: GetAction,
requestBody: []byte(`{"collectionName": "book", "id": [2, 4, 6, 8, "0"], "outputFields": ["book_id", "word_count", "book_intro"]}`),
errMsg: "",
errCode: 65535,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: InsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
errMsg: "",
errCode: 65535,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: UpsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
errMsg: "",
errCode: 65535,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: DeleteAction,
requestBody: []byte(`{"collectionName": "book", "filter": "book_id in [0]"}`),
errMsg: "",
errCode: 65535,
})
// upsert when autoid==true
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: UpsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
})
validateTestCases(t, testEngine, queryTestCases, false)
}
func TestQueryOrderByFields(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(2)
orderByValues := []string{}
mp.EXPECT().Query(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, req *milvuspb.QueryRequest) (*milvuspb.QueryResults, error) {
for _, pair := range req.GetQueryParams() {
if pair.GetKey() == proxy.OrderByFieldsKey {
orderByValues = append(orderByValues, pair.GetValue())
}
}
return &milvuspb.QueryResults{Status: commonSuccessStatus, OutputFields: []string{}, FieldsData: []*schemapb.FieldData{}}, nil
}).Times(2)
testEngine := initHTTPServerV2(mp, false)
validateTestCases(t, testEngine, []requestBodyTestCase{
{
path: QueryAction,
requestBody: []byte(`{"collectionName": "book", "filter": "book_id > 0", "limit": 10, "orderByFields": ["word_count:desc", "book_id:asc"]}`),
},
{
// no orderByFields -> no order_by_fields query param forwarded
path: QueryAction,
requestBody: []byte(`{"collectionName": "book", "filter": "book_id > 0", "limit": 10}`),
},
}, false)
assert.Equal(t, []string{"word_count:desc,book_id:asc"}, orderByValues)
}
func TestAllowInt64(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []requestBodyTestCase{}
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: InsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: UpsertAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"book_id": 0, "word_count": 0, "book_intro": [0.11825, 0.6]}]}`),
})
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Twice()
mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once()
mp.EXPECT().Upsert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, UpsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Once()
validateTestCases(t, testEngine, queryTestCases, true)
}
func generateCollectionSchemaWithVectorFields() *schemapb.CollectionSchema {
collSchema := generateCollectionSchema(schemapb.DataType_Int64, false, true)
binaryVectorField := generateVectorFieldSchema(schemapb.DataType_BinaryVector)
binaryVectorField.Name = "binaryVector"
float16VectorField := generateVectorFieldSchema(schemapb.DataType_Float16Vector)
float16VectorField.Name = "float16Vector"
bfloat16VectorField := generateVectorFieldSchema(schemapb.DataType_BFloat16Vector)
bfloat16VectorField.Name = "bfloat16Vector"
sparseFloatVectorField := generateVectorFieldSchema(schemapb.DataType_SparseFloatVector)
sparseFloatVectorField.Name = "sparseFloatVector"
int8VectorField := generateVectorFieldSchema(schemapb.DataType_Int8Vector)
int8VectorField.Name = "int8Vector"
collSchema.Fields = append(collSchema.Fields, binaryVectorField)
collSchema.Fields = append(collSchema.Fields, float16VectorField)
collSchema.Fields = append(collSchema.Fields, bfloat16VectorField)
collSchema.Fields = append(collSchema.Fields, sparseFloatVectorField)
collSchema.Fields = append(collSchema.Fields, int8VectorField)
return collSchema
}
func TestFp16Bf16VectorsV2(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
collSchema := generateCollectionSchemaWithVectorFields()
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []requestBodyTestCase{}
for _, path := range []string{InsertAction, UpsertAction} {
queryTestCases = append(queryTestCases,
requestBodyTestCase{
path: path,
requestBody: []byte(
`{
"collectionName": "book",
"data": [
{
"book_id": 0,
"word_count": 0,
"book_intro": [0.11825, 0.6],
"binaryVector": "AQ==",
"float16Vector": [3.0],
"bfloat16Vector": [4.4, 442],
"sparseFloatVector": {"1": 0.1, "2": 0.44},
"int8Vector": [1, 2]
}
]
}`),
errCode: 1804,
errMsg: "fail to deal the insert data, error: []byte size 2 doesn't equal to vector dimension 2 of Float16Vector",
}, requestBodyTestCase{
path: path,
requestBody: []byte(
`{
"collectionName": "book",
"data": [
{
"book_id": 0,
"word_count": 0,
"book_intro": [0.11825, 0.6],
"binaryVector": "AQ==",
"float16Vector": [3, 3.0],
"bfloat16Vector": [4.4, 442],
"sparseFloatVector": {"1": 0.1, "2": 0.44},
"int8Vector": [1, 2]
}
]
}`),
}, requestBodyTestCase{
path: path,
// [3, 3] shouble be converted to [float(3), float(3)]
requestBody: []byte(
`{
"collectionName": "book",
"data": [
{
"book_id": 0,
"word_count": 0,
"book_intro": [0.11825, 0.6],
"binaryVector": "AQ==",
"float16Vector": [3, 3],
"bfloat16Vector": [4.4, 442],
"sparseFloatVector": {"1": 0.1, "2": 0.44},
"int8Vector": [1, 2]
}
]
}`),
}, requestBodyTestCase{
path: path,
requestBody: []byte(
`{
"collectionName": "book",
"data": [
{
"book_id": 0,
"word_count": 0,
"book_intro": [0.11825, 0.6],
"binaryVector": "AQ==",
"float16Vector": "AQIDBA==",
"bfloat16Vector": "AQIDBA==",
"sparseFloatVector": {"1": 0.1, "2": 0.44},
"int8Vector": [1, 2]
}
]
}`),
}, requestBodyTestCase{
path: path,
requestBody: []byte(
`{
"collectionName": "book",
"data": [
{
"book_id": 0,
"word_count": 0,
"book_intro": [0.11825, 0.6],
"binaryVector": "AQ==",
"float16Vector": [3, 3.0, 3],
"bfloat16Vector": [4.4, 44],
"sparseFloatVector": {"1": 0.1, "2": 0.44},
"int8Vector": [1, 2]
}
]
}`),
errMsg: "fail to deal the insert data, error: []byte size 6 doesn't equal to vector dimension 2 of Float16Vector",
errCode: 1804,
}, requestBodyTestCase{
path: path,
requestBody: []byte(
`{
"collectionName": "book",
"data": [
{
"book_id": 0,
"word_count": 0,
"book_intro": [0.11825, 0.6],
"binaryVector": "AQ==",
"float16Vector": [3, 3.0],
"bfloat16Vector": [4.4, 442, 44],
"sparseFloatVector": {"1": 0.1, "2": 0.44},
"int8Vector": [1, 2]
}
]
}`),
errMsg: "fail to deal the insert data, error: []byte size 6 doesn't equal to vector dimension 2 of BFloat16Vector",
errCode: 1804,
}, requestBodyTestCase{
path: path,
requestBody: []byte(
`{
"collectionName": "book",
"data": [
{
"book_id": 0,
"word_count": 0,
"book_intro": [0.11825, 0.6],
"binaryVector": "AQ==",
"float16Vector": "AQIDBA==",
"bfloat16Vector": [4.4, 442],
"sparseFloatVector": {"1": 0.1, "2": 0.44},
"int8Vector": [1, 2]
},
{
"book_id": 1,
"word_count": 0,
"book_intro": [0.11825, 0.6],
"binaryVector": "AQ==",
"float16Vector": [3.1, 3.1],
"bfloat16Vector": "AQIDBA==",
"sparseFloatVector": {"3": 1.1, "2": 0.44},
"int8Vector": [1, 2]
}
]
}`),
})
}
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: collSchema,
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(len(queryTestCases))
mp.EXPECT().Insert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Times(4)
mp.EXPECT().Upsert(mock.Anything, mock.Anything).Return(&milvuspb.MutationResult{Status: commonSuccessStatus, InsertCnt: int64(0), IDs: &schemapb.IDs{IdField: &schemapb.IDs_IntId{IntId: &schemapb.LongArray{Data: []int64{}}}}}, nil).Times(4)
validateTestCases(t, testEngine, queryTestCases, false)
}
func TestSearchV2(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
outputFields := []string{FieldBookID, FieldWordCount, "author", "date"}
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(11)
mp.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{
TopK: int64(3),
OutputFields: outputFields,
FieldsData: generateFieldData(),
Ids: generateIDs(schemapb.DataType_Int64, 3),
Scores: DefaultScores,
}}, nil).Once()
mp.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{TopK: int64(0)}}, nil).Times(3)
mp.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: &commonpb.Status{
ErrorCode: 1700, // ErrFieldNotFound
Reason: "groupBy field not found in schema: field not found[field=test]",
}}, nil).Once()
mp.EXPECT().HybridSearch(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{
TopK: int64(3),
OutputFields: outputFields,
FieldsData: generateFieldData(),
Ids: generateIDs(schemapb.DataType_Int64, 3),
Scores: DefaultScores,
}}, nil).Once()
mp.EXPECT().HybridSearch(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{TopK: int64(0)}}, nil).Times(5)
collSchema := generateCollectionSchemaWithVectorFields()
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: collSchema,
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(15)
mp.EXPECT().Search(mock.Anything, mock.Anything).Return(&milvuspb.SearchResults{Status: commonSuccessStatus, Results: &schemapb.SearchResultData{TopK: int64(0)}}, nil).Times(3)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
Status: &commonpb.Status{
Code: 1100,
Reason: "mock",
},
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []requestBodyTestCase{}
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "filter": "book_id in {list}", "exprParams":{"list": [2, 4, 6, 8]}, "limit": 4, "outputFields": ["word_count"]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"],"consistencyLevel": "Strong"}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"], "params": {"radius":0.9}}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"], "params": {"radius":0.9, "range_filter": 0.1}, "groupingField": "word_count"}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"], "params": {"radius":0.9, "range_filter": 0.1}, "groupingField": "test"}`),
errMsg: "groupBy field not found in schema: field not found[field=test]",
errCode: 65535,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [["0.1", "0.2"]], "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"], "params": {"radius":0.9, "range_filter": 0.1}, "groupingField": "test"}`),
errMsg: "can only accept json format request, error: Mismatch type float32 with value",
errCode: 1801,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3}, {"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3}], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3}, {"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3}], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: HybridSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3}, {"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3}], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
})
// annsField
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "annsField": "word_count", "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"], "params": {"radius":0.9, "range_filter": 0.1}, "groupingField": "test"}`),
errMsg: "cannot find a vector field",
errCode: 1801,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [{"data": [[0.1, 0.2]], "annsField": "float_vector1", "metricType": "L2", "limit": 3}, {"data": [[0.1, 0.2]], "annsField": "float_vector2", "metricType": "L2", "limit": 3}], "rerank": {"strategy": "rrf", "params": {"k": 1}}}`),
errMsg: "cannot find a vector field",
errCode: 1801,
})
// multiple annsFields
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"]}`),
errMsg: "can only accept json format request, error: search without annsField, but already found multiple vector fields: [book_intro, binaryVector,,,]",
errCode: 1801,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "annsField": "book_intro", "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "annsField": "binaryVector", "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"]}`),
errMsg: "can only accept json format request, error: Mismatch type uint8",
errCode: 1801,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": ["AQ=="], "annsField": "binaryVector", "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "filter": "book_id in {list}", "exprParams":{"list": [2, 4, 6, 8]},"metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": [[0.1, 0.23]], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": [[0.1, 0.43]], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
})
// -2, -1, 1, 3 should be float32
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": [[-2, -1]], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": [[1, 3]], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
})
// invalid fp32 vectors for fp16/bf16
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": [[0.23]], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": [[0.1, 0.43]], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
errCode: 1801,
errMsg: "can only accept json format request, error: dimension: 2, but length of []float: 1: invalid parameter[expected=Float16Vector][actual=[[0.23]]]",
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": [[0.23, 4.4]], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": [[0.1]], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
errCode: 1801,
errMsg: "can only accept json format request, error: dimension: 2, but length of []float: 1: invalid parameter[expected=BFloat16Vector][actual=[[0.1]]]",
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "consistencyLevel":"unknown","rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
errMsg: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded, err:parameter:'unknown' is incorrect, please check it: invalid parameter",
errCode: 1100, // ErrParameterInvalid
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2, 0.3]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
errMsg: "can only accept json format request, error: dimension: 2, but length of []float: 3: invalid parameter[expected=FloatVector][actual=[[0.1, 0.2, 0.3]]]",
errCode: 1801,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQID"], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
errMsg: "can only accept json format request, error: dimension: 8, bytesLen: 1, but length of []byte: 3: invalid parameter[expected=BinaryVector][actual=\x01\x02\x03]",
errCode: 1801,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQID"], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
errMsg: "can only accept json format request, error: dimension: 2, bytesLen: 4, but length of []byte: 3: invalid parameter[expected=Float16Vector][actual=\x01\x02\x03]",
errCode: 1801,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName": "hello_milvus", "search": [` +
`{"data": [[0.1, 0.2]], "annsField": "book_intro", "metricType": "L2", "limit": 3},` +
`{"data": ["AQ=="], "annsField": "binaryVector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQIDBA=="], "annsField": "float16Vector", "metricType": "L2", "limit": 3},` +
`{"data": ["AQID"], "annsField": "bfloat16Vector", "metricType": "L2", "limit": 3}` +
`], "rerank": {"strategy": "weighted", "params": {"weights": [0.9, 0.8]}}}`),
errMsg: "can only accept json format request, error: dimension: 2, bytesLen: 4, but length of []byte: 3: invalid parameter[expected=BFloat16Vector][actual=\x01\x02\x03]",
errCode: 1801,
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [{"1": 0.1}], "annsField": "sparseFloatVector", "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"]}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"], "searchParams": {"params":"a"}}`),
errMsg: "searchParams.params must be a dict: invalid parameter",
errCode: 1100, // ErrParameterInvalid
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": [[0.1, 0.2]], "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"],"consistencyLevel": "unknown"}`),
errMsg: "consistencyLevel can only be [Strong, Session, Bounded, Eventually, Customized], default: Bounded, err:parameter:'unknown' is incorrect, please check it: invalid parameter",
errCode: 1100, // ErrParameterInvalid
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "data": ["AQ=="], "annsField": "binaryVector", "filter": "book_id in [2, 4, 6, 8]", "limit": 4, "outputFields": ["word_count"]}`),
errMsg: "mock",
errCode: 1100, // ErrParameterInvalid
})
validateTestCases(t, testEngine, queryTestCases, false)
}
func TestGetQuotaMetrics(t *testing.T) {
paramtable.Init()
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetQuotaMetrics(mock.Anything, mock.Anything).Return(&internalpb.GetQuotaMetricsResponse{
Status: &StatusSuccess,
MetricsInfo: `{"proxy1": "1000", "proxy2": "2000"}`,
}, nil)
testEngine := initHTTPServerV2(mp, false)
testcase := requestBodyTestCase{
path: DescribeAction,
requestBody: []byte(`{}`),
}
bodyReader := bytes.NewReader(testcase.requestBody)
req := httptest.NewRequest(http.MethodPost, versionalV2(QuotaCenterCategory, testcase.path), bodyReader)
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code, "case %d: ", string(testcase.requestBody))
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code, "case: %d, request body: %s ", string(testcase.requestBody))
if testcase.errCode != 0 {
assert.Contains(t, returnBody.Message, testcase.errMsg, "case: %d, request body: %s", string(testcase.requestBody))
}
fmt.Println(w.Body.String())
}
func TestSearchAggregationV2(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
paramtable.Get().Save(proxy.Params.QueryNodeCfg.StorageUsageTrackingEnabled.Key, "true")
defer paramtable.Get().Reset(proxy.Params.QueryNodeCfg.StorageUsageTrackingEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Once()
mp.EXPECT().Search(mock.Anything, mock.MatchedBy(func(req *milvuspb.SearchRequest) bool {
agg := req.GetSearchAggregation()
return agg != nil &&
len(agg.GetFields()) == 1 &&
agg.GetFields()[0] == "brand" &&
agg.GetSize() == 2 &&
agg.GetSearchSize() == 4 &&
agg.GetMetrics()["avg_price"].GetOp() == "avg" &&
agg.GetMetrics()["avg_price"].GetFieldName() == "price" &&
agg.GetOrder()[0].GetKey() == "avg_price" &&
agg.GetOrder()[0].GetDirection() == "desc" &&
agg.GetTopHits().GetSize() == 1 &&
agg.GetTopHits().GetSort()[0].GetFieldName() == "_score" &&
agg.GetTopHits().GetSort()[0].GetDirection() == "asc" &&
agg.GetSubAggregation().GetFields()[0] == "color" &&
req.GetIds() == nil
})).Return(&milvuspb.SearchResults{
Status: &commonpb.Status{
ErrorCode: commonpb.ErrorCode_Success,
Code: merr.Code(nil),
ExtraInfo: map[string]string{
"scanned_remote_bytes": "11",
"scanned_total_bytes": "22",
"cache_hit_ratio": "0.5",
},
},
Results: &schemapb.SearchResultData{
TopK: 0,
NumQueries: 1,
AggTopks: []int64{1},
Recalls: []float32{0.75},
AggBuckets: []*schemapb.AggBucket{
{
Key: []*schemapb.BucketKeyEntry{{FieldId: 101, FieldName: "brand", Value: &schemapb.BucketKeyEntry_StringVal{StringVal: "acme"}}},
Count: 3,
Metrics: map[string]*schemapb.MetricValue{"avg_price": {Value: &schemapb.MetricValue_DoubleVal{DoubleVal: 12.5}}},
Hits: []*schemapb.AggHit{
{
Pk: &schemapb.AggHit_IntPk{IntPk: 10},
Score: 0.9,
Fields: []*schemapb.AggHitField{
{FieldId: 201, FieldName: "price", Value: &schemapb.AggHitField_IntVal{IntVal: 99}},
},
},
},
},
},
},
}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
req := httptest.NewRequest(http.MethodPost, versionalV2(EntityCategory, SearchAction), bytes.NewReader([]byte(`{
"collectionName": "book",
"data": [[0.1, 0.2]],
"annsField": "book_intro",
"limit": 10,
"searchAggregation": {
"fields": [" brand "],
"size": 2,
"searchSize": 4,
"metrics": {" avg_price ": {"op": " avg ", "fieldName": " price "}},
"order": [{"key": " avg_price ", "direction": " desc "}],
"topHits": {"size": 1, "sort": [{"fieldName": " _score ", "direction": " asc "}]},
"subAggregation": {"fields": ["color"], "size": 1}
}
}`)))
req.Header.Set(HTTPHeaderAllowInt64, "true")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
resp := map[string]interface{}{}
err := json.Unmarshal(w.Body.Bytes(), &resp)
assert.NoError(t, err)
assert.Equal(t, float64(0), resp[HTTPReturnCode])
assert.Equal(t, []interface{}{float64(1)}, resp[HTTPReturnAggTopks])
assert.Equal(t, []interface{}{float64(0.75)}, resp[HTTPReturnRecalls])
assert.Equal(t, float64(11), resp[HTTPReturnScannedRemoteBytes])
assert.Equal(t, float64(22), resp[HTTPReturnScannedTotalBytes])
assert.Equal(t, float64(0.5), resp[HTTPReturnCacheHitRatio])
data := resp[HTTPReturnData].([]interface{})
buckets := data[0].(map[string]interface{})["buckets"].([]interface{})
bucket := buckets[0].(map[string]interface{})
assert.Equal(t, float64(3), bucket["count"])
assert.Equal(t, map[string]interface{}{"avg_price": float64(12.5)}, bucket["metrics"])
hits := bucket["hits"].([]interface{})
assert.Equal(t, float64(10), hits[0].(map[string]interface{})[FieldBookID])
assert.Equal(t, float64(99), hits[0].(map[string]interface{})["price"])
}
func TestSearchAggregationV2UnsupportedCombinations(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(5)
testEngine := initHTTPServerV2(mp, false)
searchAggregation := `"searchAggregation": {"fields": ["brand"], "size": 1}`
testCases := []requestBodyTestCase{
{
path: SearchAction,
requestBody: []byte(`{"collectionName":"book","ids":[1],` + searchAggregation + `}`),
errCode: 1100,
errMsg: "ids and searchAggregation cannot be used simultaneously",
},
{
path: SearchAction,
requestBody: []byte(`{"collectionName":"book","data":[[0.1,0.2]],"annsField":"book_intro","limit":10,"offset":1,` + searchAggregation + `}`),
errCode: 1100,
errMsg: "offset is not supported with searchAggregation",
},
{
path: SearchAction,
requestBody: []byte(`{"collectionName":"book","data":[[0.1,0.2]],"annsField":"book_intro","limit":10,"groupingField":"brand",` + searchAggregation + `}`),
errCode: 1100,
errMsg: "groupingField/groupSize/strictGroupSize and searchAggregation cannot be used simultaneously",
},
{
path: SearchAction,
requestBody: []byte(`{"collectionName":"book","data":[[0.1,0.2]],"annsField":"book_intro","limit":10,"searchParams":{"offset":1},` + searchAggregation + `}`),
errCode: 1100,
errMsg: "searchParams.offset is not supported with searchAggregation",
},
{
path: SearchAction,
requestBody: []byte(`{"collectionName":"book","data":[[0.1,0.2]],"annsField":"book_intro","limit":10,"searchParams":{"params":{"group_by_fields":"[\"brand\"]"}},` + searchAggregation + `}`),
errCode: 1100,
errMsg: "searchParams.group_by_field(s) and searchAggregation cannot be used simultaneously",
},
{
path: HybridSearchAction,
requestBody: []byte(`{"collectionName":"book","search":[{"data":[[0.1,0.2]],"annsField":"book_intro","limit":10}],"rerank":{"strategy":"rrf","params":{}},` + searchAggregation + `}`),
errCode: 1100,
errMsg: "searchAggregation is not supported for hybrid search",
},
{
path: AdvancedSearchAction,
requestBody: []byte(`{"collectionName":"book","search":[{"data":[[0.1,0.2]],"annsField":"book_intro","limit":10,` + searchAggregation + `}],"rerank":{"strategy":"rrf","params":{}}}`),
errCode: 1100,
errMsg: "searchAggregation is not supported for hybrid search",
},
}
validateTestCases(t, testEngine, testCases, false)
}
type AddCollectionFieldSuite struct {
suite.Suite
testEngine *gin.Engine
mp *mocks.MockProxy
}
func (s *AddCollectionFieldSuite) SetupSuite() {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
}
func (s *AddCollectionFieldSuite) TearDownSuite() {
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
}
func (s *AddCollectionFieldSuite) SetupTest() {
s.mp = mocks.NewMockProxy(s.T())
s.testEngine = initHTTPServerV2(s.mp, false)
}
func (s *AddCollectionFieldSuite) TestAddCollectionFieldNormal() {
addFieldTestCases := []requestBodyTestCase{
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "Int64", "nullable": true, "elementTypeParams": {}}}`),
},
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "VarChar", "nullable": true, "elementTypeParams": {"max_length": "256"}}}`),
},
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "Int64", "nullable": true, "defaultValue": 42}}`),
},
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "Array", "elementDataType": "Int64", "nullable": true}}`),
},
}
s.mp.EXPECT().AddCollectionField(mock.Anything, mock.Anything).Return(merr.Success(), nil)
validateRequestBodyTestCases(s.T(), s.testEngine, addFieldTestCases, false)
}
func (s *AddCollectionFieldSuite) TestAddCollectionStructFieldNormal() {
testCases := []struct {
name string
requestBody []byte
}{
{
name: "array_with_struct_element",
requestBody: []byte(`{"collectionName": "book", "schema": {
"fieldName": "clips",
"description": "clip metadata",
"dataType": "Array",
"elementDataType": "Struct",
"nullable": true,
"elementTypeParams": {"max_capacity": 16},
"fields": [
{"fieldName": "tag", "dataType": "Array", "elementDataType": "VarChar", "elementTypeParams": {"max_length": 64}},
{"fieldName": "emb", "dataType": "ArrayOfVector", "elementDataType": "FloatVector", "elementTypeParams": {"dim": 8}}
]
}}`),
},
{
name: "array_of_struct",
requestBody: []byte(`{"collectionName": "book", "schema": {
"fieldName": "clips",
"dataType": "ArrayOfStruct",
"nullable": true,
"typeParams": {"max_capacity": 16},
"fields": [
{"fieldName": "tag", "dataType": "Array", "elementDataType": "VarChar", "elementTypeParams": {"max_length": 64}},
{"fieldName": "emb", "dataType": "ArrayOfVector", "elementDataType": "FloatVector", "elementTypeParams": {"dim": 8}}
]
}}`),
},
}
for _, tc := range testCases {
s.Run(tc.name, func() {
s.mp.EXPECT().AddCollectionStructField(mock.Anything, mock.MatchedBy(func(req *milvuspb.AddCollectionStructFieldRequest) bool {
if req.GetCollectionName() != "book" {
return false
}
structSchema := req.GetStructArrayFieldSchema()
if structSchema.GetName() != "clips" || !structSchema.GetNullable() || len(structSchema.GetFields()) != 2 {
return false
}
if kvPairsToMap(structSchema.GetTypeParams())[common.MaxCapacityKey] != "16" {
return false
}
tag := structSchema.GetFields()[0]
tagParams := kvPairsToMap(tag.GetTypeParams())
if tag.GetName() != "tag" || tag.GetDataType() != schemapb.DataType_Array || tag.GetElementType() != schemapb.DataType_VarChar {
return false
}
if tagParams[common.MaxLengthKey] != "64" || tagParams[common.MaxCapacityKey] != "16" {
return false
}
emb := structSchema.GetFields()[1]
embParams := kvPairsToMap(emb.GetTypeParams())
if emb.GetName() != "emb" || emb.GetDataType() != schemapb.DataType_ArrayOfVector || emb.GetElementType() != schemapb.DataType_FloatVector {
return false
}
return embParams[common.DimKey] == "8" && embParams[common.MaxCapacityKey] == "16"
})).Return(merr.Success(), nil).Once()
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionStructFieldCategory, AddAction), bytes.NewReader(tc.requestBody))
w := httptest.NewRecorder()
s.testEngine.ServeHTTP(w, req)
s.Equal(http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
s.NoError(err)
s.Equal(int32(0), returnBody.Code)
})
}
}
func (s *AddCollectionFieldSuite) TestAddCollectionStructFieldFail() {
s.Run("reject_non_struct_array_schema", func() {
requestBody := []byte(`{"collectionName": "book", "schema": {
"fieldName": "bad_scalar",
"dataType": "Int64",
"fields": [
{"fieldName": "tag", "dataType": "Array", "elementDataType": "VarChar", "elementTypeParams": {"max_length": 64}}
]
}}`)
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionStructFieldCategory, AddAction), bytes.NewReader(requestBody))
w := httptest.NewRecorder()
s.testEngine.ServeHTTP(w, req)
s.Equal(http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
s.NoError(err)
s.Equal(merr.Code(merr.ErrParameterInvalid), returnBody.Code)
s.Contains(returnBody.Message, "must use ArrayOfStruct or Array with Struct element")
})
s.Run("invalid_struct_schema", func() {
requestBody := []byte(`{"collectionName": "book", "schema": {
"fieldName": "clips",
"dataType": "ArrayOfStruct",
"nullable": true,
"typeParams": {"max_capacity": 16}
}}`)
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionStructFieldCategory, AddAction), bytes.NewReader(requestBody))
w := httptest.NewRecorder()
s.testEngine.ServeHTTP(w, req)
s.Equal(http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
s.NoError(err)
s.Equal(merr.Code(merr.ErrParameterInvalid), returnBody.Code)
s.Contains(returnBody.Message, "must contain at least one sub-field")
})
s.Run("server_error", func() {
requestBody := []byte(`{"collectionName": "book", "schema": {
"fieldName": "clips",
"dataType": "ArrayOfStruct",
"nullable": true,
"typeParams": {"max_capacity": 16},
"fields": [
{"fieldName": "tag", "dataType": "Array", "elementDataType": "VarChar", "elementTypeParams": {"max_length": 64}}
]
}}`)
s.mp.EXPECT().AddCollectionStructField(mock.Anything, mock.Anything).Return(merr.Status(merr.WrapErrServiceInternal("mock error")), nil).Once()
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionStructFieldCategory, AddAction), bytes.NewReader(requestBody))
w := httptest.NewRecorder()
s.testEngine.ServeHTTP(w, req)
s.Equal(http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
s.NoError(err)
s.Equal(merr.Code(merr.ErrServiceInternal), returnBody.Code)
s.Contains(returnBody.Message, "mock error")
})
}
func (s *AddCollectionFieldSuite) TestAddCollectionFieldRejectsStructArray() {
requestBody := []byte(`{"collectionName": "book", "schema": {
"fieldName": "clips",
"dataType": "ArrayOfStruct",
"nullable": true,
"typeParams": {"max_capacity": 16},
"fields": [
{"fieldName": "tag", "dataType": "Array", "elementDataType": "VarChar", "elementTypeParams": {"max_length": 64}}
]
}}`)
req := httptest.NewRequest(http.MethodPost, versionalV2(CollectionFieldCategory, AddAction), bytes.NewReader(requestBody))
w := httptest.NewRecorder()
s.testEngine.ServeHTTP(w, req)
s.Equal(http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
s.NoError(err)
s.Equal(merr.Code(merr.ErrParameterInvalid), returnBody.Code)
s.Contains(returnBody.Message, "/v2/vectordb/collections/struct_fields/add")
}
func (s *AddCollectionFieldSuite) TestAddCollectionFieldFail() {
s.Run("bad_request", func() {
addFieldTestCases := []requestBodyTestCase{
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "nullable": true, "elementTypeParams": {}}}`),
errCode: 1802, // missing param
errMsg: "missing required parameters, error: Key: 'CollectionFieldReqWithSchema.Schema.DataType' Error:Field validation for 'DataType' failed on the 'required' tag",
},
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "Int64", "nullable": true, "defaultValue": "aaa"}}`),
errCode: 1100, // invalid param
errMsg: "convert defaultValue fail, err: cannot use \"aaa\"(type: string) as long default value: invalid parameter: invalid parameter",
},
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "LONGLONGLONGLONGTEXT", "nullable": true}}`),
errCode: 1100, // invalid param
errMsg: "data type LONGLONGLONGLONGTEXT is invalid(case sensitive): invalid parameter",
},
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "Array", "nullable": true}}`),
errCode: 1100, // invalid param
errMsg: "element data type is invalid(case sensitive): invalid parameter",
},
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "Array", "elementDataType": "MYBLOB", "nullable": true}}`),
errCode: 1100, // missing param
errMsg: "element data type MYBLOB is invalid(case sensitive): invalid parameter",
},
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field"`),
errCode: 1801, // bad request
errMsg: "can only accept json format request, error: unexpected EOF",
},
}
validateRequestBodyTestCases(s.T(), s.testEngine, addFieldTestCases, false)
})
s.Run("server_error", func() {
addFieldTestCases := []requestBodyTestCase{
{
path: versionalV2(CollectionFieldCategory, AddAction),
requestBody: []byte(`{"collectionName": "book", "schema": {"fieldName": "new_field", "dataType": "Int64", "nullable": true, "elementTypeParams": {}}}`),
errCode: 5,
errMsg: "service internal error: mock error",
},
}
s.mp.EXPECT().AddCollectionField(mock.Anything, mock.Anything).Return(merr.Status(merr.WrapErrServiceInternal("mock error")), nil).Maybe()
validateRequestBodyTestCases(s.T(), s.testEngine, addFieldTestCases, false)
})
}
func TestAddCollectionFieldSuite(t *testing.T) {
suite.Run(t, new(AddCollectionFieldSuite))
}
func kvPairsToMap(kvs []*commonpb.KeyValuePair) map[string]string {
ret := make(map[string]string, len(kvs))
for _, kv := range kvs {
ret[kv.GetKey()] = kv.GetValue()
}
return ret
}
func TestCollectionFunctionSuite(t *testing.T) {
suite.Run(t, new(CollectionFunctionSuite))
}
type CollectionFunctionSuite struct {
suite.Suite
testEngine *gin.Engine
mp *mocks.MockProxy
}
func (s *CollectionFunctionSuite) SetupSuite() {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
}
func (s *CollectionFunctionSuite) TearDownSuite() {
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
}
func (s *CollectionFunctionSuite) SetupTest() {
s.mp = mocks.NewMockProxy(s.T())
s.testEngine = initHTTPServerV2(s.mp, false)
}
func (s *CollectionFunctionSuite) TestAddCollectionFunctionNormal() {
s.Run("success", func() {
addFunctionTestCases := []requestBodyTestCase{
{
path: versionalV2(CollectionCategory, AddFunctionAction),
requestBody: []byte(`{"dbName": "db", "collectionName": "coll", "function": {"name": "test_function", "type": "TextEmbedding", "inputFieldNames": [], "OutputFieldNames": []}}`),
errCode: 0,
errMsg: "",
},
}
s.mp.EXPECT().AddCollectionFunction(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil).Maybe()
validateRequestBodyTestCases(s.T(), s.testEngine, addFunctionTestCases, false)
})
s.Run("bad_request", func() {
addFunctionTestCases := []requestBodyTestCase{
{
path: versionalV2(CollectionCategory, AddFunctionAction),
requestBody: []byte(`{"dbName": "db", "collectionName": "", "function": {"name": "test_function", "type": "BM25", "inputFieldNames": [], "OutputFieldNames": []}}`),
errCode: 1802,
errMsg: "missing required parameters, error: Key: 'CollectionAddFunction.CollectionName' Error:Field validation for 'CollectionName' failed on the 'required' tag",
},
{
path: versionalV2(CollectionCategory, AddFunctionAction),
requestBody: []byte(`invalid json`),
errCode: 1801,
errMsg: "can only accept json format request, error: invalid character 'i' looking for beginning of value",
},
}
validateRequestBodyTestCases(s.T(), s.testEngine, addFunctionTestCases, false)
})
}
func (s *CollectionFunctionSuite) TestAlterCollectionFunctionNormal() {
s.Run("success", func() {
alterFunctionTestCases := []requestBodyTestCase{
{
path: versionalV2(CollectionCategory, AlterFunctionAction),
requestBody: []byte(`{"dbName": "db", "collectionName": "coll", "functionName": "test_function", "function": {"name": "test_function", "type": "TextEmbedding", "inputFieldNames": [], "OutputFieldNames": []}}`),
errCode: 0,
errMsg: "",
},
}
s.mp.EXPECT().AlterCollectionFunction(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil).Maybe()
validateRequestBodyTestCases(s.T(), s.testEngine, alterFunctionTestCases, false)
})
s.Run("bad_request", func() {
alterFunctionTestCases := []requestBodyTestCase{
{
path: versionalV2(CollectionCategory, AlterFunctionAction),
requestBody: []byte(`{"dbName": "db", "collectionName": "", "functionName": "test_function", "function": {"name": "test_function", "type": "BM25", "inputFieldNames": [], "OutputFieldNames": []}}`),
errCode: 1802,
errMsg: "missing required parameters, error: Key: 'CollectionAlterFunction.CollectionName' Error:Field validation for 'CollectionName' failed on the 'required' tag",
},
{
path: versionalV2(CollectionCategory, AlterFunctionAction),
requestBody: []byte(`invalid json`),
errCode: 1801,
errMsg: "can only accept json format request, error: invalid character 'i' looking for beginning of value",
},
}
validateRequestBodyTestCases(s.T(), s.testEngine, alterFunctionTestCases, false)
})
}
func (s *CollectionFunctionSuite) TestDropCollectionFunctionNormal() {
s.Run("success", func() {
addFunctionTestCases := []requestBodyTestCase{
{
path: versionalV2(CollectionCategory, DropFunctionAction),
requestBody: []byte(`{"dbName": "db", "collectionName": "coll", "functionName": "test"}`),
errCode: 0,
errMsg: "",
},
}
s.mp.EXPECT().DropCollectionFunction(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil).Maybe()
validateRequestBodyTestCases(s.T(), s.testEngine, addFunctionTestCases, false)
})
}
func TestTruncateCollection(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().TruncateCollection(mock.Anything, mock.Anything).Return(&milvuspb.TruncateCollectionResponse{
Status: commonSuccessStatus,
}, nil).Times(2)
testEngine := initHTTPServerV2(mp, false)
testCases := []requestBodyTestCase{
{
path: versionalV2(CollectionCategory, TruncateAction),
requestBody: []byte(`{"dbName": "default", "collectionName": "` + DefaultCollectionName + `"}`),
errCode: 0,
errMsg: "",
},
{
path: versionalV2(CollectionCategory, TruncateAction),
requestBody: []byte(`{"collectionName": "` + DefaultCollectionName + `"}`),
errCode: 0,
errMsg: "",
},
{
path: versionalV2(CollectionCategory, TruncateAction),
requestBody: []byte(`{"dbName": "db", "collectionName": ""}`),
errCode: 1802,
errMsg: "missing required parameters",
},
{
path: versionalV2(CollectionCategory, TruncateAction),
requestBody: []byte(`invalid json`),
errCode: 1801,
errMsg: "can only accept json format request",
},
}
for _, testcase := range testCases {
sendReqAndVerify(t, testEngine, testcase.path, http.MethodPost, testcase)
}
}
func TestSearchByPK(t *testing.T) {
paramtable.Init()
// disable rate limit
paramtable.Get().Save(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QuotaConfig.QuotaAndLimitsEnabled.Key)
outputFields := []string{FieldBookID, FieldWordCount}
mp := mocks.NewMockProxy(t)
// Mock for successful search by PK with int64 IDs
mp.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&milvuspb.DescribeCollectionResponse{
CollectionName: DefaultCollectionName,
Schema: generateCollectionSchema(schemapb.DataType_Int64, false, true),
ShardsNum: ShardNumDefault,
Status: &StatusSuccess,
}, nil).Times(5)
mp.EXPECT().Search(mock.Anything, mock.MatchedBy(func(req *milvuspb.SearchRequest) bool {
// Verify the SearchInput is set with IDs
return req.GetIds() != nil
})).Return(&milvuspb.SearchResults{
Status: commonSuccessStatus,
Results: &schemapb.SearchResultData{
TopK: int64(3),
OutputFields: outputFields,
FieldsData: generateFieldData(),
Ids: generateIDs(schemapb.DataType_Int64, 3),
Scores: DefaultScores,
},
}, nil).Times(2)
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []requestBodyTestCase{}
// Test case 1: Search by PK with int64 IDs (JSON numbers decoded as float64)
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "ids": [1, 2, 3], "limit": 10, "outputFields": ["word_count"]}`),
})
// Test case 2: Search by PK with string IDs for int64 PK
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "ids": ["1", "2", "3"], "limit": 10, "outputFields": ["word_count"]}`),
})
// Test case 3: Search by PK with fractional float64 should fail
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "ids": [1.5, 2.9], "limit": 10, "outputFields": ["word_count"]}`),
errMsg: "has fractional part",
errCode: 1100, // ErrParameterInvalid
})
// Test case 4: Search by PK with empty ids array should fail
// Empty array is treated as "no ids provided" at request validation level
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "ids": [], "limit": 10, "outputFields": ["word_count"]}`),
errMsg: "either 'ids' (for primary key search) or 'data' (for vector search) must be provided",
errCode: 1802, // ErrIncorrectParameterFormat
})
// Test case 5: Search by PK with invalid string value should fail
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: SearchAction,
requestBody: []byte(`{"collectionName": "book", "ids": ["not_a_number"], "limit": 10, "outputFields": ["word_count"]}`),
errMsg: "invalid int64 id",
errCode: 1100, // ErrParameterInvalid
})
validateTestCases(t, testEngine, queryTestCases, false)
}
func TestCommitImportJob(t *testing.T) {
paramtable.Init()
t.Run("authorization disabled", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "false")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().CommitImport(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []requestBodyTestCase{}
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: versionalV2(ImportJobCategory, CommitAction),
requestBody: []byte(`{"jobId": "123"}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: versionalV2(ImportJobCategory, CommitAction),
requestBody: []byte(`{"jobId": "not-a-number"}`),
errCode: 1100, // ErrParameterInvalid
})
for _, testcase := range queryTestCases {
t.Run(testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
})
}
})
t.Run("authorization enabled denies user without import privilege", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetImportProgress(mock.Anything, mock.Anything).Return(&internalpb.GetImportProgressResponse{
Status: &StatusSuccess,
CollectionName: DefaultCollectionName,
}, nil).Once()
testEngine := initHTTPServerV2(mp, true)
req := httptest.NewRequest(http.MethodPost, versionalV2(ImportJobCategory, CommitAction), bytes.NewReader([]byte(`{"jobId": "123"}`)))
req.SetBasicAuth("test", "test")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusForbidden, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, int32(65535), returnBody.Code)
assert.Contains(t, returnBody.Message, "PrivilegeImport: permission deny to test")
mp.AssertNotCalled(t, "CommitImport", mock.Anything, mock.Anything)
})
t.Run("authorization enabled returns get progress error", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetImportProgress(mock.Anything, mock.Anything).Return(nil, merr.ErrImportFailed).Once()
testEngine := initHTTPServerV2(mp, true)
req := httptest.NewRequest(http.MethodPost, versionalV2(ImportJobCategory, CommitAction), bytes.NewReader([]byte(`{"jobId": "123"}`)))
req.SetBasicAuth(util.UserRoot, getDefaultRootPassword())
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, merr.Code(merr.ErrImportFailed), returnBody.Code)
mp.AssertNotCalled(t, "CommitImport", mock.Anything, mock.Anything)
})
t.Run("authorization enabled allows root", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetImportProgress(mock.Anything, mock.Anything).Return(&internalpb.GetImportProgressResponse{
Status: &StatusSuccess,
CollectionName: DefaultCollectionName,
}, nil).Once()
mp.EXPECT().CommitImport(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
testEngine := initHTTPServerV2(mp, true)
req := httptest.NewRequest(http.MethodPost, versionalV2(ImportJobCategory, CommitAction), bytes.NewReader([]byte(`{"jobId": "123"}`)))
req.SetBasicAuth(util.UserRoot, getDefaultRootPassword())
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, int32(0), returnBody.Code)
})
}
func TestCreateImportJobAutoCommitOptionValidation(t *testing.T) {
paramtable.Init()
t.Run("accepts omitted true and false strings", func(t *testing.T) {
mp := mocks.NewMockProxy(t)
mp.EXPECT().ImportV2(mock.Anything, mock.MatchedBy(func(req *internalpb.ImportRequest) bool {
return len(req.GetOptions()) == 0
})).Return(&internalpb.ImportResponse{Status: commonSuccessStatus, JobID: "0"}, nil).Once()
mp.EXPECT().ImportV2(mock.Anything, mock.MatchedBy(func(req *internalpb.ImportRequest) bool {
return kvPairsToMap(req.GetOptions())["auto_commit"] == "true"
})).Return(&internalpb.ImportResponse{Status: commonSuccessStatus, JobID: "1"}, nil).Once()
mp.EXPECT().ImportV2(mock.Anything, mock.MatchedBy(func(req *internalpb.ImportRequest) bool {
return kvPairsToMap(req.GetOptions())["auto_commit"] == "false"
})).Return(&internalpb.ImportResponse{Status: commonSuccessStatus, JobID: "2"}, nil).Once()
mp.EXPECT().ImportV2(mock.Anything, mock.MatchedBy(func(req *internalpb.ImportRequest) bool {
return kvPairsToMap(req.GetOptions())["auto_commit"] == "False"
})).Return(&internalpb.ImportResponse{Status: commonSuccessStatus, JobID: "3"}, nil).Once()
mp.EXPECT().ImportV2(mock.Anything, mock.MatchedBy(func(req *internalpb.ImportRequest) bool {
return kvPairsToMap(req.GetOptions())["priority"] == ""
})).Return(&internalpb.ImportResponse{Status: commonSuccessStatus, JobID: "4"}, nil).Once()
testEngine := initHTTPServerV2(mp, false)
for _, body := range []string{
`{"collectionName": "book", "files": [["a.parquet"]]}`,
`{"collectionName": "book", "files": [["a.parquet"]], "options": {"auto_commit": "true"}}`,
`{"collectionName": "book", "files": [["a.parquet"]], "options": {"auto_commit": "false"}}`,
`{"collectionName": "book", "files": [["a.parquet"]], "options": {"auto_commit": "False"}}`,
`{"collectionName": "book", "files": [["a.parquet"]], "options": {"priority": null}}`,
} {
req := httptest.NewRequest(http.MethodPost, versionalV2(ImportJobCategory, CreateAction), bytes.NewReader([]byte(body)))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, merr.Code(nil), returnBody.Code)
}
})
for _, testcase := range []requestBodyTestCase{
{
path: versionalV2(ImportJobCategory, CreateAction),
requestBody: []byte(`{"collectionName": "book", "files": [["a.parquet"]], "options": {"auto_commit": null}}`),
errCode: merr.Code(merr.ErrIncorrectParameterFormat),
errMsg: "options.auto_commit must be one of \"true\" or \"false\"",
},
{
path: versionalV2(ImportJobCategory, CreateAction),
requestBody: []byte(`{"collectionName": "book", "files": [["a.parquet"]], "options": {"auto_commit": ""}}`),
errCode: merr.Code(merr.ErrIncorrectParameterFormat),
errMsg: "options.auto_commit must be one of \"true\" or \"false\"",
},
{
path: versionalV2(ImportJobCategory, CreateAction),
requestBody: []byte(`{"collectionName": "book", "files": [["a.parquet"]], "options": {"auto_commit": "yes"}}`),
errCode: merr.Code(merr.ErrIncorrectParameterFormat),
errMsg: "options.auto_commit must be one of \"true\" or \"false\"",
},
} {
t.Run(string(testcase.requestBody), func(t *testing.T) {
mp := mocks.NewMockProxy(t)
testEngine := initHTTPServerV2(mp, false)
sendReqAndVerify(t, testEngine, testcase.path, http.MethodPost, testcase)
})
}
}
func TestAbortImportJob(t *testing.T) {
paramtable.Init()
t.Run("authorization disabled", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "false")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().AbortImport(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
testEngine := initHTTPServerV2(mp, false)
queryTestCases := []requestBodyTestCase{}
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: versionalV2(ImportJobCategory, AbortAction),
requestBody: []byte(`{"jobId": "123"}`),
})
queryTestCases = append(queryTestCases, requestBodyTestCase{
path: versionalV2(ImportJobCategory, AbortAction),
requestBody: []byte(`{"jobId": "not-a-number"}`),
errCode: 1100, // ErrParameterInvalid
})
for _, testcase := range queryTestCases {
t.Run(testcase.path, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, testcase.path, bytes.NewReader(testcase.requestBody))
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, testcase.errCode, returnBody.Code)
})
}
})
t.Run("authorization enabled denies user without import privilege", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetImportProgress(mock.Anything, mock.Anything).Return(&internalpb.GetImportProgressResponse{
Status: &StatusSuccess,
CollectionName: DefaultCollectionName,
}, nil).Once()
testEngine := initHTTPServerV2(mp, true)
req := httptest.NewRequest(http.MethodPost, versionalV2(ImportJobCategory, AbortAction), bytes.NewReader([]byte(`{"jobId": "123"}`)))
req.SetBasicAuth("test", "test")
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusForbidden, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, int32(65535), returnBody.Code)
assert.Contains(t, returnBody.Message, "PrivilegeImport: permission deny to test")
mp.AssertNotCalled(t, "AbortImport", mock.Anything, mock.Anything)
})
t.Run("authorization enabled returns get progress error", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetImportProgress(mock.Anything, mock.Anything).Return(nil, merr.ErrImportFailed).Once()
testEngine := initHTTPServerV2(mp, true)
req := httptest.NewRequest(http.MethodPost, versionalV2(ImportJobCategory, AbortAction), bytes.NewReader([]byte(`{"jobId": "123"}`)))
req.SetBasicAuth(util.UserRoot, getDefaultRootPassword())
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, merr.Code(merr.ErrImportFailed), returnBody.Code)
mp.AssertNotCalled(t, "AbortImport", mock.Anything, mock.Anything)
})
t.Run("authorization enabled allows root", func(t *testing.T) {
paramtable.Get().Save(proxy.Params.CommonCfg.AuthorizationEnabled.Key, "true")
defer paramtable.Get().Reset(proxy.Params.CommonCfg.AuthorizationEnabled.Key)
mp := mocks.NewMockProxy(t)
mp.EXPECT().GetImportProgress(mock.Anything, mock.Anything).Return(&internalpb.GetImportProgressResponse{
Status: &StatusSuccess,
CollectionName: DefaultCollectionName,
}, nil).Once()
mp.EXPECT().AbortImport(mock.Anything, mock.Anything).Return(commonSuccessStatus, nil).Once()
testEngine := initHTTPServerV2(mp, true)
req := httptest.NewRequest(http.MethodPost, versionalV2(ImportJobCategory, AbortAction), bytes.NewReader([]byte(`{"jobId": "123"}`)))
req.SetBasicAuth(util.UserRoot, getDefaultRootPassword())
w := httptest.NewRecorder()
testEngine.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
returnBody := &ReturnErrMsg{}
err := json.Unmarshal(w.Body.Bytes(), returnBody)
assert.Nil(t, err)
assert.Equal(t, int32(0), returnBody.Code)
})
}
Reference in New Issue View Git Blame Copy Permalink