admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-06 17:18:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use suite for integration test (#24253)

Browse Source 
Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2023-05-19 15:57:24 +08:00 committed by GitHub
parent 4500d63248
commit 3a66e1de65
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
16 changed files with 1027 additions and 972 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
pkg/config/etcd_source.go
View File

@ -19,6 +19,7 @@ package config
import ( import (
"context" "context"
"fmt" "fmt"
"path"
"strings" "strings"
"sync" "sync"
"time" "time"
@ -104,7 +105,7 @@ func (es *EtcdSource) GetSourceName() string {
} }
func (es *EtcdSource) Close() { func (es *EtcdSource) Close() {
es.etcdCli.Close() // cannot close client here, since client is shared with components
es.configRefresher.stop() es.configRefresher.stop()
} }
@ -128,7 +129,7 @@ func (es *EtcdSource) UpdateOptions(opts Options) {
func (es *EtcdSource) refreshConfigurations() error { func (es *EtcdSource) refreshConfigurations() error {
es.RLock() es.RLock()
prefix := es.keyPrefix + "/config" prefix := path.Join(es.keyPrefix, "config")
es.RUnlock() es.RUnlock()
ctx, cancel := context.WithTimeout(es.ctx, ReadConfigTimeout) ctx, cancel := context.WithTimeout(es.ctx, ReadConfigTimeout)

2
pkg/config/manager.go
View File

@ -162,6 +162,8 @@ func (m *Manager) FileConfigs() map[string]string {
} }
func (m *Manager) Close() { func (m *Manager) Close() {
m.Lock()
defer m.Unlock()
for _, s := range m.sources { for _, s := range m.sources {
s.Close() s.Close()
} }

2
pkg/util/etcd/etcd_util.go
View File

@ -175,7 +175,7 @@ func buildKvGroup(keys, values []string) (map[string]string, error) {
// StartTestEmbedEtcdServer returns a newly created embed etcd server. // StartTestEmbedEtcdServer returns a newly created embed etcd server.
// ### USED FOR UNIT TEST ONLY ### // ### USED FOR UNIT TEST ONLY ###
func StartTestEmbedEtcdServer() (*embed.Etcd, string, error) { func StartTestEmbedEtcdServer() (*embed.Etcd, string, error) {
dir, err := ioutil.TempDir(os.TempDir(), "milvus_datanode_ut") dir, err := ioutil.TempDir(os.TempDir(), "milvus_ut")
if err != nil { if err != nil {
return nil, "", err return nil, "", err
} }

2
scripts/run_intergration_test.sh
View File

@ -17,7 +17,7 @@
# limitations under the License. # limitations under the License.
# run integration test # run integration test
echo "Running integration test uner ./tests/integration" echo "Running integration test under ./tests/integration"
BASEDIR=$(dirname "$0") BASEDIR=$(dirname "$0")
source $BASEDIR/setenv.sh source $BASEDIR/setenv.sh

66
tests/integration/bulkinsert_test.go
View File

@ -27,6 +27,7 @@ import (
"github.com/golang/protobuf/proto" "github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
"go.uber.org/zap" "go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/commonpb" "github.com/milvus-io/milvus-proto/go-api/commonpb"
@ -44,6 +45,10 @@ const (
Dim = 128 Dim = 128
) )
type BulkInsertSuite struct {
MiniClusterSuite
}
// test bulk insert E2E // test bulk insert E2E
// 1, create collection with a vector column and a varchar column // 1, create collection with a vector column and a varchar column
// 2, generate numpy files // 2, generate numpy files
@ -51,17 +56,10 @@ const (
// 4, create index // 4, create index
// 5, load // 5, load
// 6, search // 6, search
func TestBulkInsert(t *testing.T) { func (s *BulkInsertSuite) TestBulkInsert() {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*180) c := s.Cluster
c, err := StartMiniCluster(ctx) ctx, cancel := context.WithCancel(c.GetContext())
assert.NoError(t, err) defer cancel()
err = c.Start()
assert.NoError(t, err)
defer func() {
err = c.Stop()
assert.NoError(t, err)
cancel()
}()
prefix := "TestBulkInsert" prefix := "TestBulkInsert"
dbName := "" dbName := ""
@ -75,7 +73,7 @@ func TestBulkInsert(t *testing.T) {
&schemapb.FieldSchema{Name: "embeddings", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "128"}}}, &schemapb.FieldSchema{Name: "embeddings", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "128"}}},
) )
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err) s.NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName, DbName: dbName,
@ -83,17 +81,17 @@ func TestBulkInsert(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: common.DefaultShardsNum, ShardsNum: common.DefaultShardsNum,
}) })
assert.NoError(t, err) s.NoError(err)
if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason())) log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason()))
t.FailNow() s.FailNow("failed to create collection")
} }
assert.Equal(t, createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus)) log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus))
showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{}) showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{})
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, showCollectionsResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(showCollectionsResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp)) log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp))
err = GenerateNumpyFile(c.chunkManager.RootPath()+"/"+"embeddings.npy", 100, schemapb.DataType_FloatVector, []*commonpb.KeyValuePair{ err = GenerateNumpyFile(c.chunkManager.RootPath()+"/"+"embeddings.npy", 100, schemapb.DataType_FloatVector, []*commonpb.KeyValuePair{
@ -102,14 +100,14 @@ func TestBulkInsert(t *testing.T) {
Value: strconv.Itoa(Dim), Value: strconv.Itoa(Dim),
}, },
}) })
assert.NoError(t, err) s.NoError(err)
err = GenerateNumpyFile(c.chunkManager.RootPath()+"/"+"image_path.npy", 100, schemapb.DataType_VarChar, []*commonpb.KeyValuePair{ err = GenerateNumpyFile(c.chunkManager.RootPath()+"/"+"image_path.npy", 100, schemapb.DataType_VarChar, []*commonpb.KeyValuePair{
{ {
Key: common.MaxLengthKey, Key: common.MaxLengthKey,
Value: strconv.Itoa(65535), Value: strconv.Itoa(65535),
}, },
}) })
assert.NoError(t, err) s.NoError(err)
bulkInsertFiles := []string{ bulkInsertFiles := []string{
c.chunkManager.RootPath() + "/" + "embeddings.npy", c.chunkManager.RootPath() + "/" + "embeddings.npy",
@ -117,13 +115,13 @@ func TestBulkInsert(t *testing.T) {
} }
health1, err := c.dataCoord.CheckHealth(ctx, &milvuspb.CheckHealthRequest{}) health1, err := c.dataCoord.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
assert.NoError(t, err) s.NoError(err)
log.Info("dataCoord health", zap.Any("health1", health1)) log.Info("dataCoord health", zap.Any("health1", health1))
importResp, err := c.proxy.Import(ctx, &milvuspb.ImportRequest{ importResp, err := c.proxy.Import(ctx, &milvuspb.ImportRequest{
CollectionName: collectionName, CollectionName: collectionName,
Files: bulkInsertFiles, Files: bulkInsertFiles,
}) })
assert.NoError(t, err) s.NoError(err)
log.Info("Import result", zap.Any("importResp", importResp), zap.Int64s("tasks", importResp.GetTasks())) log.Info("Import result", zap.Any("importResp", importResp), zap.Int64s("tasks", importResp.GetTasks()))
tasks := importResp.GetTasks() tasks := importResp.GetTasks()
@ -133,7 +131,7 @@ func TestBulkInsert(t *testing.T) {
importTaskState, err := c.proxy.GetImportState(ctx, &milvuspb.GetImportStateRequest{ importTaskState, err := c.proxy.GetImportState(ctx, &milvuspb.GetImportStateRequest{
Task: task, Task: task,
}) })
assert.NoError(t, err) s.NoError(err)
switch importTaskState.GetState() { switch importTaskState.GetState() {
case commonpb.ImportState_ImportCompleted: case commonpb.ImportState_ImportCompleted:
break loop break loop
@ -150,12 +148,12 @@ func TestBulkInsert(t *testing.T) {
} }
health2, err := c.dataCoord.CheckHealth(ctx, &milvuspb.CheckHealthRequest{}) health2, err := c.dataCoord.CheckHealth(ctx, &milvuspb.CheckHealthRequest{})
assert.NoError(t, err) s.NoError(err)
log.Info("dataCoord health", zap.Any("health2", health2)) log.Info("dataCoord health", zap.Any("health2", health2))
segments, err := c.metaWatcher.ShowSegments() segments, err := c.metaWatcher.ShowSegments()
assert.NoError(t, err) s.NoError(err)
assert.NotEmpty(t, segments) s.NotEmpty(segments)
for _, segment := range segments { for _, segment := range segments {
log.Info("ShowSegments result", zap.String("segment", segment.String())) log.Info("ShowSegments result", zap.String("segment", segment.String()))
} }
@ -170,21 +168,21 @@ func TestBulkInsert(t *testing.T) {
if createIndexStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createIndexStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createIndexStatus fail reason", zap.String("reason", createIndexStatus.GetReason())) log.Warn("createIndexStatus fail reason", zap.String("reason", createIndexStatus.GetReason()))
} }
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode())
waitingForIndexBuilt(ctx, c, t, collectionName, "embeddings") waitingForIndexBuilt(ctx, c, s.T(), collectionName, "embeddings")
// load // load
loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{ loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
DbName: dbName, DbName: dbName,
CollectionName: collectionName, CollectionName: collectionName,
}) })
assert.NoError(t, err) s.NoError(err)
if loadStatus.GetErrorCode() != commonpb.ErrorCode_Success { if loadStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("loadStatus fail reason", zap.String("reason", loadStatus.GetReason())) log.Warn("loadStatus fail reason", zap.String("reason", loadStatus.GetReason()))
} }
assert.Equal(t, commonpb.ErrorCode_Success, loadStatus.GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, loadStatus.GetErrorCode())
waitingForLoad(ctx, c, collectionName) waitingForLoad(ctx, c, collectionName)
// search // search
@ -202,8 +200,8 @@ func TestBulkInsert(t *testing.T) {
if searchResult.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success { if searchResult.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("searchResult fail reason", zap.String("reason", searchResult.GetStatus().GetReason())) log.Warn("searchResult fail reason", zap.String("reason", searchResult.GetStatus().GetReason()))
} }
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, commonpb.ErrorCode_Success, searchResult.GetStatus().GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, searchResult.GetStatus().GetErrorCode())
log.Info("======================") log.Info("======================")
log.Info("======================") log.Info("======================")
@ -212,6 +210,10 @@ func TestBulkInsert(t *testing.T) {
log.Info("======================") log.Info("======================")
} }
func TestBulkInsert(t *testing.T) {
suite.Run(t, new(BulkInsertSuite))
}
func GenerateNumpyFile(filePath string, rowCount int, dType schemapb.DataType, typeParams []*commonpb.KeyValuePair) error { func GenerateNumpyFile(filePath string, rowCount int, dType schemapb.DataType, typeParams []*commonpb.KeyValuePair) error {
if dType == schemapb.DataType_VarChar { if dType == schemapb.DataType_VarChar {
var data []string var data []string

72
tests/integration/get_index_statistics_test.go
View File

@ -2,10 +2,8 @@ package integration
import ( import (
"context" "context"
"testing"
"github.com/golang/protobuf/proto" "github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert"
"go.uber.org/zap" "go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/commonpb" "github.com/milvus-io/milvus-proto/go-api/commonpb"
@ -16,14 +14,14 @@ import (
"github.com/milvus-io/milvus/pkg/util/funcutil" "github.com/milvus-io/milvus/pkg/util/funcutil"
) )
func TestGetIndexStatistics(t *testing.T) { type GetIndexStatisticsSuite struct {
ctx := context.Background() MiniClusterSuite
c, err := StartMiniCluster(ctx) }
assert.NoError(t, err)
err = c.Start() func (s *GetIndexStatisticsSuite) TestGetIndexStatistics() {
assert.NoError(t, err) c := s.Cluster
defer c.Stop() ctx, cancel := context.WithCancel(c.GetContext())
assert.NoError(t, err) defer cancel()
prefix := "TestGetIndexStatistics" prefix := "TestGetIndexStatistics"
dbName := "" dbName := ""
@ -33,7 +31,7 @@ func TestGetIndexStatistics(t *testing.T) {
schema := constructSchema(collectionName, dim, true) schema := constructSchema(collectionName, dim, true)
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err) s.NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName, DbName: dbName,
@ -41,11 +39,11 @@ func TestGetIndexStatistics(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: 2, ShardsNum: 2,
}) })
assert.NoError(t, err) s.NoError(err)
if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason())) log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason()))
} }
assert.Equal(t, createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim) fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
hashKeys := generateHashKeys(rowNum) hashKeys := generateHashKeys(rowNum)
@ -56,19 +54,19 @@ func TestGetIndexStatistics(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, insertResult.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(insertResult.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
// flush // flush
flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{ flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{
DbName: dbName, DbName: dbName,
CollectionNames: []string{collectionName}, CollectionNames: []string{collectionName},
}) })
assert.NoError(t, err) s.NoError(err)
segmentIDs, has := flushResp.GetCollSegIDs()[collectionName] segmentIDs, has := flushResp.GetCollSegIDs()[collectionName]
ids := segmentIDs.GetData() ids := segmentIDs.GetData()
assert.NotEmpty(t, segmentIDs) s.NotEmpty(segmentIDs)
assert.Equal(t, true, has) s.Equal(true, has)
waitingForFlush(ctx, c, ids) waitingForFlush(ctx, c, ids)
// create index // create index
@ -82,20 +80,20 @@ func TestGetIndexStatistics(t *testing.T) {
if createIndexStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createIndexStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createIndexStatus fail reason", zap.String("reason", createIndexStatus.GetReason())) log.Warn("createIndexStatus fail reason", zap.String("reason", createIndexStatus.GetReason()))
} }
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode())
waitingForIndexBuilt(ctx, c, t, collectionName, floatVecField) waitingForIndexBuilt(ctx, c, s.T(), collectionName, floatVecField)
getIndexStatisticsResponse, err := c.proxy.GetIndexStatistics(ctx, &milvuspb.GetIndexStatisticsRequest{ getIndexStatisticsResponse, err := c.proxy.GetIndexStatistics(ctx, &milvuspb.GetIndexStatisticsRequest{
CollectionName: collectionName, CollectionName: collectionName,
IndexName: indexName, IndexName: indexName,
}) })
assert.NoError(t, err) s.NoError(err)
indexInfos := getIndexStatisticsResponse.GetIndexDescriptions() indexInfos := getIndexStatisticsResponse.GetIndexDescriptions()
assert.Equal(t, 1, len(indexInfos)) s.Equal(1, len(indexInfos))
assert.Equal(t, int64(3000), indexInfos[0].IndexedRows) s.Equal(int64(3000), indexInfos[0].IndexedRows)
assert.Equal(t, int64(3000), indexInfos[0].TotalRows) s.Equal(int64(3000), indexInfos[0].TotalRows)
// skip second insert case for now // skip second insert case for now
// the result is not certain // the result is not certain
@ -107,32 +105,32 @@ func TestGetIndexStatistics(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, insertResult2.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(insertResult2.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
_, err = c.proxy.Flush(ctx, &milvuspb.FlushRequest{ _, err = c.proxy.Flush(ctx, &milvuspb.FlushRequest{
DbName: dbName, DbName: dbName,
CollectionNames: []string{collectionName}, CollectionNames: []string{collectionName},
}) })
assert.NoError(t, err) s.NoError(err)
segmentIDs2, has2 := flushResp.GetCollSegIDs()[collectionName] segmentIDs2, has2 := flushResp.GetCollSegIDs()[collectionName]
ids2 := segmentIDs2.GetData() ids2 := segmentIDs2.GetData()
assert.NotEmpty(t, segmentIDs) s.NotEmpty(segmentIDs)
assert.Equal(t, true, has2) s.Equal(true, has2)
waitingForFlush(ctx, c, ids2) waitingForFlush(ctx, c, ids2)
loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{ loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
DbName: dbName, DbName: dbName,
CollectionName: collectionName, CollectionName: collectionName,
}) })
assert.NoError(t, err) s.NoError(err)
if loadStatus.GetErrorCode() != commonpb.ErrorCode_Success { if loadStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("loadStatus fail reason", zap.String("reason", loadStatus.GetReason())) log.Warn("loadStatus fail reason", zap.String("reason", loadStatus.GetReason()))
} }
assert.Equal(t, commonpb.ErrorCode_Success, loadStatus.GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, loadStatus.GetErrorCode())
waitingForLoad(ctx, c, collectionName) waitingForLoad(ctx, c, collectionName)
assert.NoError(t, err) s.NoError(err)
waitingForIndexBuilt(ctx, c, t, collectionName, floatVecField) waitingForIndexBuilt(ctx, c, t, collectionName, floatVecField)
@ -140,11 +138,11 @@ func TestGetIndexStatistics(t *testing.T) {
CollectionName: collectionName, CollectionName: collectionName,
IndexName: indexName, IndexName: indexName,
}) })
assert.NoError(t, err) s.NoError(err)
indexInfos2 := getIndexStatisticsResponse2.GetIndexDescriptions() indexInfos2 := getIndexStatisticsResponse2.GetIndexDescriptions()
assert.Equal(t, 1, len(indexInfos2)) s.Equal(1, len(indexInfos2))
assert.Equal(t, int64(6000), indexInfos2[0].IndexedRows) s.Equal(int64(6000), indexInfos2[0].IndexedRows)
assert.Equal(t, int64(6000), indexInfos2[0].TotalRows) s.Equal(int64(6000), indexInfos2[0].TotalRows)
*/ */
log.Info("TestGetIndexStatistics succeed") log.Info("TestGetIndexStatistics succeed")

