Skip to content

Commit

Permalink
test: add cases got gosdk v2 search and databases (#34204)
Browse files Browse the repository at this point in the history
issue: #33952 #33627 #33419

---------

Signed-off-by: ThreadDao <yufen.zong@zilliz.com>
  • Loading branch information
ThreadDao committed Jun 27, 2024
1 parent 9c2eeff commit be23495
Show file tree
Hide file tree
Showing 11 changed files with 1,337 additions and 34 deletions.
2 changes: 1 addition & 1 deletion client/entity/sparse.go
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ func (e sliceSparseEmbedding) Serialize() []byte {
for idx := 0; idx < e.Len(); idx++ {
pos, value, _ := e.Get(idx)
binary.LittleEndian.PutUint32(row[idx*8:], pos)
binary.LittleEndian.PutUint32(row[pos*8+4:], math.Float32bits(value))
binary.LittleEndian.PutUint32(row[idx*8+4:], math.Float32bits(value))
}
return row
}
Expand Down
10 changes: 10 additions & 0 deletions client/read.go
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,16 @@ type ResultSet struct {
// DataSet is an alias type for column slice.
type DataSet []column.Column

// GetColumn returns column with provided field name.
func (rs ResultSet) GetColumn(fieldName string) column.Column {
for _, column := range rs.Fields {
if column.Name() == fieldName {
return column
}
}
return nil
}

func (c *Client) Search(ctx context.Context, option SearchOption, callOptions ...grpc.CallOption) ([]ResultSet, error) {
req := option.Request()
collection, err := c.getCollection(ctx, req.GetCollectionName())
Expand Down
2 changes: 1 addition & 1 deletion tests/go_client/common/consts.go
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ const (

// const default value from milvus config
const (
MaxPartitionNum = 4096
MaxPartitionNum = 1024
DefaultDynamicFieldName = "$meta"
QueryCountFieldName = "count(*)"
DefaultPartition = "_default"
Expand Down
2 changes: 2 additions & 0 deletions tests/go_client/common/response_checker.go
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,8 @@ func CheckSearchResult(t *testing.T, actualSearchResults []clientv2.ResultSet, e
require.Len(t, actualSearchResults, expNq)
for _, actualSearchResult := range actualSearchResults {
require.Equal(t, actualSearchResult.ResultCount, expTopK)
require.Equal(t, actualSearchResult.IDs.Len(), expTopK)
require.Equal(t, len(actualSearchResult.Scores), expTopK)
}
}

Expand Down
4 changes: 2 additions & 2 deletions tests/go_client/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ go 1.21
toolchain go1.21.10

require (
github.com/milvus-io/milvus/client/v2 v2.0.0-20240621033600-e653ad27e2d5
github.com/milvus-io/milvus/client/v2 v2.0.0-20240625063004-b12c34a8baf2
github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3
github.com/quasilyte/go-ruleguard/dsl v0.3.22
github.com/stretchr/testify v1.9.0
Expand All @@ -14,7 +14,7 @@ require (
google.golang.org/grpc v1.64.0
)

//replace github.com/milvus-io/milvus/client/v2 v2.0.0-20240621033600-e653ad27e2d5 => ../../../milvus/client
// replace github.com/milvus-io/milvus/client/v2 v2.0.0-20240625063004-b12c34a8baf2 => ../../../milvus/client

require (
github.com/beorn7/perks v1.0.1 // indirect
Expand Down
4 changes: 2 additions & 2 deletions tests/go_client/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -403,8 +403,8 @@ github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/le
github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3 h1:KUSaWVePVlHMIluAXf2qmNffI1CMlGFLLiP+4iy9014=
github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240621033600-e653ad27e2d5 h1:KNE4Smy6HxIpoJHLpds5BI2ZyvhrE5FcLhUcsjxAxAk=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240621033600-e653ad27e2d5/go.mod h1:thfuEkUztRRmQ+qu4hCoO/6uxDJoUVNNx4vHqx9yh5I=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240625063004-b12c34a8baf2 h1:Eb3E5TQwNAImS2M1yRNc1/IzlfD8iQJ9HZt8Lf41xVc=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240625063004-b12c34a8baf2/go.mod h1:thfuEkUztRRmQ+qu4hCoO/6uxDJoUVNNx4vHqx9yh5I=
github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3 h1:ZBpRWhBa7FTFxW4YYVv9AUESoW1Xyb3KNXTzTqfkZmw=
github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3/go.mod h1:jQ2BUZny1COsgv1Qbcv8dmbppW+V9J/c4YQZNb3EOm8=
github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=
Expand Down
3 changes: 1 addition & 2 deletions tests/go_client/testcases/collection_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ func TestCreateCollection(t *testing.T) {
}
}

// func TestCreateCollection(t *testing.T) {}
func TestCreateAutoIdCollectionField(t *testing.T) {
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
mc := createDefaultMilvusClient(ctx, t)
Expand Down Expand Up @@ -626,7 +625,7 @@ func TestPartitionKeyInvalidNumPartition(t *testing.T) {
numPartitions int64
errMsg string
}{
{common.MaxPartitionNum + 1, "exceeds max configuration (4096)"},
{common.MaxPartitionNum + 1, "exceeds max configuration (1024)"},
{-1, "the specified partitions should be greater than 0 if partition key is used"},
}
for _, npStruct := range invalidNumPartitionStruct {
Expand Down
281 changes: 281 additions & 0 deletions tests/go_client/testcases/database_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,281 @@
package testcases

import (
"fmt"
"testing"
"time"

"github.com/stretchr/testify/require"
"go.uber.org/zap"

clientv2 "github.com/milvus-io/milvus/client/v2"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/tests/go_client/base"
"github.com/milvus-io/milvus/tests/go_client/common"
hp "github.com/milvus-io/milvus/tests/go_client/testcases/helper"
)

// teardownTest
func teardownTest(t *testing.T) func(t *testing.T) {
log.Info("setup test func")
return func(t *testing.T) {
log.Info("teardown func drop all non-default db")
// drop all db
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
mc := createDefaultMilvusClient(ctx, t)
dbs, _ := mc.ListDatabases(ctx, clientv2.NewListDatabaseOption())
for _, db := range dbs {
if db != common.DefaultDb {
_ = mc.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(db))
collections, _ := mc.ListCollections(ctx, clientv2.NewListCollectionOption())
for _, coll := range collections {
_ = mc.DropCollection(ctx, clientv2.NewDropCollectionOption(coll))
}
_ = mc.DropDatabase(ctx, clientv2.NewDropDatabaseOption(db))
}
}
}
}

func TestDatabase(t *testing.T) {
teardownSuite := teardownTest(t)
defer teardownSuite(t)

ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
clientDefault := createMilvusClient(ctx, t, &defaultCfg)

// create db1
dbName1 := common.GenRandomString("db1", 4)
err := clientDefault.CreateDatabase(ctx, clientv2.NewCreateDatabaseOption(dbName1))
common.CheckErr(t, err, true)

// list db and verify db1 in dbs
dbs, errList := clientDefault.ListDatabases(ctx, clientv2.NewListDatabaseOption())
common.CheckErr(t, errList, true)
require.Containsf(t, dbs, dbName1, fmt.Sprintf("%s db not in dbs: %v", dbName1, dbs))

// new client with db1 -> using db
clientDB1 := createMilvusClient(ctx, t, &clientv2.ClientConfig{Address: *addr, DBName: dbName1})
t.Log("https://github.com/milvus-io/milvus/issues/34137")
err = clientDB1.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(dbName1))
common.CheckErr(t, err, true)

// create collections -> verify collections contains
_, db1Col1 := hp.CollPrepare.CreateCollection(ctx, t, clientDB1, hp.NewCreateCollectionParams(hp.Int64Vec), hp.TNewFieldsOption(), hp.TNewSchemaOption())
_, db1Col2 := hp.CollPrepare.CreateCollection(ctx, t, clientDB1, hp.NewCreateCollectionParams(hp.Int64Vec), hp.TNewFieldsOption(), hp.TNewSchemaOption())
collections, errListCollections := clientDB1.ListCollections(ctx, clientv2.NewListCollectionOption())
common.CheckErr(t, errListCollections, true)
require.Containsf(t, collections, db1Col1.CollectionName, fmt.Sprintf("The collection %s not in: %v", db1Col1.CollectionName, collections))
require.Containsf(t, collections, db1Col2.CollectionName, fmt.Sprintf("The collection %s not in: %v", db1Col2.CollectionName, collections))

// create db2
dbName2 := common.GenRandomString("db2", 4)
err = clientDefault.CreateDatabase(ctx, clientv2.NewCreateDatabaseOption(dbName2))
common.CheckErr(t, err, true)
dbs, err = clientDefault.ListDatabases(ctx, clientv2.NewListDatabaseOption())
common.CheckErr(t, err, true)
require.Containsf(t, dbs, dbName2, fmt.Sprintf("%s db not in dbs: %v", dbName2, dbs))

// using db2 -> create collection -> drop collection
err = clientDefault.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(dbName2))
common.CheckErr(t, err, true)
_, db2Col1 := hp.CollPrepare.CreateCollection(ctx, t, clientDefault, hp.NewCreateCollectionParams(hp.Int64Vec), hp.TNewFieldsOption(), hp.TNewSchemaOption())
err = clientDefault.DropCollection(ctx, clientv2.NewDropCollectionOption(db2Col1.CollectionName))
common.CheckErr(t, err, true)

// using empty db -> drop db2
clientDefault.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(""))
err = clientDefault.DropDatabase(ctx, clientv2.NewDropDatabaseOption(dbName2))
common.CheckErr(t, err, true)

// list db and verify db drop success
dbs, err = clientDefault.ListDatabases(ctx, clientv2.NewListDatabaseOption())
common.CheckErr(t, err, true)
require.NotContains(t, dbs, dbName2)

// drop db1 which has some collections
err = clientDB1.DropDatabase(ctx, clientv2.NewDropDatabaseOption(dbName1))
common.CheckErr(t, err, false, "must drop all collections before drop database")

// drop all db1's collections -> drop db1
clientDB1.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(dbName1))
err = clientDB1.DropCollection(ctx, clientv2.NewDropCollectionOption(db1Col1.CollectionName))
common.CheckErr(t, err, true)

err = clientDB1.DropCollection(ctx, clientv2.NewDropCollectionOption(db1Col2.CollectionName))
common.CheckErr(t, err, true)

err = clientDB1.DropDatabase(ctx, clientv2.NewDropDatabaseOption(dbName1))
common.CheckErr(t, err, true)

// drop default db
err = clientDefault.DropDatabase(ctx, clientv2.NewDropDatabaseOption(common.DefaultDb))
common.CheckErr(t, err, false, "can not drop default database")

dbs, err = clientDefault.ListDatabases(ctx, clientv2.NewListDatabaseOption())
common.CheckErr(t, err, true)
require.Containsf(t, dbs, common.DefaultDb, fmt.Sprintf("The db %s not in: %v", common.DefaultDb, dbs))
}

// test create with invalid db name
func TestCreateDb(t *testing.T) {
teardownSuite := teardownTest(t)
defer teardownSuite(t)

// create db
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
mc := createDefaultMilvusClient(ctx, t)
dbName := common.GenRandomString("db", 4)
err := mc.CreateDatabase(ctx, clientv2.NewCreateDatabaseOption(dbName))
common.CheckErr(t, err, true)

// create existed db
err = mc.CreateDatabase(ctx, clientv2.NewCreateDatabaseOption(dbName))
common.CheckErr(t, err, false, fmt.Sprintf("database already exist: %s", dbName))

// create default db
err = mc.CreateDatabase(ctx, clientv2.NewCreateDatabaseOption(common.DefaultDb))
common.CheckErr(t, err, false, fmt.Sprintf("database already exist: %s", common.DefaultDb))

emptyErr := mc.CreateDatabase(ctx, clientv2.NewCreateDatabaseOption(""))
common.CheckErr(t, emptyErr, false, "database name couldn't be empty")
}

// test drop db
func TestDropDb(t *testing.T) {
teardownSuite := teardownTest(t)
defer teardownSuite(t)

// create collection in default db
listCollOpt := clientv2.NewListCollectionOption()
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
mc := createDefaultMilvusClient(ctx, t)
_, defCol := hp.CollPrepare.CreateCollection(ctx, t, mc, hp.NewCreateCollectionParams(hp.Int64Vec), hp.TNewFieldsOption(), hp.TNewSchemaOption())
collections, _ := mc.ListCollections(ctx, listCollOpt)
require.Contains(t, collections, defCol.CollectionName)

// create db
dbName := common.GenRandomString("db", 4)
err := mc.CreateDatabase(ctx, clientv2.NewCreateDatabaseOption(dbName))
common.CheckErr(t, err, true)

// using db and drop the db
err = mc.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(dbName))
common.CheckErr(t, err, true)
err = mc.DropDatabase(ctx, clientv2.NewDropDatabaseOption(dbName))
common.CheckErr(t, err, true)

// verify current db
_, err = mc.ListCollections(ctx, listCollOpt)
common.CheckErr(t, err, false, fmt.Sprintf("database not found[database=%s]", dbName))

// using default db and verify collections
err = mc.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(common.DefaultDb))
common.CheckErr(t, err, true)
collections, _ = mc.ListCollections(ctx, listCollOpt)
require.Contains(t, collections, defCol.CollectionName)

// drop not existed db
err = mc.DropDatabase(ctx, clientv2.NewDropDatabaseOption(common.GenRandomString("db", 4)))
common.CheckErr(t, err, true)

// drop empty db
err = mc.DropDatabase(ctx, clientv2.NewDropDatabaseOption(""))
common.CheckErr(t, err, false, "database name couldn't be empty")

// drop default db
err = mc.DropDatabase(ctx, clientv2.NewDropDatabaseOption(common.DefaultDb))
common.CheckErr(t, err, false, "can not drop default database")
}

// test using db
func TestUsingDb(t *testing.T) {
teardownSuite := teardownTest(t)
defer teardownSuite(t)

// create collection in default db
listCollOpt := clientv2.NewListCollectionOption()
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
mc := createDefaultMilvusClient(ctx, t)

_, col := hp.CollPrepare.CreateCollection(ctx, t, mc, hp.NewCreateCollectionParams(hp.Int64Vec), hp.TNewFieldsOption(), hp.TNewSchemaOption())
// collName := createDefaultCollection(ctx, t, mc, true, common.DefaultShards)
collections, _ := mc.ListCollections(ctx, listCollOpt)
require.Contains(t, collections, col.CollectionName)

// using not existed db
dbName := common.GenRandomString("db", 4)
err := mc.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(dbName))
common.CheckErr(t, err, false, fmt.Sprintf("database not found[database=%s]", dbName))

// using empty db
err = mc.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(""))
common.CheckErr(t, err, true)
collections, _ = mc.ListCollections(ctx, listCollOpt)
require.Contains(t, collections, col.CollectionName)

