Add highlevel cases and update utils func (#546)

Signed-off-by: ThreadDao <[email protected]>
milvus-io · Aug 9, 2023 · 0c4a500 · 0c4a500
1 parent 43f08ee
commit 0c4a500
Show file tree

Hide file tree

Showing 12 changed files with 480 additions and 229 deletions.
diff --git a/.github/workflows/test-ci.yaml b/.github/workflows/test-ci.yaml
@@ -70,7 +70,7 @@ jobs:
         run: |
           nc -vz 127.0.0.1 19530
           curl -sSL "https://github.com/gotestyourself/gotestsum/releases/download/v1.8.2/gotestsum_1.8.2_linux_amd64.tar.gz" | sudo tar -xz -C /usr/local/bin gotestsum
-          gotestsum --format testname --hide-summary=output ./testcases/... --tags L0 --addr=127.0.0.1:19530
+          gotestsum --format testname --hide-summary=output ./testcases/... --tags L0 --addr=127.0.0.1:19530 -timeout=20m
 
       - name: Export logs
         if: ${{ !success() }}

diff --git a/test/base/milvus_client.go b/test/base/milvus_client.go
@@ -132,6 +132,14 @@ func (mc *MilvusClient) CreateCollection(ctx context.Context, collSchema *entity
 	return err
 }
 
+// NewCollection highlevel new Collection
+func (mc *MilvusClient) NewCollection(ctx context.Context, collName string, dim int64, opts ...client.CreateCollectionOption) error {
+	preRequest("NewCollection", ctx, collName, dim, opts)
+	err := mc.mClient.NewCollection(ctx, collName, dim, opts...)
+	postResponse("NewCollection", err)
+	return err
+}
+
 // ListCollections list collections
 func (mc *MilvusClient) ListCollections(ctx context.Context) ([]*entity.Collection, error) {
 	preRequest("ListCollections", ctx)
@@ -378,18 +386,17 @@ func (mc *MilvusClient) Search(ctx context.Context, collName string, partitions
 	outputFields []string, vectors []entity.Vector, vectorField string, metricType entity.MetricType, topK int, sp entity.SearchParam, opts ...client.SearchQueryOptionFunc,
 ) ([]client.SearchResult, error) {
 	funcName := "Search"
-	preRequest(funcName, ctx, collName, partitions, expr, outputFields, vectors, vectorField, metricType, topK, sp, opts)
+	preRequest(funcName, ctx, collName, partitions, expr, outputFields, len(vectors), vectorField, metricType, topK, sp, opts)
 
 	searchResult, err := mc.mClient.Search(ctx, collName, partitions, expr, outputFields, vectors, vectorField, metricType, topK, sp, opts...)
 	postResponse(funcName, err, searchResult)
 
 	return searchResult, err
 }
 
-// Query query from collection
-func (mc *MilvusClient) Query(ctx context.Context, collName string, partitions []string, ids entity.Column,
-	outputFields []string, opts ...client.SearchQueryOptionFunc,
-) ([]entity.Column, error) {
+// QueryByPks query from collection
+func (mc *MilvusClient) QueryByPks(ctx context.Context, collName string, partitions []string, ids entity.Column,
+	outputFields []string, opts ...client.SearchQueryOptionFunc) (client.ResultSet, error) {
 	funcName := "QueryByPks"
 	preRequest(funcName, ctx, collName, partitions, ids, outputFields, opts)
 
@@ -399,6 +406,28 @@ func (mc *MilvusClient) Query(ctx context.Context, collName string, partitions [
 	return queryResults, err
 }
 
+// Query query from collection
+func (mc *MilvusClient) Query(ctx context.Context, collName string, partitions []string, expr string, outputFields []string, opts ...client.SearchQueryOptionFunc,
+) (client.ResultSet, error) {
+	funcName := "Query"
+	preRequest(funcName, ctx, collName, partitions, expr, outputFields, opts)
+
+	queryResults, err := mc.mClient.Query(ctx, collName, partitions, expr, outputFields, opts...)
+
+	postResponse(funcName, err, queryResults)
+	return queryResults, err
+}
+
+// Get query from collection
+func (mc *MilvusClient) Get(ctx context.Context, collName string, ids entity.Column, opts ...client.GetOption,
+) (client.ResultSet, error) {
+	funcName := "Get"
+	preRequest(funcName, ctx, collName, ids, opts)
+	queryResults, err := mc.mClient.Get(ctx, collName, ids, opts...)
+	postResponse(funcName, err, queryResults)
+	return queryResults, err
+}
+
 // -- row based apis --
 
 // CreateCollectionByRow Create Collection By Row
@@ -465,6 +494,22 @@ func (mc *MilvusClient) ListBulkInsertTasks(ctx context.Context, collName string
 	return bulkInsertTaskStates, err
 }
 
+// GetLoadingProgress
+func (mc *MilvusClient) GetLoadingProgress(ctx context.Context, collName string, partitionNames []string) (int64, error) {
+	preRequest("GetLoadingProgress", ctx, collName, partitionNames)
+	loadingProgress, err := mc.mClient.GetLoadingProgress(ctx, collName, partitionNames)
+	postResponse("GetLoadingProgress", err, loadingProgress)
+	return loadingProgress, err
+}
+
+// GetLoadState
+func (mc *MilvusClient) GetLoadState(ctx context.Context, collName string, partitionNames []string) (entity.LoadState, error) {
+	preRequest("GetLoadState", ctx, collName, partitionNames)
+	loadState, err := mc.mClient.GetLoadState(ctx, collName, partitionNames)
+	postResponse("GetLoadState", err, loadState)
+	return loadState, err
+}
+
 // ListResourceGroups List Resource Groups
 func (mc *MilvusClient) ListResourceGroups(ctx context.Context) ([]string, error) {
 	preRequest("ListResourceGroups", ctx)