302
tests/integration/get_vector_test.go
View File

@ -21,7 +21,6 @@ import (
"fmt" "fmt"
"strconv" "strconv"
"testing" "testing"
"time"
"github.com/golang/protobuf/proto" "github.com/golang/protobuf/proto"
"github.com/stretchr/testify/suite" "github.com/stretchr/testify/suite"
@ -36,11 +35,7 @@ import (
) )
type TestGetVectorSuite struct { type TestGetVectorSuite struct {
suite.Suite MiniClusterSuite
ctx context.Context
cancel context.CancelFunc
cluster *MiniCluster
// test params // test params
nq int nq int
@ -51,25 +46,12 @@ type TestGetVectorSuite struct {
vecType schemapb.DataType vecType schemapb.DataType
} }
func (suite *TestGetVectorSuite) SetupTest() { func (s *TestGetVectorSuite) run() {
suite.ctx, suite.cancel = context.WithTimeout(context.Background(), time.Second*180) ctx, cancel := context.WithCancel(s.Cluster.ctx)
defer cancel()
var err error
suite.cluster, err = StartMiniCluster(suite.ctx)
suite.Require().NoError(err)
err = suite.cluster.Start()
suite.Require().NoError(err)
}
func (suite *TestGetVectorSuite) TearDownTest() {
err := suite.cluster.Stop()
suite.Require().NoError(err)
suite.cancel()
}
func (suite *TestGetVectorSuite) run() {
collection := fmt.Sprintf("TestGetVector_%d_%d_%s_%s_%s", collection := fmt.Sprintf("TestGetVector_%d_%d_%s_%s_%s",
suite.nq, suite.topK, suite.indexType, suite.metricType, funcutil.GenRandomStr()) s.nq, s.topK, s.indexType, s.metricType, funcutil.GenRandomStr())
const ( const (
NB = 10000 NB = 10000
@ -83,7 +65,7 @@ func (suite *TestGetVectorSuite) run() {
Name: pkFieldName, Name: pkFieldName,
IsPrimaryKey: true, IsPrimaryKey: true,
Description: "", Description: "",
DataType: suite.pkType, DataType: s.pkType,
TypeParams: []*commonpb.KeyValuePair{ TypeParams: []*commonpb.KeyValuePair{
{ {
Key: common.MaxLengthKey, Key: common.MaxLengthKey,
@ -98,7 +80,7 @@ func (suite *TestGetVectorSuite) run() {
Name: vecFieldName, Name: vecFieldName,
IsPrimaryKey: false, IsPrimaryKey: false,
Description: "", Description: "",
DataType: suite.vecType, DataType: s.vecType,
TypeParams: []*commonpb.KeyValuePair{ TypeParams: []*commonpb.KeyValuePair{
{ {
Key: common.DimKey, Key: common.DimKey,
@ -109,96 +91,96 @@ func (suite *TestGetVectorSuite) run() {
} }
schema := constructSchema(collection, dim, false, pk, fVec) schema := constructSchema(collection, dim, false, pk, fVec)
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
suite.Require().NoError(err) s.Require().NoError(err)
createCollectionStatus, err := suite.cluster.proxy.CreateCollection(suite.ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := s.Cluster.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
CollectionName: collection, CollectionName: collection,
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: 2, ShardsNum: 2,
}) })
suite.Require().NoError(err) s.Require().NoError(err)
suite.Require().Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Require().Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
fieldsData := make([]*schemapb.FieldData, 0) fieldsData := make([]*schemapb.FieldData, 0)
if suite.pkType == schemapb.DataType_Int64 { if s.pkType == schemapb.DataType_Int64 {
fieldsData = append(fieldsData, newInt64FieldData(pkFieldName, NB)) fieldsData = append(fieldsData, newInt64FieldData(pkFieldName, NB))
} else { } else {
fieldsData = append(fieldsData, newStringFieldData(pkFieldName, NB)) fieldsData = append(fieldsData, newStringFieldData(pkFieldName, NB))
} }
var vecFieldData *schemapb.FieldData var vecFieldData *schemapb.FieldData
if suite.vecType == schemapb.DataType_FloatVector { if s.vecType == schemapb.DataType_FloatVector {
vecFieldData = newFloatVectorFieldData(vecFieldName, NB, dim) vecFieldData = newFloatVectorFieldData(vecFieldName, NB, dim)
} else { } else {
vecFieldData = newBinaryVectorFieldData(vecFieldName, NB, dim) vecFieldData = newBinaryVectorFieldData(vecFieldName, NB, dim)
} }
fieldsData = append(fieldsData, vecFieldData) fieldsData = append(fieldsData, vecFieldData)
hashKeys := generateHashKeys(NB) hashKeys := generateHashKeys(NB)
_, err = suite.cluster.proxy.Insert(suite.ctx, &milvuspb.InsertRequest{ _, err = s.Cluster.proxy.Insert(ctx, &milvuspb.InsertRequest{
CollectionName: collection, CollectionName: collection,
FieldsData: fieldsData, FieldsData: fieldsData,
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(NB), NumRows: uint32(NB),
}) })
suite.Require().NoError(err) s.Require().NoError(err)
suite.Require().Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Require().Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
// flush // flush
flushResp, err := suite.cluster.proxy.Flush(suite.ctx, &milvuspb.FlushRequest{ flushResp, err := s.Cluster.proxy.Flush(ctx, &milvuspb.FlushRequest{
CollectionNames: []string{collection}, CollectionNames: []string{collection},
}) })
suite.Require().NoError(err) s.Require().NoError(err)
segmentIDs, has := flushResp.GetCollSegIDs()[collection] segmentIDs, has := flushResp.GetCollSegIDs()[collection]
ids := segmentIDs.GetData() ids := segmentIDs.GetData()
suite.Require().NotEmpty(segmentIDs) s.Require().NotEmpty(segmentIDs)
suite.Require().True(has) s.Require().True(has)
segments, err := suite.cluster.metaWatcher.ShowSegments() segments, err := s.Cluster.metaWatcher.ShowSegments()
suite.Require().NoError(err) s.Require().NoError(err)
suite.Require().NotEmpty(segments) s.Require().NotEmpty(segments)
waitingForFlush(suite.ctx, suite.cluster, ids) waitingForFlush(ctx, s.Cluster, ids)
// create index // create index
_, err = suite.cluster.proxy.CreateIndex(suite.ctx, &milvuspb.CreateIndexRequest{ _, err = s.Cluster.proxy.CreateIndex(ctx, &milvuspb.CreateIndexRequest{
CollectionName: collection, CollectionName: collection,
FieldName: vecFieldName, FieldName: vecFieldName,
IndexName: "_default", IndexName: "_default",
ExtraParams: constructIndexParam(dim, suite.indexType, suite.metricType), ExtraParams: constructIndexParam(dim, s.indexType, s.metricType),
}) })
suite.Require().NoError(err) s.Require().NoError(err)
suite.Require().Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Require().Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
waitingForIndexBuilt(suite.ctx, suite.cluster, suite.T(), collection, vecFieldName) waitingForIndexBuilt(ctx, s.Cluster, s.T(), collection, vecFieldName)
// load // load
_, err = suite.cluster.proxy.LoadCollection(suite.ctx, &milvuspb.LoadCollectionRequest{ _, err = s.Cluster.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
CollectionName: collection, CollectionName: collection,
}) })
suite.Require().NoError(err) s.Require().NoError(err)
suite.Require().Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Require().Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
waitingForLoad(suite.ctx, suite.cluster, collection) waitingForLoad(ctx, s.Cluster, collection)
// search // search
nq := suite.nq nq := s.nq
topk := suite.topK topk := s.topK
outputFields := []string{vecFieldName} outputFields := []string{vecFieldName}
params := getSearchParams(suite.indexType, suite.metricType) params := getSearchParams(s.indexType, s.metricType)
searchReq := constructSearchRequest("", collection, "", searchReq := constructSearchRequest("", collection, "",
vecFieldName, suite.vecType, outputFields, suite.metricType, params, nq, dim, topk, -1) vecFieldName, s.vecType, outputFields, s.metricType, params, nq, dim, topk, -1)
searchResp, err := suite.cluster.proxy.Search(suite.ctx, searchReq) searchResp, err := s.Cluster.proxy.Search(ctx, searchReq)
suite.Require().NoError(err) s.Require().NoError(err)
suite.Require().Equal(searchResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Require().Equal(searchResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
result := searchResp.GetResults() result := searchResp.GetResults()
if suite.pkType == schemapb.DataType_Int64 { if s.pkType == schemapb.DataType_Int64 {
suite.Require().Len(result.GetIds().GetIntId().GetData(), nq*topk) s.Require().Len(result.GetIds().GetIntId().GetData(), nq*topk)
} else { } else {
suite.Require().Len(result.GetIds().GetStrId().GetData(), nq*topk) s.Require().Len(result.GetIds().GetStrId().GetData(), nq*topk)
} }
suite.Require().Len(result.GetScores(), nq*topk) s.Require().Len(result.GetScores(), nq*topk)
suite.Require().GreaterOrEqual(len(result.GetFieldsData()), 1) s.Require().GreaterOrEqual(len(result.GetFieldsData()), 1)
var vecFieldIndex = -1 var vecFieldIndex = -1
for i, fieldData := range result.GetFieldsData() { for i, fieldData := range result.GetFieldsData() {
if typeutil.IsVectorType(fieldData.GetType()) { if typeutil.IsVectorType(fieldData.GetType()) {
@ -206,162 +188,162 @@ func (suite *TestGetVectorSuite) run() {
break break
} }
} }
suite.Require().EqualValues(nq, result.GetNumQueries()) s.Require().EqualValues(nq, result.GetNumQueries())
suite.Require().EqualValues(topk, result.GetTopK()) s.Require().EqualValues(topk, result.GetTopK())
// check output vectors // check output vectors
if suite.vecType == schemapb.DataType_FloatVector { if s.vecType == schemapb.DataType_FloatVector {
suite.Require().Len(result.GetFieldsData()[vecFieldIndex].GetVectors().GetFloatVector().GetData(), nq*topk*dim) s.Require().Len(result.GetFieldsData()[vecFieldIndex].GetVectors().GetFloatVector().GetData(), nq*topk*dim)
rawData := vecFieldData.GetVectors().GetFloatVector().GetData() rawData := vecFieldData.GetVectors().GetFloatVector().GetData()
resData := result.GetFieldsData()[vecFieldIndex].GetVectors().GetFloatVector().GetData() resData := result.GetFieldsData()[vecFieldIndex].GetVectors().GetFloatVector().GetData()
if suite.pkType == schemapb.DataType_Int64 { if s.pkType == schemapb.DataType_Int64 {
for i, id := range result.GetIds().GetIntId().GetData() { for i, id := range result.GetIds().GetIntId().GetData() {
expect := rawData[int(id)*dim : (int(id)+1)*dim] expect := rawData[int(id)*dim : (int(id)+1)*dim]
actual := resData[i*dim : (i+1)*dim] actual := resData[i*dim : (i+1)*dim]
suite.Require().ElementsMatch(expect, actual) s.Require().ElementsMatch(expect, actual)
} }
} else { } else {
for i, idStr := range result.GetIds().GetStrId().GetData() { for i, idStr := range result.GetIds().GetStrId().GetData() {
id, err := strconv.Atoi(idStr) id, err := strconv.Atoi(idStr)
suite.Require().NoError(err) s.Require().NoError(err)
expect := rawData[id*dim : (id+1)*dim] expect := rawData[id*dim : (id+1)*dim]
actual := resData[i*dim : (i+1)*dim] actual := resData[i*dim : (i+1)*dim]
suite.Require().ElementsMatch(expect, actual) s.Require().ElementsMatch(expect, actual)
} }
} }
} else { } else {
suite.Require().Len(result.GetFieldsData()[vecFieldIndex].GetVectors().GetBinaryVector(), nq*topk*dim/8) s.Require().Len(result.GetFieldsData()[vecFieldIndex].GetVectors().GetBinaryVector(), nq*topk*dim/8)
rawData := vecFieldData.GetVectors().GetBinaryVector() rawData := vecFieldData.GetVectors().GetBinaryVector()
resData := result.GetFieldsData()[vecFieldIndex].GetVectors().GetBinaryVector() resData := result.GetFieldsData()[vecFieldIndex].GetVectors().GetBinaryVector()
if suite.pkType == schemapb.DataType_Int64 { if s.pkType == schemapb.DataType_Int64 {
for i, id := range result.GetIds().GetIntId().GetData() { for i, id := range result.GetIds().GetIntId().GetData() {
dataBytes := dim / 8 dataBytes := dim / 8
for j := 0; j < dataBytes; j++ { for j := 0; j < dataBytes; j++ {
expect := rawData[int(id)*dataBytes+j] expect := rawData[int(id)*dataBytes+j]
actual := resData[i*dataBytes+j] actual := resData[i*dataBytes+j]
suite.Require().Equal(expect, actual) s.Require().Equal(expect, actual)
} }
} }
} else { } else {
for i, idStr := range result.GetIds().GetStrId().GetData() { for i, idStr := range result.GetIds().GetStrId().GetData() {
dataBytes := dim / 8 dataBytes := dim / 8
id, err := strconv.Atoi(idStr) id, err := strconv.Atoi(idStr)
suite.Require().NoError(err) s.Require().NoError(err)
for j := 0; j < dataBytes; j++ { for j := 0; j < dataBytes; j++ {
expect := rawData[id*dataBytes+j] expect := rawData[id*dataBytes+j]
actual := resData[i*dataBytes+j] actual := resData[i*dataBytes+j]
suite.Require().Equal(expect, actual) s.Require().Equal(expect, actual)
} }
} }
} }
} }
status, err := suite.cluster.proxy.DropCollection(suite.ctx, &milvuspb.DropCollectionRequest{ status, err := s.Cluster.proxy.DropCollection(ctx, &milvuspb.DropCollectionRequest{
CollectionName: collection, CollectionName: collection,
}) })
suite.Require().NoError(err) s.Require().NoError(err)
suite.Require().Equal(status.GetErrorCode(), commonpb.ErrorCode_Success) s.Require().Equal(status.GetErrorCode(), commonpb.ErrorCode_Success)
} }
func (suite *TestGetVectorSuite) TestGetVector_FLAT() { func (s *TestGetVectorSuite) TestGetVector_FLAT() {
suite.nq = 10 s.nq = 10
suite.topK = 10 s.topK = 10
suite.indexType = IndexFaissIDMap s.indexType = IndexFaissIDMap
suite.metricType = distance.L2 s.metricType = distance.L2
suite.pkType = schemapb.DataType_Int64 s.pkType = schemapb.DataType_Int64
suite.vecType = schemapb.DataType_FloatVector s.vecType = schemapb.DataType_FloatVector
suite.run() s.run()
} }
func (suite *TestGetVectorSuite) TestGetVector_IVF_FLAT() { func (s *TestGetVectorSuite) TestGetVector_IVF_FLAT() {
suite.nq = 10 s.nq = 10
suite.topK = 10 s.topK = 10
suite.indexType = IndexFaissIvfFlat s.indexType = IndexFaissIvfFlat
suite.metricType = distance.L2 s.metricType = distance.L2
suite.pkType = schemapb.DataType_Int64 s.pkType = schemapb.DataType_Int64
suite.vecType = schemapb.DataType_FloatVector s.vecType = schemapb.DataType_FloatVector
suite.run() s.run()
} }
func (suite *TestGetVectorSuite) TestGetVector_IVF_PQ() { func (s *TestGetVectorSuite) TestGetVector_IVF_PQ() {
suite.nq = 10 s.nq = 10
suite.topK = 10 s.topK = 10
suite.indexType = IndexFaissIvfPQ s.indexType = IndexFaissIvfPQ
suite.metricType = distance.L2 s.metricType = distance.L2
suite.pkType = schemapb.DataType_Int64 s.pkType = schemapb.DataType_Int64
suite.vecType = schemapb.DataType_FloatVector s.vecType = schemapb.DataType_FloatVector
suite.run() s.run()
} }
func (suite *TestGetVectorSuite) TestGetVector_IVF_SQ8() { func (s *TestGetVectorSuite) TestGetVector_IVF_SQ8() {
suite.nq = 10 s.nq = 10
suite.topK = 10 s.topK = 10
suite.indexType = IndexFaissIvfSQ8 s.indexType = IndexFaissIvfSQ8
suite.metricType = distance.L2 s.metricType = distance.L2
suite.pkType = schemapb.DataType_Int64 s.pkType = schemapb.DataType_Int64
suite.vecType = schemapb.DataType_FloatVector s.vecType = schemapb.DataType_FloatVector
suite.run() s.run()
} }
func (suite *TestGetVectorSuite) TestGetVector_HNSW() { func (s *TestGetVectorSuite) TestGetVector_HNSW() {
suite.nq = 10 s.nq = 10
suite.topK = 10 s.topK = 10
suite.indexType = IndexHNSW s.indexType = IndexHNSW
suite.metricType = distance.L2 s.metricType = distance.L2
suite.pkType = schemapb.DataType_Int64 s.pkType = schemapb.DataType_Int64
suite.vecType = schemapb.DataType_FloatVector s.vecType = schemapb.DataType_FloatVector
suite.run() s.run()
} }
func (suite *TestGetVectorSuite) TestGetVector_IP() { func (s *TestGetVectorSuite) TestGetVector_IP() {
suite.nq = 10 s.nq = 10
suite.topK = 10 s.topK = 10
suite.indexType = IndexHNSW s.indexType = IndexHNSW
suite.metricType = distance.IP s.metricType = distance.IP
suite.pkType = schemapb.DataType_Int64 s.pkType = schemapb.DataType_Int64
suite.vecType = schemapb.DataType_FloatVector s.vecType = schemapb.DataType_FloatVector
suite.run() s.run()
} }
func (suite *TestGetVectorSuite) TestGetVector_StringPK() { func (s *TestGetVectorSuite) TestGetVector_StringPK() {
suite.nq = 10 s.nq = 10
suite.topK = 10 s.topK = 10
suite.indexType = IndexHNSW s.indexType = IndexHNSW
suite.metricType = distance.L2 s.metricType = distance.L2
suite.pkType = schemapb.DataType_VarChar s.pkType = schemapb.DataType_VarChar
suite.vecType = schemapb.DataType_FloatVector s.vecType = schemapb.DataType_FloatVector
suite.run() s.run()
} }
func (suite *TestGetVectorSuite) TestGetVector_BinaryVector() { func (s *TestGetVectorSuite) TestGetVector_BinaryVector() {
suite.nq = 10 s.nq = 10
suite.topK = 10 s.topK = 10
suite.indexType = IndexFaissBinIvfFlat s.indexType = IndexFaissBinIvfFlat
suite.metricType = distance.JACCARD s.metricType = distance.JACCARD
suite.pkType = schemapb.DataType_Int64 s.pkType = schemapb.DataType_Int64
suite.vecType = schemapb.DataType_BinaryVector s.vecType = schemapb.DataType_BinaryVector
suite.run() s.run()
} }
func (suite *TestGetVectorSuite) TestGetVector_Big_NQ_TOPK() { func (s *TestGetVectorSuite) TestGetVector_Big_NQ_TOPK() {
suite.T().Skip("skip big NQ Top due to timeout") s.T().Skip("skip big NQ Top due to timeout")
suite.nq = 10000 s.nq = 10000
suite.topK = 200 s.topK = 200
suite.indexType = IndexHNSW s.indexType = IndexHNSW
suite.metricType = distance.L2 s.metricType = distance.L2
suite.pkType = schemapb.DataType_Int64 s.pkType = schemapb.DataType_Int64
suite.vecType = schemapb.DataType_FloatVector s.vecType = schemapb.DataType_FloatVector
suite.run() s.run()
} }
//func (suite *TestGetVectorSuite) TestGetVector_DISKANN() { //func (s *TestGetVectorSuite) TestGetVector_DISKANN() {
// suite.nq = 10 // s.nq = 10
// suite.topK = 10 // s.topK = 10
// suite.indexType = IndexDISKANN // s.indexType = IndexDISKANN
// suite.metricType = distance.L2 // s.metricType = distance.L2
// suite.pkType = schemapb.DataType_Int64 // s.pkType = schemapb.DataType_Int64
// suite.vecType = schemapb.DataType_FloatVector // s.vecType = schemapb.DataType_FloatVector
// suite.run() // s.run()
//} //}
func TestGetVector(t *testing.T) { func TestGetVector(t *testing.T) {

64
tests/integration/hello_milvus_test.go
View File

@ -20,10 +20,9 @@ import (
"context" "context"
"fmt" "fmt"
"testing" "testing"
"time"
"github.com/golang/protobuf/proto" "github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite"
"go.uber.org/zap" "go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/commonpb" "github.com/milvus-io/milvus-proto/go-api/commonpb"
@ -35,18 +34,14 @@ import (
"github.com/milvus-io/milvus/pkg/util/funcutil" "github.com/milvus-io/milvus/pkg/util/funcutil"
) )
func TestHelloMilvus(t *testing.T) { type HelloMilvusSuite struct {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*180) MiniClusterSuite
}
func (s *HelloMilvusSuite) TestHelloMilvus() {
ctx, cancel := context.WithCancel(s.Cluster.ctx)
defer cancel() defer cancel()
c, err := StartMiniCluster(ctx) c := s.Cluster
assert.NoError(t, err)
err = c.Start()
assert.NoError(t, err)
defer func() {
err = c.Stop()
assert.NoError(t, err)
cancel()
}()
const ( const (
dim = 128 dim = 128
@ -58,7 +53,7 @@ func TestHelloMilvus(t *testing.T) {
schema := constructSchema(collectionName, dim, true) schema := constructSchema(collectionName, dim, true)
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err) s.NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName, DbName: dbName,
@ -66,16 +61,16 @@ func TestHelloMilvus(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: common.DefaultShardsNum, ShardsNum: common.DefaultShardsNum,
}) })
assert.NoError(t, err) s.NoError(err)
if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason())) log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason()))
} }
assert.Equal(t, createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus)) log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus))
showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{}) showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{})
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, showCollectionsResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(showCollectionsResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp)) log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp))
fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim) fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
@ -87,23 +82,23 @@ func TestHelloMilvus(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, insertResult.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(insertResult.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
// flush // flush
flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{ flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{
DbName: dbName, DbName: dbName,
CollectionNames: []string{collectionName}, CollectionNames: []string{collectionName},
}) })
assert.NoError(t, err) s.NoError(err)
segmentIDs, has := flushResp.GetCollSegIDs()[collectionName] segmentIDs, has := flushResp.GetCollSegIDs()[collectionName]
ids := segmentIDs.GetData() ids := segmentIDs.GetData()
assert.NotEmpty(t, segmentIDs) s.NotEmpty(segmentIDs)
assert.True(t, has) s.True(has)
segments, err := c.metaWatcher.ShowSegments() segments, err := c.metaWatcher.ShowSegments()
assert.NoError(t, err) s.NoError(err)
assert.NotEmpty(t, segments) s.NotEmpty(segments)
for _, segment := range segments { for _, segment := range segments {
log.Info("ShowSegments result", zap.String("segment", segment.String())) log.Info("ShowSegments result", zap.String("segment", segment.String()))
} }
@ -119,21 +114,21 @@ func TestHelloMilvus(t *testing.T) {
if createIndexStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createIndexStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createIndexStatus fail reason", zap.String("reason", createIndexStatus.GetReason())) log.Warn("createIndexStatus fail reason", zap.String("reason", createIndexStatus.GetReason()))
} }
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode())
waitingForIndexBuilt(ctx, c, t, collectionName, floatVecField) waitingForIndexBuilt(ctx, c, s.T(), collectionName, floatVecField)
// load // load
loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{ loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
DbName: dbName, DbName: dbName,
CollectionName: collectionName, CollectionName: collectionName,
}) })
assert.NoError(t, err) s.NoError(err)
if loadStatus.GetErrorCode() != commonpb.ErrorCode_Success { if loadStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("loadStatus fail reason", zap.String("reason", loadStatus.GetReason())) log.Warn("loadStatus fail reason", zap.String("reason", loadStatus.GetReason()))
} }
assert.Equal(t, commonpb.ErrorCode_Success, loadStatus.GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, loadStatus.GetErrorCode())
waitingForLoad(ctx, c, collectionName) waitingForLoad(ctx, c, collectionName)
// search // search
@ -151,8 +146,13 @@ func TestHelloMilvus(t *testing.T) {
if searchResult.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success { if searchResult.GetStatus().GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("searchResult fail reason", zap.String("reason", searchResult.GetStatus().GetReason())) log.Warn("searchResult fail reason", zap.String("reason", searchResult.GetStatus().GetReason()))
} }
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, commonpb.ErrorCode_Success, searchResult.GetStatus().GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, searchResult.GetStatus().GetErrorCode())
log.Info("TestHelloMilvus succeed") log.Info("TestHelloMilvus succeed")
}
func TestHelloMilvus(t *testing.T) {
suite.Run(t, new(HelloMilvusSuite))
} }