// using current db
err = mc.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(common.DefaultDb))
common.CheckErr(t, err, true)
collections, _ = mc.ListCollections(ctx, listCollOpt)
require.Contains(t, collections, col.CollectionName)
}

func TestClientWithDb(t *testing.T) {
t.Skip("https://github.com/milvus-io/milvus/issues/34137")
teardownSuite := teardownTest(t)
defer teardownSuite(t)

listCollOpt := clientv2.NewListCollectionOption()
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)

// connect with not existed db
_, err := base.NewMilvusClient(ctx, &clientv2.ClientConfig{Address: *addr, DBName: "dbName"})
common.CheckErr(t, err, false, "database not found")

// connect default db -> create a collection in default db
mcDefault, errDefault := base.NewMilvusClient(ctx, &clientv2.ClientConfig{
Address: *addr,
// DBName: common.DefaultDb,
})
common.CheckErr(t, errDefault, true)
_, defCol1 := hp.CollPrepare.CreateCollection(ctx, t, mcDefault, hp.NewCreateCollectionParams(hp.Int64Vec), hp.TNewFieldsOption(), hp.TNewSchemaOption())
defCollections, _ := mcDefault.ListCollections(ctx, listCollOpt)
require.Contains(t, defCollections, defCol1.CollectionName)
log.Debug("default db collections:", zap.Any("default collections", defCollections))

