mirror of https://github.com/milvus-io/milvus.git
433 lines
9.9 KiB
Go
433 lines
9.9 KiB
Go
package test
|
|
|
|
import (
|
|
"encoding/json"
|
|
"github.com/milvus-io/milvus-sdk-go/milvus"
|
|
"testing"
|
|
)
|
|
|
|
var TABLENAME string = "go_test"
|
|
var client milvus.MilvusClient = GetClient()
|
|
|
|
type segmentsInfo struct {
|
|
DataSize int64 `json:"data_size"`
|
|
IndexName string `json:"index_name"`
|
|
Name string `json:"name"`
|
|
RowCount int64 `json:"row_count"`
|
|
}
|
|
|
|
type partitionsInfo struct {
|
|
RowCount int64 `json:"row_count"`
|
|
Segments []segmentsInfo `json:"segments"`
|
|
Tag string `json:"tag"`
|
|
}
|
|
|
|
type collectionInfo struct {
|
|
Partitions []partitionsInfo `json:"partitions"`
|
|
RowCount int64 `json:"row_count"`
|
|
}
|
|
|
|
func GetClient() milvus.MilvusClient {
|
|
var grpcClient milvus.Milvusclient
|
|
client := milvus.NewMilvusClient(grpcClient.Instance)
|
|
connectParam := milvus.ConnectParam{"127.0.0.1", "19530"}
|
|
err := client.Connect(connectParam)
|
|
if err != nil {
|
|
println("Connect failed")
|
|
return nil
|
|
}
|
|
return client
|
|
}
|
|
|
|
func CreateCollection() error {
|
|
boolReply, status, err := client.HasCollection(TABLENAME)
|
|
if boolReply == true {
|
|
return err
|
|
}
|
|
|
|
collectionParam := milvus.CollectionParam{TABLENAME, 128, 1024, int64(milvus.L2)}
|
|
status, err = client.CreateCollection(collectionParam)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !status.Ok() {
|
|
println("CreateCollection failed")
|
|
}
|
|
return err
|
|
}
|
|
|
|
func TestConnection(t *testing.T) {
|
|
var grpcClient milvus.Milvusclient
|
|
testClient := milvus.NewMilvusClient(grpcClient.Instance)
|
|
connectParam := milvus.ConnectParam{"127.0.0.1", "19530"}
|
|
err := testClient.Connect(connectParam)
|
|
if err != nil {
|
|
t.Error("Connect error: " + err.Error())
|
|
}
|
|
|
|
// test wrong uri connect
|
|
connectParam = milvus.ConnectParam{"12345", "111"}
|
|
err = testClient.Connect(connectParam)
|
|
if err == nil {
|
|
t.Error("Use wrong uri to connect, return true")
|
|
}
|
|
}
|
|
|
|
func TestCollection(t *testing.T) {
|
|
param := milvus.CollectionParam{"test_1", 128, 1024, int64(milvus.L2)}
|
|
status, err := client.CreateCollection(param)
|
|
if err != nil {
|
|
t.Error("CreateCollection error")
|
|
}
|
|
|
|
if !status.Ok() {
|
|
t.Error("CreateCollection return status wrong!")
|
|
}
|
|
|
|
// test ShowCollections
|
|
collections, status, err := client.ListCollections()
|
|
if err != nil {
|
|
t.Error("ShowCollections error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("ShowCollections status check error")
|
|
}
|
|
if len(collections) != 1 && collections[0] != TABLENAME {
|
|
t.Error("ShowCollections result check error")
|
|
}
|
|
|
|
// test normal hascollection
|
|
hasCollection, status, err := client.HasCollection("test_1")
|
|
if err != nil {
|
|
t.Error("HasCollection error")
|
|
return
|
|
}
|
|
|
|
if !status.Ok() {
|
|
t.Error("HasCollection status check error")
|
|
}
|
|
|
|
if !hasCollection {
|
|
t.Error("HasCollection result check error")
|
|
}
|
|
|
|
// test HasCollection with collection not exist
|
|
hasCollection, status, err = client.HasCollection("aaa")
|
|
if err != nil {
|
|
t.Error("HasCollection error")
|
|
}
|
|
|
|
if hasCollection == true {
|
|
t.Error("HasCollection result check error")
|
|
}
|
|
|
|
// test DropCollection
|
|
status, err = client.DropCollection("test_1")
|
|
if err != nil {
|
|
t.Error("DropCollection error")
|
|
}
|
|
|
|
if !status.Ok() {
|
|
t.Error("DropCollection status check error")
|
|
}
|
|
|
|
hasCollection, status, err = client.HasCollection("test_1")
|
|
if hasCollection == true {
|
|
t.Error("DropCollection result check error")
|
|
}
|
|
|
|
// test DropCollection with collection not exist
|
|
status, err = client.DropCollection("aaa")
|
|
if err != nil {
|
|
t.Error("DropCollection error")
|
|
}
|
|
|
|
if status.Ok() {
|
|
t.Error("DropCollection status check error")
|
|
}
|
|
}
|
|
|
|
func TestEntity(t *testing.T) {
|
|
err := CreateCollection()
|
|
if err != nil {
|
|
t.Error("Create collection error")
|
|
}
|
|
// test insert
|
|
var i, j int
|
|
|
|
nb := 10000
|
|
dimension := 128
|
|
records := make([]milvus.Entity, nb)
|
|
recordArray := make([][]float32, 10000)
|
|
for i = 0; i < nb; i++ {
|
|
recordArray[i] = make([]float32, dimension)
|
|
for j = 0; j < dimension; j++ {
|
|
recordArray[i][j] = float32(i % (j + 1))
|
|
}
|
|
records[i].FloatData = recordArray[i]
|
|
}
|
|
insertParam := milvus.InsertParam{TABLENAME, "", records, nil}
|
|
ret_id_array, status, err := client.Insert(&insertParam)
|
|
if err != nil {
|
|
t.Error("Insert error")
|
|
return
|
|
}
|
|
|
|
if ret_id_array == nil {
|
|
t.Error("Insert vecto result check error")
|
|
}
|
|
|
|
if !status.Ok() {
|
|
t.Error("Insert status check error")
|
|
}
|
|
|
|
// Flush
|
|
collection_array := make([]string, 1)
|
|
collection_array[0] = TABLENAME
|
|
status, err = client.Flush(collection_array)
|
|
if err != nil {
|
|
t.Error("Flush error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("Flush status check error")
|
|
}
|
|
|
|
// test ShowCollectionInfos
|
|
collectionInfos, status, err := client.GetCollectionStats(TABLENAME)
|
|
if err != nil {
|
|
t.Error("ShowCollectionInfo error")
|
|
return
|
|
}
|
|
|
|
if !status.Ok() {
|
|
t.Error("ShowCollectionInfo status check error")
|
|
}
|
|
|
|
if len(collectionInfos) == 0 {
|
|
t.Error("ShowCollectionInfo result check error")
|
|
}
|
|
|
|
collections := collectionInfo{}
|
|
if err = json.Unmarshal([]byte(collectionInfos), &collections); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
segmentName := collections.Partitions[0].Segments[0].Name
|
|
|
|
// test GetEntityIds
|
|
listIDInSegment := milvus.ListIDInSegmentParam{TABLENAME, string(segmentName)}
|
|
entityIDs, status, err := client.ListIDInSegment(listIDInSegment)
|
|
if err != nil {
|
|
t.Error("GetEntityIDs error")
|
|
return
|
|
}
|
|
|
|
if len(entityIDs) == 0 {
|
|
t.Error("GetEntityIDs result check error")
|
|
}
|
|
|
|
// test GetEntityById
|
|
rowRecord, status, err := client.GetEntityByID(TABLENAME, entityIDs[0:10])
|
|
if err != nil {
|
|
t.Error("GetEntityByID error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("GetEntitiesByID status check error: ", status.GetMessage())
|
|
}
|
|
if len(rowRecord) == 0 {
|
|
t.Error("GetEntityiesByID result check error: ", len(rowRecord))
|
|
}
|
|
|
|
// test DeleteByID
|
|
id_array := make([]int64, 1)
|
|
id_array[0] = entityIDs[0]
|
|
status, err = client.DeleteEntityByID(TABLENAME, id_array)
|
|
if err != nil {
|
|
t.Error("DeleteByID error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("DeleteByID status check error")
|
|
}
|
|
}
|
|
|
|
func TestIndex(t *testing.T) {
|
|
extraParam := "{\"nlist\" : 16384}"
|
|
indexParam := milvus.IndexParam{TABLENAME, milvus.IVFFLAT, extraParam}
|
|
status, err := client.CreateIndex(&indexParam)
|
|
if err != nil {
|
|
t.Error("CreateIndex error")
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("CreateIndex status check error")
|
|
}
|
|
|
|
// test DescribeIndex
|
|
indexInfo, status, err := client.GetIndexInfo(TABLENAME)
|
|
if err != nil {
|
|
t.Error("DescribeIndex error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("DescribeIndex status check error")
|
|
}
|
|
if indexInfo.CollectionName != TABLENAME || indexInfo.IndexType != milvus.IVFFLAT {
|
|
t.Error("DescribeIndex result chck error")
|
|
}
|
|
|
|
// test DropIndex
|
|
status, err = client.DropIndex(TABLENAME)
|
|
if err != nil {
|
|
t.Error("DropIndex error")
|
|
return
|
|
}
|
|
|
|
if !status.Ok() {
|
|
t.Error("DropIndex status check erro")
|
|
}
|
|
|
|
status, err = client.CreateIndex(&indexParam)
|
|
if err != nil {
|
|
t.Error("CreateIndex error")
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("CreateIndex status check error")
|
|
}
|
|
}
|
|
|
|
func TestSearch(t *testing.T) {
|
|
var i, j int
|
|
//Construct query entities
|
|
nq := 10
|
|
dimension := 128
|
|
topk := 10
|
|
queryRecords := make([]milvus.Entity, nq)
|
|
queryEntitys := make([][]float32, nq)
|
|
for i = 0; i < nq; i++ {
|
|
queryEntitys[i] = make([]float32, dimension)
|
|
for j = 0; j < dimension; j++ {
|
|
queryEntitys[i][j] = float32(i % (j + 1))
|
|
}
|
|
queryRecords[i].FloatData = queryEntitys[i]
|
|
}
|
|
|
|
var topkQueryResult milvus.TopkQueryResult
|
|
extraParam := "{\"nprobe\" : 32}"
|
|
searchParam := milvus.SearchParam{TABLENAME, queryRecords, int64(topk), nil, extraParam}
|
|
topkQueryResult, status, err := client.Search(searchParam)
|
|
if err != nil {
|
|
t.Error(err.Error())
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("Search status check error")
|
|
}
|
|
if len(topkQueryResult.QueryResultList) != nq {
|
|
t.Error("Search result check error")
|
|
}
|
|
}
|
|
|
|
func TestCmd(t *testing.T) {
|
|
// test ServerStatus
|
|
serverStatus, status, err := client.ServerStatus()
|
|
if err != nil {
|
|
t.Error("ServerStatus error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("ServerStatus status check error")
|
|
}
|
|
if serverStatus != "server alive" {
|
|
t.Error("ServerStatus result check error: " + serverStatus)
|
|
}
|
|
|
|
// test ServerVersion
|
|
serverVersion, status, err := client.ServerVersion()
|
|
if err != nil {
|
|
t.Error("ServerVersion error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("ServerVersion status check error")
|
|
}
|
|
if len(serverVersion) == 0 {
|
|
t.Error("ServerVersion result check error")
|
|
}
|
|
|
|
// test SetConfig and GetConfig
|
|
nodeName := "cache.cache_size"
|
|
nodeValue := "2"
|
|
status, err = client.SetConfig(nodeName, nodeValue)
|
|
if err != nil {
|
|
t.Error("SetConfig error")
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("SetConfig status check error: " + status.GetMessage())
|
|
}
|
|
|
|
value, status, err := client.GetConfig(nodeName)
|
|
if err != nil {
|
|
t.Error("GetConfig error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("GetConfig status check error")
|
|
}
|
|
if value != nodeValue {
|
|
t.Error("GetConfig or SetConfig result check error")
|
|
}
|
|
}
|
|
|
|
func TestPartition(t *testing.T) {
|
|
// test CreatePartition
|
|
partitionTag := "part_1"
|
|
status, err := client.CreatePartition(milvus.PartitionParam{TABLENAME, partitionTag})
|
|
if err != nil {
|
|
t.Error("CreatePartition error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("CreatePartition status check error")
|
|
}
|
|
|
|
// test ListPartitions
|
|
partitionParam, status, err := client.ListPartitions(TABLENAME)
|
|
if !status.Ok() {
|
|
t.Error("ShowPartitions status check error")
|
|
}
|
|
if len(partitionParam) == 0 {
|
|
t.Error("ShowPartitions result check error")
|
|
}
|
|
|
|
// test DropPartition
|
|
status, err = client.DropPartition(milvus.PartitionParam{TABLENAME, partitionTag})
|
|
if !status.Ok() {
|
|
t.Error("DropPartition status check error")
|
|
}
|
|
}
|
|
|
|
func TestFlush(t *testing.T) {
|
|
collectionNameArray := make([]string, 1)
|
|
collectionNameArray[0] = TABLENAME
|
|
status, err := client.Flush(collectionNameArray)
|
|
if err != nil {
|
|
t.Error("Flush error")
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("Flush status check error")
|
|
}
|
|
}
|
|
|
|
func TestCompact(t *testing.T) {
|
|
status, err := client.Compact(TABLENAME)
|
|
if err != nil {
|
|
t.Error("Compact error")
|
|
return
|
|
}
|
|
if !status.Ok() {
|
|
t.Error("Compact status check error")
|
|
}
|
|
}
|