847
tests/integration/json_expr_test.go
View File

File diff suppressed because it is too large Load Diff

103
tests/integration/meta_watcher_test.go
View File

@ -24,7 +24,7 @@ import (
"github.com/cockroachdb/errors" "github.com/cockroachdb/errors"
"github.com/golang/protobuf/proto" "github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite"
"go.uber.org/zap" "go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/commonpb" "github.com/milvus-io/milvus-proto/go-api/commonpb"
@ -36,32 +36,23 @@ import (
"github.com/milvus-io/milvus/pkg/util/funcutil" "github.com/milvus-io/milvus/pkg/util/funcutil"
) )
func TestShowSessions(t *testing.T) { type MetaWatcherSuite struct {
ctx := context.Background() MiniClusterSuite
c, err := StartMiniCluster(ctx) }
assert.NoError(t, err)
err = c.Start()
assert.NoError(t, err)
defer c.Stop()
assert.NoError(t, err)
sessions, err := c.metaWatcher.ShowSessions() func (s *MetaWatcherSuite) TestShowSessions() {
assert.NoError(t, err) sessions, err := s.Cluster.metaWatcher.ShowSessions()
assert.NotEmpty(t, sessions) s.NoError(err)
s.NotEmpty(sessions)
for _, session := range sessions { for _, session := range sessions {
log.Info("ShowSessions result", zap.String("session", session.String())) log.Info("ShowSessions result", zap.String("session", session.String()))
} }
} }
func TestShowSegments(t *testing.T) { func (s *MetaWatcherSuite) TestShowSegments() {
ctx := context.Background() c := s.Cluster
c, err := StartMiniCluster(ctx) ctx, cancel := context.WithCancel(c.GetContext())
assert.NoError(t, err) defer cancel()
err = c.Start()
assert.NoError(t, err)
defer c.Stop()
assert.NoError(t, err)
prefix := "TestShowSegments" prefix := "TestShowSegments"
dbName := "" dbName := ""
@ -109,7 +100,7 @@ func TestShowSegments(t *testing.T) {
} }
schema := constructCollectionSchema() schema := constructCollectionSchema()
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err) s.NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName, DbName: dbName,
@ -117,13 +108,13 @@ func TestShowSegments(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: common.DefaultShardsNum, ShardsNum: common.DefaultShardsNum,
}) })
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus)) log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus))
showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{}) showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{})
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, showCollectionsResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(showCollectionsResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp)) log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp))
fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim) fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
@ -135,25 +126,21 @@ func TestShowSegments(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, insertResult.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(insertResult.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
segments, err := c.metaWatcher.ShowSegments() segments, err := c.metaWatcher.ShowSegments()
assert.NoError(t, err) s.NoError(err)
assert.NotEmpty(t, segments) s.NotEmpty(segments)
for _, segment := range segments { for _, segment := range segments {
log.Info("ShowSegments result", zap.String("segment", segment.String())) log.Info("ShowSegments result", zap.String("segment", segment.String()))
} }
} }
func TestShowReplicas(t *testing.T) { func (s *MetaWatcherSuite) TestShowReplicas() {
ctx := context.Background() c := s.Cluster
c, err := StartMiniCluster(ctx) ctx, cancel := context.WithCancel(c.GetContext())
assert.NoError(t, err) defer cancel()
err = c.Start()
assert.NoError(t, err)
defer c.Stop()
assert.NoError(t, err)
prefix := "TestShowReplicas" prefix := "TestShowReplicas"
dbName := "" dbName := ""
@ -201,7 +188,7 @@ func TestShowReplicas(t *testing.T) {
} }
schema := constructCollectionSchema() schema := constructCollectionSchema()
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err) s.NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName, DbName: dbName,
@ -209,16 +196,16 @@ func TestShowReplicas(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: common.DefaultShardsNum, ShardsNum: common.DefaultShardsNum,
}) })
assert.NoError(t, err) s.NoError(err)
if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason())) log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason()))
} }
assert.Equal(t, createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus)) log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus))
showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{}) showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{})
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, showCollectionsResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(showCollectionsResp.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp)) log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp))
fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim) fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
@ -230,22 +217,22 @@ func TestShowReplicas(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, insertResult.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(insertResult.GetStatus().GetErrorCode(), commonpb.ErrorCode_Success)
// flush // flush
flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{ flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{
DbName: dbName, DbName: dbName,
CollectionNames: []string{collectionName}, CollectionNames: []string{collectionName},
}) })
assert.NoError(t, err) s.NoError(err)
segmentIDs, has := flushResp.GetCollSegIDs()[collectionName] segmentIDs, has := flushResp.GetCollSegIDs()[collectionName]
ids := segmentIDs.GetData() ids := segmentIDs.GetData()
assert.NotEmpty(t, segmentIDs) s.NotEmpty(segmentIDs)
segments, err := c.metaWatcher.ShowSegments() segments, err := c.metaWatcher.ShowSegments()
assert.NoError(t, err) s.NoError(err)
assert.NotEmpty(t, segments) s.NotEmpty(segments)
for _, segment := range segments { for _, segment := range segments {
log.Info("ShowSegments result", zap.String("segment", segment.String())) log.Info("ShowSegments result", zap.String("segment", segment.String()))
} }
@ -296,24 +283,24 @@ func TestShowReplicas(t *testing.T) {
}, },
}, },
}) })
assert.NoError(t, err) s.NoError(err)
if createIndexStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createIndexStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createIndexStatus fail reason", zap.String("reason", createIndexStatus.GetReason())) log.Warn("createIndexStatus fail reason", zap.String("reason", createIndexStatus.GetReason()))
} }
assert.Equal(t, commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, createIndexStatus.GetErrorCode())
waitingForIndexBuilt(ctx, c, t, collectionName, floatVecField) waitingForIndexBuilt(ctx, c, s.T(), collectionName, floatVecField)
// load // load
loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{ loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
DbName: dbName, DbName: dbName,
CollectionName: collectionName, CollectionName: collectionName,
}) })
assert.NoError(t, err) s.NoError(err)
if loadStatus.GetErrorCode() != commonpb.ErrorCode_Success { if loadStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("loadStatus fail reason", zap.String("reason", loadStatus.GetReason())) log.Warn("loadStatus fail reason", zap.String("reason", loadStatus.GetReason()))
} }
assert.Equal(t, commonpb.ErrorCode_Success, loadStatus.GetErrorCode()) s.Equal(commonpb.ErrorCode_Success, loadStatus.GetErrorCode())
for { for {
loadProgress, err := c.proxy.GetLoadingProgress(ctx, &milvuspb.GetLoadingProgressRequest{ loadProgress, err := c.proxy.GetLoadingProgress(ctx, &milvuspb.GetLoadingProgressRequest{
CollectionName: collectionName, CollectionName: collectionName,
@ -328,11 +315,15 @@ func TestShowReplicas(t *testing.T) {
} }
replicas, err := c.metaWatcher.ShowReplicas() replicas, err := c.metaWatcher.ShowReplicas()
assert.NoError(t, err) s.NoError(err)
assert.NotEmpty(t, replicas) s.NotEmpty(replicas)
for _, replica := range replicas { for _, replica := range replicas {
log.Info("ShowReplicas result", zap.String("replica", PrettyReplica(replica))) log.Info("ShowReplicas result", zap.String("replica", PrettyReplica(replica)))
} }
log.Info("TestShowReplicas succeed") log.Info("TestShowReplicas succeed")
} }
func TestMetaWatcher(t *testing.T) {
suite.Run(t, new(MetaWatcherSuite))
}

4
tests/integration/minicluster.go
View File

@ -310,6 +310,10 @@ func StartMiniCluster(ctx context.Context, opts ...Option) (cluster *MiniCluster
return cluster, nil return cluster, nil
} }
func (cluster *MiniCluster) GetContext() context.Context {
return cluster.ctx
}
func (cluster *MiniCluster) Start() error { func (cluster *MiniCluster) Start() error {
log.Info("mini cluster start") log.Info("mini cluster start")
err := cluster.rootCoord.Init() err := cluster.rootCoord.Init()

179
tests/integration/minicluster_test.go
View File

@ -20,183 +20,170 @@ import (
"context" "context"
"testing" "testing"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus/internal/datanode" "github.com/milvus-io/milvus/internal/datanode"
"github.com/milvus-io/milvus/internal/indexnode" "github.com/milvus-io/milvus/internal/indexnode"
"github.com/milvus-io/milvus/internal/querynodev2" "github.com/milvus-io/milvus/internal/querynodev2"
) )
func TestMiniClusterStartAndStop(t *testing.T) { type MiniClusterMethodsSuite struct {
ctx := context.Background() MiniClusterSuite
c, err := StartMiniCluster(ctx)
assert.NoError(t, err)
err = c.Start()
assert.NoError(t, err)
err = c.Stop()
assert.NoError(t, err)
} }
func TestAddRemoveDataNode(t *testing.T) { func (s *MiniClusterMethodsSuite) TestStartAndStop() {
ctx := context.Background() //Do nothing
c, err := StartMiniCluster(ctx) }
assert.NoError(t, err)
err = c.Start() func (s *MiniClusterMethodsSuite) TestRemoveDataNode() {
assert.NoError(t, err) c := s.Cluster
defer c.Stop() ctx, cancel := context.WithCancel(c.GetContext())
assert.NoError(t, err) defer cancel()
datanode := datanode.NewDataNode(ctx, c.factory) datanode := datanode.NewDataNode(ctx, c.factory)
datanode.SetEtcdClient(c.etcdCli) datanode.SetEtcdClient(c.etcdCli)
//datanode := c.CreateDefaultDataNode() //datanode := c.CreateDefaultDataNode()
err = c.AddDataNode(datanode) err := c.AddDataNode(datanode)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 2, c.clusterConfig.DataNodeNum) s.Equal(2, c.clusterConfig.DataNodeNum)
assert.Equal(t, 2, len(c.dataNodes)) s.Equal(2, len(c.dataNodes))
err = c.RemoveDataNode(datanode) err = c.RemoveDataNode(datanode)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 1, c.clusterConfig.DataNodeNum) s.Equal(1, c.clusterConfig.DataNodeNum)
assert.Equal(t, 1, len(c.dataNodes)) s.Equal(1, len(c.dataNodes))
// add default node and remove randomly // add default node and remove randomly
err = c.AddDataNode(nil) err = c.AddDataNode(nil)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 2, c.clusterConfig.DataNodeNum) s.Equal(2, c.clusterConfig.DataNodeNum)
assert.Equal(t, 2, len(c.dataNodes)) s.Equal(2, len(c.dataNodes))
err = c.RemoveDataNode(nil) err = c.RemoveDataNode(nil)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 1, c.clusterConfig.DataNodeNum) s.Equal(1, c.clusterConfig.DataNodeNum)
assert.Equal(t, 1, len(c.dataNodes)) s.Equal(1, len(c.dataNodes))
} }
func TestAddRemoveQueryNode(t *testing.T) { func (s *MiniClusterMethodsSuite) TestRemoveQueryNode() {
ctx := context.Background() c := s.Cluster
c, err := StartMiniCluster(ctx) ctx, cancel := context.WithCancel(c.GetContext())
assert.NoError(t, err) defer cancel()
err = c.Start()
assert.NoError(t, err)
defer c.Stop()
assert.NoError(t, err)
queryNode := querynodev2.NewQueryNode(ctx, c.factory) queryNode := querynodev2.NewQueryNode(ctx, c.factory)
queryNode.SetEtcdClient(c.etcdCli) queryNode.SetEtcdClient(c.etcdCli)
//queryNode := c.CreateDefaultQueryNode() //queryNode := c.CreateDefaultQueryNode()
err = c.AddQueryNode(queryNode) err := c.AddQueryNode(queryNode)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 2, c.clusterConfig.QueryNodeNum) s.Equal(2, c.clusterConfig.QueryNodeNum)
assert.Equal(t, 2, len(c.queryNodes)) s.Equal(2, len(c.queryNodes))
err = c.RemoveQueryNode(queryNode) err = c.RemoveQueryNode(queryNode)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 1, c.clusterConfig.QueryNodeNum) s.Equal(1, c.clusterConfig.QueryNodeNum)
assert.Equal(t, 1, len(c.queryNodes)) s.Equal(1, len(c.queryNodes))
// add default node and remove randomly // add default node and remove randomly
err = c.AddQueryNode(nil) err = c.AddQueryNode(nil)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 2, c.clusterConfig.QueryNodeNum) s.Equal(2, c.clusterConfig.QueryNodeNum)
assert.Equal(t, 2, len(c.queryNodes)) s.Equal(2, len(c.queryNodes))
err = c.RemoveQueryNode(nil) err = c.RemoveQueryNode(nil)
assert.NoError(t, err) s.NoError(err)
s.Equal(1, c.clusterConfig.QueryNodeNum)
s.Equal(1, len(c.queryNodes))
assert.Equal(t, 1, c.clusterConfig.QueryNodeNum)
assert.Equal(t, 1, len(c.queryNodes))
} }
func TestAddRemoveIndexNode(t *testing.T) { func (s *MiniClusterMethodsSuite) TestRemoveIndexNode() {
ctx := context.Background() c := s.Cluster
c, err := StartMiniCluster(ctx) ctx, cancel := context.WithCancel(c.GetContext())
assert.NoError(t, err) defer cancel()
err = c.Start()
assert.NoError(t, err)
defer c.Stop()
assert.NoError(t, err)
indexNode := indexnode.NewIndexNode(ctx, c.factory) indexNode := indexnode.NewIndexNode(ctx, c.factory)
indexNode.SetEtcdClient(c.etcdCli) indexNode.SetEtcdClient(c.etcdCli)
//indexNode := c.CreateDefaultIndexNode() //indexNode := c.CreateDefaultIndexNode()
err = c.AddIndexNode(indexNode) err := c.AddIndexNode(indexNode)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 2, c.clusterConfig.IndexNodeNum) s.Equal(2, c.clusterConfig.IndexNodeNum)
assert.Equal(t, 2, len(c.indexNodes)) s.Equal(2, len(c.indexNodes))
err = c.RemoveIndexNode(indexNode) err = c.RemoveIndexNode(indexNode)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 1, c.clusterConfig.IndexNodeNum) s.Equal(1, c.clusterConfig.IndexNodeNum)
assert.Equal(t, 1, len(c.indexNodes)) s.Equal(1, len(c.indexNodes))
// add default node and remove randomly // add default node and remove randomly
err = c.AddIndexNode(nil) err = c.AddIndexNode(nil)
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 2, c.clusterConfig.IndexNodeNum) s.Equal(2, c.clusterConfig.IndexNodeNum)
assert.Equal(t, 2, len(c.indexNodes)) s.Equal(2, len(c.indexNodes))
err = c.RemoveIndexNode(nil) err = c.RemoveIndexNode(nil)
assert.NoError(t, err) s.NoError(err)
s.Equal(1, c.clusterConfig.IndexNodeNum)
s.Equal(1, len(c.indexNodes))
assert.Equal(t, 1, c.clusterConfig.IndexNodeNum)
assert.Equal(t, 1, len(c.indexNodes))
} }
func TestUpdateClusterSize(t *testing.T) { func (s *MiniClusterMethodsSuite) TestUpdateClusterSize() {
ctx := context.Background()
c, err := StartMiniCluster(ctx)
assert.NoError(t, err)
err = c.Start()
assert.NoError(t, err)
defer c.Stop()
assert.NoError(t, err)
err = c.UpdateClusterSize(ClusterConfig{ c := s.Cluster
err := c.UpdateClusterSize(ClusterConfig{
QueryNodeNum: -1, QueryNodeNum: -1,
DataNodeNum: -1, DataNodeNum: -1,
IndexNodeNum: -1, IndexNodeNum: -1,
}) })
assert.Error(t, err) s.Error(err)
err = c.UpdateClusterSize(ClusterConfig{ err = c.UpdateClusterSize(ClusterConfig{
QueryNodeNum: 2, QueryNodeNum: 2,
DataNodeNum: 2, DataNodeNum: 2,
IndexNodeNum: 2, IndexNodeNum: 2,
}) })
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 2, c.clusterConfig.DataNodeNum) s.Equal(2, c.clusterConfig.DataNodeNum)
assert.Equal(t, 2, c.clusterConfig.QueryNodeNum) s.Equal(2, c.clusterConfig.QueryNodeNum)
assert.Equal(t, 2, c.clusterConfig.IndexNodeNum) s.Equal(2, c.clusterConfig.IndexNodeNum)
assert.Equal(t, 2, len(c.dataNodes)) s.Equal(2, len(c.dataNodes))
assert.Equal(t, 2, len(c.queryNodes)) s.Equal(2, len(c.queryNodes))
assert.Equal(t, 2, len(c.indexNodes)) s.Equal(2, len(c.indexNodes))
err = c.UpdateClusterSize(ClusterConfig{ err = c.UpdateClusterSize(ClusterConfig{
DataNodeNum: 3, DataNodeNum: 3,
QueryNodeNum: 2, QueryNodeNum: 2,
IndexNodeNum: 1, IndexNodeNum: 1,
}) })
assert.NoError(t, err) s.NoError(err)
assert.Equal(t, 3, c.clusterConfig.DataNodeNum) s.Equal(3, c.clusterConfig.DataNodeNum)
assert.Equal(t, 2, c.clusterConfig.QueryNodeNum) s.Equal(2, c.clusterConfig.QueryNodeNum)
assert.Equal(t, 1, c.clusterConfig.IndexNodeNum) s.Equal(1, c.clusterConfig.IndexNodeNum)
assert.Equal(t, 3, len(c.dataNodes)) s.Equal(3, len(c.dataNodes))
assert.Equal(t, 2, len(c.queryNodes)) s.Equal(2, len(c.queryNodes))
assert.Equal(t, 1, len(c.indexNodes)) s.Equal(1, len(c.indexNodes))
}
func TestMiniCluster(t *testing.T) {
suite.Run(t, new(MiniClusterMethodsSuite))
} }