// create a db and create collection in db
dbName := common.GenRandomString("db", 5)
err = mcDefault.CreateDatabase(ctx, clientv2.NewCreateDatabaseOption(dbName))
common.CheckErr(t, err, true)

// and connect with db
mcDb, err := base.NewMilvusClient(ctx, &clientv2.ClientConfig{
Address: *addr,
DBName: dbName,
})
common.CheckErr(t, err, true)
_, dbCol1 := hp.CollPrepare.CreateCollection(ctx, t, mcDb, hp.NewCreateCollectionParams(hp.Int64Vec), hp.TNewFieldsOption(), hp.TNewSchemaOption())

dbCollections, _ := mcDb.ListCollections(ctx, listCollOpt)
log.Debug("db collections:", zap.Any("db collections", dbCollections))
require.Containsf(t, dbCollections, dbCol1.CollectionName, fmt.Sprintf("The collection %s not in: %v", dbCol1.CollectionName, dbCollections))

// using default db and collection not in
_ = mcDb.UsingDatabase(ctx, clientv2.NewUsingDatabaseOption(common.DefaultDb))
defCollections, _ = mcDb.ListCollections(ctx, listCollOpt)
require.NotContains(t, defCollections, dbCol1.CollectionName)

// connect empty db (actually default db)
mcEmpty, err := base.NewMilvusClient(ctx, &clientv2.ClientConfig{
Address: *addr,
DBName: "",
})
common.CheckErr(t, err, true)
defCollections, _ = mcEmpty.ListCollections(ctx, listCollOpt)
require.Contains(t, defCollections, defCol1.CollectionName)
}

func TestAlterDatabase(t *testing.T) {
t.Skip("waiting for AlterDatabase and DescribeDatabase")
}
18 changes: 18 additions & 0 deletions tests/go_client/testcases/main_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,24 @@ func createDefaultMilvusClient(ctx context.Context, t *testing.T) *base.MilvusCl
return mc
}

// create connect
func createMilvusClient(ctx context.Context, t *testing.T, cfg *clientv2.ClientConfig) *base.MilvusClient {
t.Helper()

var (
mc *base.MilvusClient
err error
)
mc, err = base.NewMilvusClient(ctx, cfg)
common.CheckErr(t, err, true)

t.Cleanup(func() {
mc.Close(ctx)
})

return mc
}

func TestMain(m *testing.M) {
flag.Parse()
log.Info("Parser Milvus address", zap.String("address", *addr))
Expand Down
Loading

0 comments on commit be23495

Please sign in to comment.