111
tests/integration/range_search_test.go
View File

@ -20,10 +20,9 @@ import (
"context" "context"
"fmt" "fmt"
"testing" "testing"
"time"
"github.com/golang/protobuf/proto" "github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite"
"go.uber.org/zap" "go.uber.org/zap"
"github.com/milvus-io/milvus-proto/go-api/milvuspb" "github.com/milvus-io/milvus-proto/go-api/milvuspb"
@ -35,17 +34,14 @@ import (
"github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/merr"
) )
func TestRangeSearchIP(t *testing.T) { type RangeSearchSuite struct {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*180) MiniClusterSuite
c, err := StartMiniCluster(ctx) }
assert.NoError(t, err)
err = c.Start() func (s *RangeSearchSuite) TestRangeSearchIP() {
assert.NoError(t, err) c := s.Cluster
defer func() { ctx, cancel := context.WithCancel(c.GetContext())
err = c.Stop() defer cancel()
assert.NoError(t, err)
cancel()
}()
prefix := "TestRangeSearchIP" prefix := "TestRangeSearchIP"
dbName := "" dbName := ""
@ -55,7 +51,7 @@ func TestRangeSearchIP(t *testing.T) {
schema := constructSchema(collectionName, dim, true) schema := constructSchema(collectionName, dim, true)
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err) s.NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName, DbName: dbName,
@ -63,7 +59,7 @@ func TestRangeSearchIP(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: common.DefaultShardsNum, ShardsNum: common.DefaultShardsNum,
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(createCollectionStatus) err = merr.Error(createCollectionStatus)
if err != nil { if err != nil {
@ -72,8 +68,8 @@ func TestRangeSearchIP(t *testing.T) {
log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus)) log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus))
showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{}) showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{})
assert.NoError(t, err) s.NoError(err)
assert.True(t, merr.Ok(showCollectionsResp.GetStatus())) s.True(merr.Ok(showCollectionsResp.GetStatus()))
log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp)) log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp))
fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim) fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
@ -85,23 +81,23 @@ func TestRangeSearchIP(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
assert.True(t, merr.Ok(insertResult.GetStatus())) s.True(merr.Ok(insertResult.GetStatus()))
// flush // flush
flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{ flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{
DbName: dbName, DbName: dbName,
CollectionNames: []string{collectionName}, CollectionNames: []string{collectionName},
}) })
assert.NoError(t, err) s.NoError(err)
segmentIDs, has := flushResp.GetCollSegIDs()[collectionName] segmentIDs, has := flushResp.GetCollSegIDs()[collectionName]
assert.True(t, has) s.True(has)
ids := segmentIDs.GetData() ids := segmentIDs.GetData()
assert.NotEmpty(t, segmentIDs) s.NotEmpty(segmentIDs)
segments, err := c.metaWatcher.ShowSegments() segments, err := c.metaWatcher.ShowSegments()
assert.NoError(t, err) s.NoError(err)
assert.NotEmpty(t, segments) s.NotEmpty(segments)
for _, segment := range segments { for _, segment := range segments {
log.Info("ShowSegments result", zap.String("segment", segment.String())) log.Info("ShowSegments result", zap.String("segment", segment.String()))
} }
@ -114,19 +110,19 @@ func TestRangeSearchIP(t *testing.T) {
IndexName: "_default", IndexName: "_default",
ExtraParams: constructIndexParam(dim, IndexFaissIvfFlat, distance.IP), ExtraParams: constructIndexParam(dim, IndexFaissIvfFlat, distance.IP),
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(createIndexStatus) err = merr.Error(createIndexStatus)
if err != nil { if err != nil {
log.Warn("createIndexStatus fail reason", zap.Error(err)) log.Warn("createIndexStatus fail reason", zap.Error(err))
} }
waitingForIndexBuilt(ctx, c, t, collectionName, floatVecField) waitingForIndexBuilt(ctx, c, s.T(), collectionName, floatVecField)
// load // load
loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{ loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
DbName: dbName, DbName: dbName,
CollectionName: collectionName, CollectionName: collectionName,
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(loadStatus) err = merr.Error(loadStatus)
if err != nil { if err != nil {
log.Warn("LoadCollection fail reason", zap.Error(err)) log.Warn("LoadCollection fail reason", zap.Error(err))
@ -153,7 +149,7 @@ func TestRangeSearchIP(t *testing.T) {
if err != nil { if err != nil {
log.Warn("searchResult fail reason", zap.Error(err)) log.Warn("searchResult fail reason", zap.Error(err))
} }
assert.NoError(t, err) s.NoError(err)
// pass in radius and range_filter when range search // pass in radius and range_filter when range search
params["range_filter"] = filter params["range_filter"] = filter
@ -166,7 +162,7 @@ func TestRangeSearchIP(t *testing.T) {
if err != nil { if err != nil {
log.Warn("searchResult fail reason", zap.Error(err)) log.Warn("searchResult fail reason", zap.Error(err))
} }
assert.NoError(t, err) s.NoError(err)
// pass in illegal radius and range_filter when range search // pass in illegal radius and range_filter when range search
params["radius"] = filter params["radius"] = filter
@ -180,26 +176,20 @@ func TestRangeSearchIP(t *testing.T) {
if err != nil { if err != nil {
log.Warn("searchResult fail reason", zap.Error(err)) log.Warn("searchResult fail reason", zap.Error(err))
} }
assert.Error(t, err) s.Error(err)
log.Info("=========================") log.Info("=========================")
log.Info("=========================") log.Info("=========================")
log.Info("TestRangeSearchIP succeed") log.Info("TestRangeSearchIP succeed")
log.Info("=========================") log.Info("=========================")
log.Info("=========================") log.Info("=========================")
} }
func TestRangeSearchL2(t *testing.T) { func (s *RangeSearchSuite) TestRangeSearchL2() {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*180) c := s.Cluster
c, err := StartMiniCluster(ctx) ctx, cancel := context.WithCancel(c.GetContext())
assert.NoError(t, err) defer cancel()
err = c.Start()
assert.NoError(t, err)
defer func() {
err = c.Stop()
assert.NoError(t, err)
cancel()
}()
prefix := "TestRangeSearchL2" prefix := "TestRangeSearchL2"
dbName := "" dbName := ""
@ -209,7 +199,7 @@ func TestRangeSearchL2(t *testing.T) {
schema := constructSchema(collectionName, dim, true) schema := constructSchema(collectionName, dim, true)
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err) s.NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName, DbName: dbName,
@ -217,7 +207,7 @@ func TestRangeSearchL2(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: common.DefaultShardsNum, ShardsNum: common.DefaultShardsNum,
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(createCollectionStatus) err = merr.Error(createCollectionStatus)
if err != nil { if err != nil {
@ -226,8 +216,8 @@ func TestRangeSearchL2(t *testing.T) {
log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus)) log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus))
showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{}) showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{})
assert.NoError(t, err) s.NoError(err)
assert.True(t, merr.Ok(showCollectionsResp.GetStatus())) s.True(merr.Ok(showCollectionsResp.GetStatus()))
log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp)) log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp))
fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim) fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
@ -239,23 +229,23 @@ func TestRangeSearchL2(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
assert.True(t, merr.Ok(insertResult.GetStatus())) s.True(merr.Ok(insertResult.GetStatus()))
// flush // flush
flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{ flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{
DbName: dbName, DbName: dbName,
CollectionNames: []string{collectionName}, CollectionNames: []string{collectionName},
}) })
assert.NoError(t, err) s.NoError(err)
segmentIDs, has := flushResp.GetCollSegIDs()[collectionName] segmentIDs, has := flushResp.GetCollSegIDs()[collectionName]
assert.True(t, has) s.True(has)
ids := segmentIDs.GetData() ids := segmentIDs.GetData()
assert.NotEmpty(t, segmentIDs) s.NotEmpty(segmentIDs)
segments, err := c.metaWatcher.ShowSegments() segments, err := c.metaWatcher.ShowSegments()
assert.NoError(t, err) s.NoError(err)
assert.NotEmpty(t, segments) s.NotEmpty(segments)
for _, segment := range segments { for _, segment := range segments {
log.Info("ShowSegments result", zap.String("segment", segment.String())) log.Info("ShowSegments result", zap.String("segment", segment.String()))
} }
@ -268,19 +258,19 @@ func TestRangeSearchL2(t *testing.T) {
IndexName: "_default", IndexName: "_default",
ExtraParams: constructIndexParam(dim, IndexFaissIvfFlat, distance.L2), ExtraParams: constructIndexParam(dim, IndexFaissIvfFlat, distance.L2),
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(createIndexStatus) err = merr.Error(createIndexStatus)
if err != nil { if err != nil {
log.Warn("createIndexStatus fail reason", zap.Error(err)) log.Warn("createIndexStatus fail reason", zap.Error(err))
} }
waitingForIndexBuilt(ctx, c, t, collectionName, floatVecField) waitingForIndexBuilt(ctx, c, s.T(), collectionName, floatVecField)
// load // load
loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{ loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
DbName: dbName, DbName: dbName,
CollectionName: collectionName, CollectionName: collectionName,
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(loadStatus) err = merr.Error(loadStatus)
if err != nil { if err != nil {
log.Warn("LoadCollection fail reason", zap.Error(err)) log.Warn("LoadCollection fail reason", zap.Error(err))
@ -306,7 +296,7 @@ func TestRangeSearchL2(t *testing.T) {
if err != nil { if err != nil {
log.Warn("searchResult fail reason", zap.Error(err)) log.Warn("searchResult fail reason", zap.Error(err))
} }
assert.NoError(t, err) s.NoError(err)
// pass in radius and range_filter when range search // pass in radius and range_filter when range search
params["range_filter"] = filter params["range_filter"] = filter
@ -319,7 +309,7 @@ func TestRangeSearchL2(t *testing.T) {
if err != nil { if err != nil {
log.Warn("searchResult fail reason", zap.Error(err)) log.Warn("searchResult fail reason", zap.Error(err))
} }
assert.NoError(t, err) s.NoError(err)
// pass in illegal radius and range_filter when range search // pass in illegal radius and range_filter when range search
params["radius"] = filter params["radius"] = filter
@ -333,11 +323,16 @@ func TestRangeSearchL2(t *testing.T) {
if err != nil { if err != nil {
log.Warn("searchResult fail reason", zap.Error(err)) log.Warn("searchResult fail reason", zap.Error(err))
} }
assert.Error(t, err) s.Error(err)
log.Info("=========================") log.Info("=========================")
log.Info("=========================") log.Info("=========================")
log.Info("TestRangeSearchL2 succeed") log.Info("TestRangeSearchL2 succeed")
log.Info("=========================") log.Info("=========================")
log.Info("=========================") log.Info("=========================")
}
func TestRangeSearch(t *testing.T) {
suite.Run(t, new(RangeSearchSuite))
} }

99
tests/integration/refresh_config_test.go
View File

@ -29,62 +29,50 @@ import (
"github.com/milvus-io/milvus/pkg/log" "github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/distance" "github.com/milvus-io/milvus/pkg/util/distance"
"github.com/milvus-io/milvus/pkg/util/paramtable" "github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite"
"go.uber.org/zap" "go.uber.org/zap"
) )
func TestRefreshPasswordLength(t *testing.T) { type RefreshConfigSuite struct {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*180) MiniClusterSuite
defer cancel()
c, err := StartMiniCluster(ctx)
assert.NoError(t, err)
err = c.Start()
assert.NoError(t, err)
defer func() {
err = c.Stop()
assert.NoError(t, err)
cancel()
}()
s, err := c.proxy.CreateCredential(ctx, &milvuspb.CreateCredentialRequest{
Username: "test",
Password: "1234",
})
log.Debug("first create result", zap.Any("state", s))
assert.Equal(t, commonpb.ErrorCode_IllegalArgument, s.GetErrorCode())
params := paramtable.Get()
c.etcdCli.KV.Put(ctx, fmt.Sprintf("%s/config/proxy/minpasswordlength", params.EtcdCfg.RootPath.GetValue()), "3")
assert.Eventually(t, func() bool {
s, err = c.proxy.CreateCredential(ctx, &milvuspb.CreateCredentialRequest{
Username: "test",
Password: "1234",
})
log.Debug("second create result", zap.Any("state", s))
return commonpb.ErrorCode_Success == s.GetErrorCode()
}, time.Second*20, time.Millisecond*500)
} }
func TestRefreshDefaultIndexName(t *testing.T) { func (s *RefreshConfigSuite) TestRefreshPasswordLength() {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*180) c := s.Cluster
ctx, cancel := context.WithCancel(c.GetContext())
defer cancel() defer cancel()
c, err := StartMiniCluster(ctx)
assert.NoError(t, err)
err = c.Start() resp, err := c.proxy.CreateCredential(ctx, &milvuspb.CreateCredentialRequest{
assert.NoError(t, err) Username: "test",
defer func() { Password: "1234",
err = c.Stop() })
assert.NoError(t, err) log.Debug("first create result", zap.Any("state", resp))
cancel() s.Require().NoError(err)
}() s.Equal(commonpb.ErrorCode_IllegalArgument, resp.GetErrorCode())
params := paramtable.Get()
key := fmt.Sprintf("%s/config/proxy/minpasswordlength", params.EtcdCfg.RootPath.GetValue())
c.etcdCli.KV.Put(ctx, key, "3")
s.Eventually(func() bool {
resp, err = c.proxy.CreateCredential(ctx, &milvuspb.CreateCredentialRequest{
Username: "test",
Password: "1234",
})
log.Debug("second create result", zap.Any("state", resp))
return commonpb.ErrorCode_Success == resp.GetErrorCode()
}, time.Second*20, time.Millisecond*500)
}
func (s *RefreshConfigSuite) TestRefreshDefaultIndexName() {
c := s.Cluster
ctx, cancel := context.WithCancel(c.GetContext())
defer cancel()
params := paramtable.Get() params := paramtable.Get()
c.etcdCli.KV.Put(ctx, fmt.Sprintf("%s/config/common/defaultIndexName", params.EtcdCfg.RootPath.GetValue()), "a_index") c.etcdCli.KV.Put(ctx, fmt.Sprintf("%s/config/common/defaultIndexName", params.EtcdCfg.RootPath.GetValue()), "a_index")
assert.Eventually(t, func() bool { s.Eventually(func() bool {
return params.CommonCfg.DefaultIndexName.GetValue() == "a_index" return params.CommonCfg.DefaultIndexName.GetValue() == "a_index"
}, time.Second*10, time.Millisecond*500) }, time.Second*10, time.Millisecond*500)
@ -95,6 +83,7 @@ func TestRefreshDefaultIndexName(t *testing.T) {
schema := constructSchema("test", 128, true) schema := constructSchema("test", 128, true)
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
s.Require().NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: "default", DbName: "default",
@ -102,11 +91,11 @@ func TestRefreshDefaultIndexName(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: 1, ShardsNum: 1,
}) })
assert.NoError(t, err) s.NoError(err)
if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success { if createCollectionStatus.GetErrorCode() != commonpb.ErrorCode_Success {
log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason())) log.Warn("createCollectionStatus fail reason", zap.String("reason", createCollectionStatus.GetReason()))
} }
assert.Equal(t, createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success) s.Equal(createCollectionStatus.GetErrorCode(), commonpb.ErrorCode_Success)
fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim) fVecColumn := newFloatVectorFieldData(floatVecField, rowNum, dim)
hashKeys := generateHashKeys(rowNum) hashKeys := generateHashKeys(rowNum)
@ -117,19 +106,25 @@ func TestRefreshDefaultIndexName(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
_, err = c.proxy.CreateIndex(ctx, &milvuspb.CreateIndexRequest{ _, err = c.proxy.CreateIndex(ctx, &milvuspb.CreateIndexRequest{
CollectionName: collectionName, CollectionName: collectionName,
FieldName: floatVecField, FieldName: floatVecField,
ExtraParams: constructIndexParam(dim, IndexFaissIvfFlat, distance.L2), ExtraParams: constructIndexParam(dim, IndexFaissIvfFlat, distance.L2),
}) })
s.NoError(err)
s, err := c.proxy.DescribeIndex(ctx, &milvuspb.DescribeIndexRequest{ resp, err := c.proxy.DescribeIndex(ctx, &milvuspb.DescribeIndexRequest{
DbName: dbName, DbName: dbName,
CollectionName: collectionName, CollectionName: collectionName,
}) })
assert.Equal(t, commonpb.ErrorCode_Success, s.Status.GetErrorCode()) s.NoError(err)
assert.Equal(t, 1, len(s.IndexDescriptions)) s.Equal(commonpb.ErrorCode_Success, resp.Status.GetErrorCode())
assert.Equal(t, "a_index_101", s.IndexDescriptions[0].GetIndexName()) s.Equal(1, len(resp.IndexDescriptions))
s.Equal("a_index_101", resp.IndexDescriptions[0].GetIndexName())
}
func TestRefreshConfig(t *testing.T) {
suite.Run(t, new(RefreshConfigSuite))
} }

103
tests/integration/suite_test.go Normal file
View File

@ -0,0 +1,103 @@
// 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 integration
import (
"context"
"os"
"strings"
"time"
"github.com/milvus-io/milvus/pkg/util/etcd"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/stretchr/testify/suite"
"go.etcd.io/etcd/server/v3/embed"
)
// EmbedEtcdSuite contains embed setup & teardown related logic
type EmbedEtcdSuite struct {
EtcdServer *embed.Etcd
EtcdDir string
}
func (s *EmbedEtcdSuite) SetupEmbedEtcd() error {
server, folder, err := etcd.StartTestEmbedEtcdServer()
if err != nil {
return err
}
s.EtcdServer = server
s.EtcdDir = folder
return nil
}
func (s *EmbedEtcdSuite) TearDownEmbedEtcd() {
if s.EtcdServer != nil {
s.EtcdServer.Server.Stop()
}
if s.EtcdDir != "" {
os.RemoveAll(s.EtcdDir)
}
}
type MiniClusterSuite struct {
suite.Suite
EmbedEtcdSuite
Cluster *MiniCluster
cancelFunc context.CancelFunc
}
func (s *MiniClusterSuite) SetupSuite() {
s.Require().NoError(s.SetupEmbedEtcd())
}
func (s *MiniClusterSuite) TearDownSuite() {
s.TearDownEmbedEtcd()
}
func (s *MiniClusterSuite) SetupTest() {
s.T().Log("Setup test...")
// setup mini cluster to use embed etcd
endpoints := etcd.GetEmbedEtcdEndpoints(s.EtcdServer)
val := strings.Join(endpoints, ",")
// setup env value to init etcd source
s.T().Setenv("etcd.endpoints", val)
params = paramtable.Get()
ctx, cancel := context.WithTimeout(context.Background(), time.Second*180)
s.cancelFunc = cancel
c, err := StartMiniCluster(ctx, func(c *MiniCluster) {
// change config etcd endpoints
c.params[params.EtcdCfg.Endpoints.Key] = val
})
s.Require().NoError(err)
s.Cluster = c
// start mini cluster
s.Require().NoError(s.Cluster.Start())
}
func (s *MiniClusterSuite) TearDownTest() {
s.T().Log("Tear Down test...")
defer s.cancelFunc()
if s.Cluster != nil {
s.Cluster.Stop()
}
}

58
tests/integration/upsert_test.go
View File

@ -20,7 +20,6 @@ import (
"context" "context"
"fmt" "fmt"
"testing" "testing"
"time"
"github.com/golang/protobuf/proto" "github.com/golang/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/milvuspb" "github.com/milvus-io/milvus-proto/go-api/milvuspb"
@ -30,22 +29,18 @@ import (
"github.com/milvus-io/milvus/pkg/util/distance" "github.com/milvus-io/milvus/pkg/util/distance"
"github.com/milvus-io/milvus/pkg/util/funcutil" "github.com/milvus-io/milvus/pkg/util/funcutil"
"github.com/milvus-io/milvus/pkg/util/merr" "github.com/milvus-io/milvus/pkg/util/merr"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/suite"
"go.uber.org/zap" "go.uber.org/zap"
) )
func TestUpsert(t *testing.T) { type UpsertSuite struct {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*180) MiniClusterSuite
c, err := StartMiniCluster(ctx) }
assert.NoError(t, err)
err = c.Start() func (s *UpsertSuite) TestUpsert() {
assert.NoError(t, err) c := s.Cluster
defer func() { ctx, cancel := context.WithCancel(c.GetContext())
err = c.Stop() defer cancel()
assert.NoError(t, err)
cancel()
}()
assert.NoError(t, err)
prefix := "TestUpsert" prefix := "TestUpsert"
dbName := "" dbName := ""
@ -55,7 +50,7 @@ func TestUpsert(t *testing.T) {
schema := constructSchema(collectionName, dim, false) schema := constructSchema(collectionName, dim, false)
marshaledSchema, err := proto.Marshal(schema) marshaledSchema, err := proto.Marshal(schema)
assert.NoError(t, err) s.NoError(err)
createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{ createCollectionStatus, err := c.proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
DbName: dbName, DbName: dbName,
@ -63,7 +58,7 @@ func TestUpsert(t *testing.T) {
Schema: marshaledSchema, Schema: marshaledSchema,
ShardsNum: common.DefaultShardsNum, ShardsNum: common.DefaultShardsNum,
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(createCollectionStatus) err = merr.Error(createCollectionStatus)
if err != nil { if err != nil {
@ -72,8 +67,8 @@ func TestUpsert(t *testing.T) {
log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus)) log.Info("CreateCollection result", zap.Any("createCollectionStatus", createCollectionStatus))
showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{}) showCollectionsResp, err := c.proxy.ShowCollections(ctx, &milvuspb.ShowCollectionsRequest{})
assert.NoError(t, err) s.NoError(err)
assert.True(t, merr.Ok(showCollectionsResp.GetStatus())) s.True(merr.Ok(showCollectionsResp.GetStatus()))
log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp)) log.Info("ShowCollections result", zap.Any("showCollectionsResp", showCollectionsResp))
pkFieldData := newInt64FieldData(int64Field, rowNum) pkFieldData := newInt64FieldData(int64Field, rowNum)
@ -86,23 +81,23 @@ func TestUpsert(t *testing.T) {
HashKeys: hashKeys, HashKeys: hashKeys,
NumRows: uint32(rowNum), NumRows: uint32(rowNum),
}) })
assert.NoError(t, err) s.NoError(err)
assert.True(t, merr.Ok(upsertResult.GetStatus())) s.True(merr.Ok(upsertResult.GetStatus()))
// flush // flush
flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{ flushResp, err := c.proxy.Flush(ctx, &milvuspb.FlushRequest{
DbName: dbName, DbName: dbName,
CollectionNames: []string{collectionName}, CollectionNames: []string{collectionName},
}) })
assert.NoError(t, err) s.NoError(err)
segmentIDs, has := flushResp.GetCollSegIDs()[collectionName] segmentIDs, has := flushResp.GetCollSegIDs()[collectionName]
assert.True(t, has) s.True(has)
ids := segmentIDs.GetData() ids := segmentIDs.GetData()
assert.NotEmpty(t, segmentIDs) s.NotEmpty(segmentIDs)
segments, err := c.metaWatcher.ShowSegments() segments, err := c.metaWatcher.ShowSegments()
assert.NoError(t, err) s.NoError(err)
assert.NotEmpty(t, segments) s.NotEmpty(segments)
for _, segment := range segments { for _, segment := range segments {
log.Info("ShowSegments result", zap.String("segment", segment.String())) log.Info("ShowSegments result", zap.String("segment", segment.String()))
} }
@ -115,20 +110,20 @@ func TestUpsert(t *testing.T) {
IndexName: "_default", IndexName: "_default",
ExtraParams: constructIndexParam(dim, IndexFaissIvfFlat, distance.IP), ExtraParams: constructIndexParam(dim, IndexFaissIvfFlat, distance.IP),
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(createIndexStatus) err = merr.Error(createIndexStatus)
if err != nil { if err != nil {
log.Warn("createIndexStatus fail reason", zap.Error(err)) log.Warn("createIndexStatus fail reason", zap.Error(err))
} }
waitingForIndexBuilt(ctx, c, t, collectionName, floatVecField) waitingForIndexBuilt(ctx, c, s.T(), collectionName, floatVecField)
// load // load
loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{ loadStatus, err := c.proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
DbName: dbName, DbName: dbName,
CollectionName: collectionName, CollectionName: collectionName,
}) })
assert.NoError(t, err) s.NoError(err)
err = merr.Error(loadStatus) err = merr.Error(loadStatus)
if err != nil { if err != nil {
log.Warn("LoadCollection fail reason", zap.Error(err)) log.Warn("LoadCollection fail reason", zap.Error(err))
@ -150,11 +145,16 @@ func TestUpsert(t *testing.T) {
if err != nil { if err != nil {
log.Warn("searchResult fail reason", zap.Error(err)) log.Warn("searchResult fail reason", zap.Error(err))
} }
assert.NoError(t, err) s.NoError(err)
log.Info("==================") log.Info("==================")
log.Info("==================") log.Info("==================")
log.Info("TestUpsert succeed") log.Info("TestUpsert succeed")
log.Info("==================") log.Info("==================")
log.Info("==================") log.Info("==================")
}
func TestUpsert(t *testing.T) {
suite.Run(t, new(UpsertSuite))
} }