Repository: milvus-io/milvus-sdk-go
Pull request #516 (2023-07-14 07:54:33)

Sync v2.2.x change to master
CP change from #493 #496 #502 #505 #506 #508 #509
/kind improvement

sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/516#" title="Author self-approved">congqixia</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
> Merging [#516](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (108b0c3) into [master](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/commit/565c5e3e4acf2c83d71e9846429dd8bb5b1328dc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (565c5e3) will **decrease** coverage by `0.43%`.
> The diff coverage is `68.58%`.

[![Impacted file tree graph](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516/graphs/tree.svg?width=650&height=150&src=pr&token=WJK43PV390&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

```diff
@@            Coverage Diff             @@
##           master     #516      +/-   ##
==========================================
- Coverage   83.62%   83.19%   -0.43%     
==========================================
  Files          38       39       +1     
  Lines        3383     3464      +81     
==========================================
+ Hits         2829     2882      +53     
- Misses        447      474      +27     
- Partials      107      108       +1     
```


| [Impacted Files](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) | Coverage Δ | |
|---|---|---|
| [client/client.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudC5nbw==) | `28.94% <ø> (-2.76%)` | :arrow_down: |
| [client/database.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2RhdGFiYXNlLmdv) | `66.66% <0.00%> (-12.13%)` | :arrow_down: |
| [client/errors.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2Vycm9ycy5nbw==) | `50.00% <ø> (ø)` | |
| [client/grpc\_client.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2dycGNfY2xpZW50Lmdv) | `53.33% <35.71%> (ø)` | |
| [client/options.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L29wdGlvbnMuZ28=) | `61.11% <46.15%> (-8.12%)` | :arrow_down: |
| [client/data.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2RhdGEuZ28=) | `69.40% <53.84%> (-0.60%)` | :arrow_down: |
| [client/results.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L3Jlc3VsdHMuZ28=) | `66.66% <66.66%> (ø)` | |
| [client/collection.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NvbGxlY3Rpb24uZ28=) | `88.58% <84.41%> (+0.77%)` | :arrow_up: |
| [client/partition.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L3BhcnRpdGlvbi5nbw==) | `83.48% <88.88%> (-0.77%)` | :arrow_down: |
| [client/config.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NvbmZpZy5nbw==) | `67.34% <100.00%> (+2.12%)` | :arrow_up: |
| ... and [1 more](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/516?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) | |


Patch:
diff --git a/client/client.go b/client/client.go
index f737cea..fc44d6f 100644
--- a/client/client.go
+++ b/client/client.go
@@ -48,6 +48,8 @@ type Client interface {
 
 	// -- collection --
 
+	// NewCollection intializeds a new collection with pre defined attributes
+	NewCollection(ctx context.Context, collName string, dimension int64, opts ...CreateCollectionOption) error
 	// ListCollections list collections from connection
 	ListCollections(ctx context.Context) ([]*entity.Collection, error)
 	// CreateCollection create collection using provided schema
@@ -138,6 +140,8 @@ type Client interface {
 	QueryByPks(ctx context.Context, collectionName string, partitionNames []string, ids entity.Column, outputFields []string, opts ...SearchQueryOptionFunc) (ResultSet, error)
 	// Query performs query records with boolean expression.
 	Query(ctx context.Context, collectionName string, partitionNames []string, expr string, outputFields []string, opts ...SearchQueryOptionFunc) (ResultSet, error)
+	// Get grabs the inserted entities using the primary key from the Collection.
+	Get(ctx context.Context, collectionName string, ids entity.Column, opts ...GetOption) (ResultSet, error)
 
 	// CalcDistance calculate the distance between vectors specified by ids or provided
 	CalcDistance(ctx context.Context, collName string, partitions []string,
@@ -203,34 +207,6 @@ type Client interface {
 	GetVersion(ctx context.Context) (string, error)
 }
 
-// SearchResult contains the result from Search api of client
-// IDs is the auto generated id values for the entities
-// Fields contains the data of `outputFieleds` specified or all columns if non
-// Scores is actually the distance between the vector current record contains and the search target vector
-type SearchResult struct {
-	ResultCount int           // the returning entry count
-	IDs         entity.Column // auto generated id, can be mapped to the columns from `Insert` API
-	Fields      ResultSet     //[]entity.Column // output field data
-	Scores      []float32     // distance to the target vector
-	Err         error         // search error if any
-}
-
-// ResultSet is an alias type for column slice.
-type ResultSet []entity.Column
-
-// GetColumn returns column with provided field name.
-func (rs ResultSet) GetColumn(fieldName string) entity.Column {
-	for _, column := range rs {
-		if column.Name() == fieldName {
-			return column
-		}
-	}
-	return nil
-}
-
-// Check if GrpcClient implement Client.
-var _ Client = &GrpcClient{}
-
 // NewClient create a client connected to remote milvus cluster.
 // More connect option can be modified by Config.
 func NewClient(ctx context.Context, config Config) (Client, error) {
diff --git a/client/collection.go b/client/collection.go
index 74f22f8..825db5e 100644
--- a/client/collection.go
+++ b/client/collection.go
@@ -13,7 +13,6 @@ package client
 
 import (
 	"context"
-	"fmt"
 	"time"
 
 	"github.com/cockroachdb/errors"
@@ -75,41 +74,92 @@ func (c *GrpcClient) ListCollections(ctx context.Context) ([]*entity.Collection,
 	return collections, nil
 }
 
-// CreateCollection create collection with specified schema
-func (c *GrpcClient) CreateCollection(ctx context.Context, collSchema *entity.Schema, shardNum int32, opts ...CreateCollectionOption) error {
+// NewCollection creates a common simple collection with pre-defined attributes.
+func (c *GrpcClient) NewCollection(ctx context.Context, collName string, dimension int64, opts ...CreateCollectionOption) error {
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
-	if err := validateSchema(collSchema); err != nil {
+
+	//	shardNum := entity.DefaultShardNumber
+	opt := &createCollOpt{
+		ConsistencyLevel:    entity.DefaultConsistencyLevel,
+		PrimaryKeyFieldName: "id",
+		PrimaryKeyFieldType: entity.FieldTypeInt64,
+		VectorFieldName:     "vector",
+		MetricsType:         entity.IP,
+		AutoID:              false,
+		EnableDynamicSchema: true,
+	}
+
+	for _, o := range opts {
+		o(opt)
+	}
+
+	pkField := entity.NewField().WithName(opt.PrimaryKeyFieldName).WithDataType(opt.PrimaryKeyFieldType).WithIsAutoID(opt.AutoID).WithIsPrimaryKey(true)
+	if opt.PrimaryKeyFieldType == entity.FieldTypeVarChar && opt.PrimaryKeyMaxLength > 0 {
+		pkField = pkField.WithMaxLength(opt.PrimaryKeyMaxLength)
+	}
+
+	sch := entity.NewSchema().WithName(collName).WithAutoID(opt.AutoID).WithDynamicFieldEnabled(opt.EnableDynamicSchema).
+		WithField(pkField).
+		WithField(entity.NewField().WithName(opt.VectorFieldName).WithDataType(entity.FieldTypeFloatVector).WithDim(dimension))
+
+	if err := c.validateSchema(sch); err != nil {
 		return err
 	}
 
-	has, err := c.HasCollection(ctx, collSchema.CollectionName)
-	if err != nil {
+	if err := c.requestCreateCollection(ctx, sch, opt, entity.DefaultShardNumber); err != nil {
 		return err
 	}
-	if has {
-		return fmt.Errorf("collection %s already exist", collSchema.CollectionName)
+
+	idx := entity.NewGenericIndex("", "", map[string]string{
+		"metric_type": string(opt.MetricsType),
+	})
+
+	if err := c.CreateIndex(ctx, collName, opt.VectorFieldName, idx, false); err != nil {
+		return err
 	}
-	sch := collSchema.ProtoMessage()
-	bs, err := proto.Marshal(sch)
-	if err != nil {
+
+	return c.LoadCollection(ctx, collName, false)
+}
+
+// CreateCollection create collection with specified schema
+func (c *GrpcClient) CreateCollection(ctx context.Context, collSchema *entity.Schema, shardNum int32, opts ...CreateCollectionOption) error {
+	if c.Service == nil {
+		return ErrClientNotReady
+	}
+	if err := c.validateSchema(collSchema); err != nil {
 		return err
 	}
-	req := &server.CreateCollectionRequest{
-		DbName:         "", // reserved fields, not used for now
-		CollectionName: collSchema.CollectionName,
-		Schema:         bs,
-		ShardsNum:      shardNum,
-		// default consistency level is strong
-		// to be consistent with previous version
-		ConsistencyLevel: common.ConsistencyLevel_Bounded,
+
+	opt := &createCollOpt{
+		ConsistencyLevel: entity.DefaultConsistencyLevel,
 		NumPartitions:    0,
 	}
 	// apply options on request
-	for _, opt := range opts {
-		opt(req)
+	for _, o := range opts {
+		o(opt)
 	}
+
+	return c.requestCreateCollection(ctx, collSchema, opt, shardNum)
+}
+
+func (c *GrpcClient) requestCreateCollection(ctx context.Context, sch *entity.Schema, opt *createCollOpt, shardNum int32) error {
+	bs, err := proto.Marshal(sch.ProtoMessage())
+	if err != nil {
+		return err
+	}
+
+	req := &server.CreateCollectionRequest{
+		DbName:           "", // reserved fields, not used for now
+		CollectionName:   sch.CollectionName,
+		Schema:           bs,
+		ShardsNum:        shardNum,
+		ConsistencyLevel: opt.ConsistencyLevel.CommonConsistencyLevel(),
+		NumPartitions:    opt.NumPartitions,
+		Properties:       entity.MapKvPairs(opt.Properties),
+	}
+
 	resp, err := c.Service.CreateCollection(ctx, req)
 	if err != nil {
 		return err
@@ -121,7 +171,7 @@ func (c *GrpcClient) CreateCollection(ctx context.Context, collSchema *entity.Sc
 	return nil
 }
 
-func validateSchema(sch *entity.Schema) error {
+func (c *GrpcClient) validateSchema(sch *entity.Schema) error {
 	if sch == nil {
 		return errors.New("nil schema")
 	}
@@ -132,6 +182,9 @@ func validateSchema(sch *entity.Schema) error {
 	primaryKey := false
 	autoID := false
 	vectors := 0
+	hasPartitionKey := false
+	hasDynamicSchema := sch.EnableDynamicField
+	hasJSON := false
 	for _, field := range sch.Fields {
 		if field.PrimaryKey {
 			if primaryKey { // another primary key found, only one primary key field for now
@@ -151,6 +204,15 @@ func validateSchema(sch *entity.Schema) error {
 			}
 			autoID = true
 		}
+		if field.DataType == entity.FieldTypeJSON {
+			hasJSON = true
+		}
+		if field.IsDynamic {
+			hasDynamicSchema = true
+		}
+		if field.IsPartitionKey {
+			hasPartitionKey = true
+		}
 		if field.DataType == entity.FieldTypeFloatVector || field.DataType == entity.FieldTypeBinaryVector {
 			vectors++
 		}
@@ -158,6 +220,14 @@ func validateSchema(sch *entity.Schema) error {
 	if vectors <= 0 {
 		return errors.New("vector field not set")
 	}
+	switch {
+	case hasJSON && c.config.hasFlags(disableJSON):
+		return ErrFeatureNotSupported
+	case hasDynamicSchema && c.config.hasFlags(disableDynamicSchema):
+		return ErrFeatureNotSupported
+	case hasPartitionKey && c.config.hasFlags(disableParitionKey):
+		return ErrFeatureNotSupported
+	}
 	return nil
 }
 
@@ -356,22 +426,21 @@ func (c *GrpcClient) LoadCollection(ctx context.Context, collName string, async
 	}
 
 	if !async {
+		ticker := time.NewTicker(200 * time.Millisecond)
+		defer ticker.Stop()
 		for {
 			select {
 			case <-ctx.Done():
-				return errors.New("context deadline exceeded")
-			default:
-			}
-
-			coll, err := c.ShowCollection(ctx, collName)
-			if err != nil {
-				return err
+				return ctx.Err()
+			case <-ticker.C:
+				progress, err := c.getLoadingProgress(ctx, collName)
+				if err != nil {
+					return err
+				}
+				if progress == 100 {
+					return nil
+				}
 			}
-			if coll.Loaded {
-				break
-			}
-
-			time.Sleep(200 * time.Millisecond) // TODO change to configuration
 		}
 	}
 	return nil
@@ -521,3 +590,21 @@ func (c *GrpcClient) AlterCollection(ctx context.Context, collName string, attrs
 	}
 	return handleRespStatus(resp)
 }
+
+func (c *GrpcClient) getLoadingProgress(ctx context.Context, collectionName string, partitionNames ...string) (int64, error) {
+	req := &server.GetLoadingProgressRequest{
+		Base:           &common.MsgBase{},
+		DbName:         "",
+		CollectionName: collectionName,
+		PartitionNames: partitionNames,
+	}
+
+	resp, err := c.Service.GetLoadingProgress(ctx, req)
+	if err != nil {
+		return -1, err
+	}
+	if err := handleRespStatus(resp.GetStatus()); err != nil {
+		return -1, err
+	}
+	return resp.GetProgress(), nil
+}
diff --git a/client/config.go b/client/config.go
index 1fa967a..1dc1789 100644
--- a/client/config.go
+++ b/client/config.go
@@ -20,6 +20,13 @@ import (
 	"google.golang.org/grpc/keepalive"
 )
 
+const (
+	disableDatabase uint64 = 1 << iota
+	disableJSON
+	disableDynamicSchema
+	disableParitionKey
+)
+
 var regexValidScheme = regexp.MustCompile(`^https?:\/\/`)
 
 // DefaultGrpcOpts is GRPC options for milvus client.
@@ -50,12 +57,15 @@ type Config struct {
 	Identifier    string // Identifier for this connection
 	EnableTLSAuth bool   // Enable TLS Auth for transport security.
 	APIKey        string // API key
+	ServerVersion string // ServerVersion
 
 	DialOptions []grpc.DialOption // Dial options for GRPC.
 
 	parsedAddress *url.URL
 
 	DisableConn bool
+
+	flags uint64 // internal flags
 }
 
 // Copy a new config, dialOption may shared with old config.
@@ -147,3 +157,17 @@ func (c *Config) getDialOption() []grpc.DialOption {
 	))
 	return options
 }
+
+// addFlags set internal flags
+func (c *Config) addFlags(flags uint64) {
+	c.flags |= flags
+}
+
+// hasFlags check flags is set
+func (c *Config) hasFlags(flags uint64) bool {
+	return (c.flags & flags) > 0
+}
+
+func (c *Config) resetFlags(flags uint64) {
+	c.flags &= ^flags
+}
diff --git a/client/data.go b/client/data.go
index 4faa027..bc7e82b 100644
--- a/client/data.go
+++ b/client/data.go
@@ -150,6 +150,30 @@ func PKs2Expr(backName string, ids entity.Column) string {
 	return expr
 }
 
+// Get grabs the inserted entities using the primary key from the Collection.
+func (c *GrpcClient) Get(ctx context.Context, collectionName string, ids entity.Column, opts ...GetOption) (ResultSet, error) {
+	if c.Service == nil {
+		return nil, ErrClientNotReady
+	}
+
+	o := &getOption{}
+	for _, opt := range opts {
+		opt(o)
+	}
+
+	if len(o.outputFields) == 0 {
+		coll, err := c.DescribeCollection(ctx, collectionName)
+		if err != nil {
+			return nil, err
+		}
+		for _, f := range coll.Schema.Fields {
+			o.outputFields = append(o.outputFields, f.Name)
+		}
+	}
+
+	return c.QueryByPks(ctx, collectionName, o.partitionNames, ids, o.outputFields)
+}
+
 // QueryByPks query record by specified primary key(s)
 func (c *GrpcClient) QueryByPks(ctx context.Context, collectionName string, partitionNames []string, ids entity.Column, outputFields []string, opts ...SearchQueryOptionFunc) (ResultSet, error) {
 	if c.Service == nil {
@@ -250,6 +274,15 @@ func getPKField(schema *entity.Schema) *entity.Field {
 	return nil
 }
 
+func getVectorField(schema *entity.Schema) *entity.Field {
+	for _, f := range schema.Fields {
+		if f.DataType == entity.FieldTypeFloatVector || f.DataType == entity.FieldTypeBinaryVector {
+			return f
+		}
+	}
+	return nil
+}
+
 func prepareSearchRequest(collName string, partitions []string,
 	expr string, outputFields []string, vectors []entity.Vector, vectorField string,
 	metricType entity.MetricType, topK int, sp entity.SearchParam, opt *SearchQueryOption) (*server.SearchRequest, error) {
diff --git a/client/database.go b/client/database.go
index ea225c7..697009a 100644
--- a/client/database.go
+++ b/client/database.go
@@ -44,6 +44,9 @@ func (c *GrpcClient) CreateDatabase(ctx context.Context, dbName string) error {
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
+	if c.config.hasFlags(disableDatabase) {
+		return ErrFeatureNotSupported
+	}
 	req := &server.CreateDatabaseRequest{
 		DbName: dbName,
 	}
@@ -59,6 +62,9 @@ func (c *GrpcClient) ListDatabases(ctx context.Context) ([]entity.Database, erro
 	if c.Service == nil {
 		return nil, ErrClientNotReady
 	}
+	if c.config.hasFlags(disableDatabase) {
+		return nil, ErrFeatureNotSupported
+	}
 
 	req := &server.ListDatabasesRequest{}
 	resp, err := c.Service.ListDatabases(ctx, req)
@@ -82,6 +88,9 @@ func (c *GrpcClient) DropDatabase(ctx context.Context, dbName string) error {
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
+	if c.config.hasFlags(disableDatabase) {
+		return ErrFeatureNotSupported
+	}
 
 	req := &server.DropDatabaseRequest{
 		DbName: dbName,
diff --git a/client/errors.go b/client/errors.go
index b431e91..68fec4f 100644
--- a/client/errors.go
+++ b/client/errors.go
@@ -23,8 +23,12 @@ type ErrServiceFailed error
 var (
 	//ErrClientNotReady error indicates client not ready
 	ErrClientNotReady = errors.New("client not ready")
+
 	//ErrStatusNil error indicates response has nil status
 	ErrStatusNil = errors.New("response status is nil")
+
+	// ErrFeatureNotSupported error indicates current server does not support this feature
+	ErrFeatureNotSupported = errors.New("feature not supported")
 )
 
 // ErrCollectionNotExists indicates the collection with specified collection name does not exist
diff --git a/client/connect.go b/client/grpc_client.go
similarity index 56%
rename from client/connect.go
rename to client/grpc_client.go
index cc46a08..9e82088 100644
--- a/client/connect.go
+++ b/client/grpc_client.go
@@ -1,14 +1,3 @@
-// Copyright (C) 2019-2021 Zilliz. All rights reserved.
-//
-// Licensed 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 client
 
 import (
@@ -18,7 +7,6 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/cockroachdb/errors"
 	common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	server "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"google.golang.org/grpc"
@@ -26,14 +14,19 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-// GrpcClient, uses default grpc Service definition to connect with Milvus2.0
+// Check if GrpcClient implement Client.
+var _ Client = &GrpcClient{}
+
+// GrpcClient  uses default grpc Service definition to connect with Milvus2.0
 type GrpcClient struct {
 	Conn    *grpc.ClientConn           // grpc connection instance
 	Service server.MilvusServiceClient // Service client stub
-	config  *Config                    // No thread safety
+
+	config *Config // No thread safety
 }
 
-func (c *GrpcClient) dial(ctx context.Context, addr string, opts ...grpc.DialOption) error {
+// connect connect to Service
+func (c *GrpcClient) connect(ctx context.Context, addr string, opts ...grpc.DialOption) error {
 	if addr == "" {
 		return fmt.Errorf("address is empty")
 	}
@@ -41,22 +34,17 @@ func (c *GrpcClient) dial(ctx context.Context, addr string, opts ...grpc.DialOpt
 	if err != nil {
 		return err
 	}
-	c.Conn = conn
-	return nil
-}
 
-// connect connect to Service
-func (c *GrpcClient) connect(ctx context.Context, addr string, opts ...grpc.DialOption) error {
-	if err := c.dial(ctx, addr, opts...); err != nil {
-		return err
-	}
+	c.Conn = conn
 	c.Service = server.NewMilvusServiceClient(c.Conn)
 
 	if !c.config.DisableConn {
-		if err := c.connectInternal(ctx); err != nil {
+		err = c.connectInternal(ctx)
+		if err != nil {
 			return err
 		}
 	}
+
 	return nil
 }
 
@@ -68,8 +56,8 @@ func (c *GrpcClient) connectInternal(ctx context.Context) error {
 
 	req := &server.ConnectRequest{
 		ClientInfo: &common.ClientInfo{
-			SdkType:    "golang",
-			SdkVersion: sdkVerion,
+			SdkType:    "Golang",
+			SdkVersion: "v2.2.4",
 			LocalTime:  time.Now().String(),
 			User:       c.config.Username,
 			Host:       hostName,
@@ -81,17 +69,24 @@ func (c *GrpcClient) connectInternal(ctx context.Context) error {
 		status, ok := status.FromError(err)
 		if ok {
 			if status.Code() == codes.Unimplemented {
-				return errors.New("this version of sdk is incompatible with server," +
-					" please downgrade your sdk or upgrade your server")
+				// disable unsupported feature
+				c.config.addFlags(
+					disableDatabase |
+						disableJSON |
+						disableParitionKey |
+						disableDynamicSchema)
 			}
+			return nil
 		}
 		return err
 	}
-	if err := handleRespStatus(resp.GetStatus()); err != nil {
-		return err
+
+	if resp.GetStatus().GetErrorCode() != common.ErrorCode_Success {
+		return fmt.Errorf("connect fail, %s", resp.Status.Reason)
 	}
 
 	c.config.Identifier = strconv.FormatInt(resp.GetIdentifier(), 10)
+	c.config.ServerVersion = resp.GetServerInfo().GetBuildTags()
 	return nil
 }
 
diff --git a/client/interceptor.go b/client/interceptor.go
index 6ef9eeb..20cac99 100644
--- a/client/interceptor.go
+++ b/client/interceptor.go
@@ -12,6 +12,10 @@ import (
 
 const (
 	authorizationHeader = `authorization`
+
+	identifierHeader = `identifier`
+
+	databaseHeader = `dbname`
 )
 
 // authenticationInterceptor appends credential into context metadata
@@ -25,16 +29,16 @@ func authenticationInterceptor(ctx context.Context, username, password string) c
 
 func apiKeyInterceptor(ctx context.Context, apiKey string) context.Context {
 	if apiKey != "" {
-		value := crypto.Base64Encode(fmt.Sprintf("Bearer: %s", apiKey))
+		value := crypto.Base64Encode(apiKey)
 		return metadata.AppendToOutgoingContext(ctx, authorizationHeader, value)
 	}
 	return ctx
 }
 
 func identifierInterceptor(ctx context.Context, identifierGetter func() string) context.Context {
-	dbname := identifierGetter()
-	if dbname != "" {
-		ctx = metadata.AppendToOutgoingContext(ctx, "identifier", dbname)
+	identifier := identifierGetter()
+	if identifier != "" {
+		ctx = metadata.AppendToOutgoingContext(ctx, identifierHeader, identifier)
 	}
 	return ctx
 }
@@ -43,7 +47,7 @@ func identifierInterceptor(ctx context.Context, identifierGetter func() string)
 func databaseNameInterceptor(ctx context.Context, dbNameGetter func() string) context.Context {
 	dbname := dbNameGetter()
 	if dbname != "" {
-		ctx = metadata.AppendToOutgoingContext(ctx, "dbname", dbname)
+		ctx = metadata.AppendToOutgoingContext(ctx, databaseHeader, dbname)
 	}
 	return ctx
 }
diff --git a/client/options.go b/client/options.go
index db85bdc..dcd1b4a 100644
--- a/client/options.go
+++ b/client/options.go
@@ -16,31 +16,88 @@ import (
 
 	"github.com/cockroachdb/errors"
 
-	common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	server "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 )
 
-// CreateCollectionOption is an option that is used to modify CreateCollectionRequest
-type CreateCollectionOption func(*server.CreateCollectionRequest)
+// CreateCollectionOption is an option that is used to alter create collection options.
+type CreateCollectionOption func(opt *createCollOpt)
+
+type createCollOpt struct {
+	ConsistencyLevel    entity.ConsistencyLevel
+	NumPartitions       int64
+	PrimaryKeyFieldName string
+	PrimaryKeyFieldType entity.FieldType
+	PrimaryKeyMaxLength int64
+	VectorFieldName     string
+	MetricsType         entity.MetricType
+	AutoID              bool
+	EnableDynamicSchema bool
+	Properties          map[string]string
+}
 
-// WithConsistencyLevel specifies a specific ConsistencyLevel, rather than using the default ReaderProperties.
-func WithConsistencyLevel(cl entity.ConsistencyLevel) CreateCollectionOption {
-	return func(req *server.CreateCollectionRequest) {
-		req.ConsistencyLevel = cl.CommonConsistencyLevel()
+func WithPKFieldName(name string) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		opt.PrimaryKeyFieldName = name
+	}
+}
+
+func WithPKFieldType(tp entity.FieldType) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		opt.PrimaryKeyFieldType = tp
+	}
+}
+
+func WithPKMaxLength(maxLength int64) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		opt.PrimaryKeyMaxLength = maxLength
+	}
+}
+
+func WithVectorFieldName(name string) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		opt.VectorFieldName = name
+	}
+}
+
+func WithMetricsType(mt entity.MetricType) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		opt.MetricsType = mt
+	}
+}
+
+func WithAutoID(autoID bool) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		opt.AutoID = autoID
 	}
 }
 
-func WithCollectionProperty(key string, value string) CreateCollectionOption {
-	return func(req *server.CreateCollectionRequest) {
-		req.Properties = append(req.Properties, &common.KeyValuePair{Key: key, Value: value})
+func WithEnableDynamicSchema(enable bool) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		opt.EnableDynamicSchema = enable
+	}
+}
+
+// WithConsistencyLevel specifies a specific ConsistencyLevel, rather than using the default ReaderProperties.
+func WithConsistencyLevel(cl entity.ConsistencyLevel) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		opt.ConsistencyLevel = cl
 	}
 }
 
 // WithPartitionNum returns a create collection options to set physical partition number when logical partition feature.
 func WithPartitionNum(partitionNums int64) CreateCollectionOption {
-	return func(req *server.CreateCollectionRequest) {
-		req.NumPartitions = partitionNums
+	return func(opt *createCollOpt) {
+		opt.NumPartitions = partitionNums
+	}
+}
+
+func WithCollectionProperty(key, value string) CreateCollectionOption {
+	return func(opt *createCollOpt) {
+		if opt.Properties == nil {
+			opt.Properties = make(map[string]string)
+		}
+		opt.Properties[key] = value
 	}
 }
 
@@ -173,3 +230,22 @@ func WithEndTs(endTs int64) BulkInsertOption {
 		req.Options = entity.MapKvPairs(optionMap)
 	}
 }
+
+type getOption struct {
+	partitionNames []string
+	outputFields   []string
+}
+
+type GetOption func(o *getOption)
+
+func GetWithPartitions(partionNames ...string) GetOption {
+	return func(o *getOption) {
+		o.partitionNames = partionNames
+	}
+}
+
+func GetWithOutputFields(outputFields ...string) GetOption {
+	return func(o *getOption) {
+		o.outputFields = outputFields
+	}
+}
diff --git a/client/partition.go b/client/partition.go
index d9cd338..36c3199 100644
--- a/client/partition.go
+++ b/client/partition.go
@@ -149,23 +149,6 @@ func (c *GrpcClient) LoadPartitions(ctx context.Context, collName string, partit
 			return err
 		}
 	}
-	partitions, err := c.ShowPartitions(ctx, collName)
-	if err != nil {
-		return err
-	}
-	m := make(map[string]int64)
-	for _, partition := range partitions {
-		m[partition.Name] = partition.ID
-	}
-	// load partitions ids
-	ids := make(map[int64]struct{})
-	for _, partitionName := range partitionNames {
-		id, has := m[partitionName]
-		if !has {
-			return fmt.Errorf("collection %s does not has partitions %s", collName, partitionName)
-		}
-		ids[id] = struct{}{}
-	}
 
 	req := &server.LoadPartitionsRequest{
 		DbName:         "", // reserved
@@ -181,34 +164,22 @@ func (c *GrpcClient) LoadPartitions(ctx context.Context, collName string, partit
 	}
 
 	if !async {
+		ticker := time.NewTicker(200 * time.Millisecond)
+		defer ticker.Stop()
+
 		for {
 			select {
 			case <-ctx.Done():
-				return errors.New("context deadline exceeded")
-			default:
-			}
-			partitions, err := c.ShowPartitions(ctx, collName)
-			if err != nil {
-				return err
-			}
-			foundLoading := false
-			loaded := 0
-			for _, partition := range partitions {
-				if _, has := ids[partition.ID]; !has {
-					continue
+				return ctx.Err()
+			case <-ticker.C:
+				progress, err := c.getLoadingProgress(ctx, collName, partitionNames...)
+				if err != nil {
+					return err
 				}
-				if !partition.Loaded {
-					//Not loaded
-					foundLoading = true
-					break
+				if progress == 100 {
+					return nil
 				}
-				loaded++
-			}
-			if foundLoading || loaded < len(partitionNames) {
-				time.Sleep(time.Millisecond * 100)
-				continue
 			}
-			break
 		}
 	}
 
diff --git a/client/results.go b/client/results.go
new file mode 100644
index 0000000..51c7c52
--- /dev/null
+++ b/client/results.go
@@ -0,0 +1,28 @@
+package client
+
+import "github.com/milvus-io/milvus-sdk-go/v2/entity"
+
+// SearchResult contains the result from Search api of client
+// IDs is the auto generated id values for the entities
+// Fields contains the data of `outputFieleds` specified or all columns if non
+// Scores is actually the distance between the vector current record contains and the search target vector
+type SearchResult struct {
+	ResultCount int           // the returning entry count
+	IDs         entity.Column // auto generated id, can be mapped to the columns from `Insert` API
+	Fields      ResultSet     //[]entity.Column // output field data
+	Scores      []float32     // distance to the target vector
+	Err         error         // search error if any
+}
+
+// ResultSet is an alias type for column slice.
+type ResultSet []entity.Column
+
+// GetColumn returns column with provided field name.
+func (rs ResultSet) GetColumn(fieldName string) entity.Column {
+	for _, column := range rs {
+		if column.Name() == fieldName {
+			return column
+		}
+	}
+	return nil
+}
diff --git a/entity/collection.go b/entity/collection.go
index 345f9c3..c9ad95a 100644
--- a/entity/collection.go
+++ b/entity/collection.go
@@ -15,6 +15,9 @@ package entity
 // DefaultShardNumber const value for using Milvus default shard number.
 const DefaultShardNumber int32 = 0
 
+// DefaultConsistencyLevel const value for using Milvus default consistency level setting.
+const DefaultConsistencyLevel ConsistencyLevel = ClBounded
+
 // Collection represent collection meta in Milvus
 type Collection struct {
 	ID               int64   // collection id


Test patch:
diff --git a/client/collection_test.go b/client/collection_test.go
index 74eef01..0c4f795 100644
--- a/client/collection_test.go
+++ b/client/collection_test.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"math/rand"
 	"testing"
-	"time"
 
 	"github.com/cockroachdb/errors"
 
@@ -205,38 +204,7 @@ func (s *CollectionSuite) TestCreateCollection() {
 		}
 	})
 
-	s.Run("collection_already_exists", func() {
-		defer s.resetMock()
-
-		ds := defaultSchema()
-		shardsNum := int32(1)
-
-		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
-			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
-
-		err := c.CreateCollection(ctx, ds, shardsNum)
-		s.Error(err)
-	})
-
 	s.Run("server_returns_error", func() {
-		s.Run("has_collection_error", func() {
-			defer s.resetMock()
-			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
-				Return(nil, errors.New("mocked grpc error"))
-
-			err := c.CreateCollection(ctx, defaultSchema(), 1)
-			s.Error(err)
-		})
-
-		s.Run("has_collection_fail", func() {
-			defer s.resetMock()
-			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
-				Return(&server.BoolResponse{Status: &common.Status{ErrorCode: common.ErrorCode_UnexpectedError}}, nil)
-
-			err := c.CreateCollection(ctx, defaultSchema(), 1)
-			s.Error(err)
-		})
-
 		s.Run("create_collection_error", func() {
 			defer s.resetMock()
 			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).Return(&server.BoolResponse{Status: &common.Status{}, Value: false}, nil)
@@ -257,6 +225,162 @@ func (s *CollectionSuite) TestCreateCollection() {
 			s.Error(err)
 		})
 	})
+
+	s.Run("feature_not_support", func() {
+		cases := []struct {
+			tag  string
+			flag uint64
+		}{
+			{tag: "json", flag: disableJSON},
+			{tag: "partition_key", flag: disableParitionKey},
+			{tag: "dyanmic_schema", flag: disableDynamicSchema},
+		}
+		sch := entity.NewSchema().WithName("all_feature").WithDynamicFieldEnabled(true).
+			WithField(entity.NewField().WithName("id").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
+			WithField(entity.NewField().WithName("embedding").WithDataType(entity.FieldTypeFloatVector).WithDim(128)).
+			WithField(entity.NewField().WithName("partition").WithDataType(entity.FieldTypeInt64).WithIsPartitionKey(true)).
+			WithField(entity.NewField().WithName("dynamic").WithDataType(entity.FieldTypeJSON).WithIsDynamic(true))
+		for _, tc := range cases {
+			s.Run(tc.tag, func() {
+				grpcClient, ok := c.(*GrpcClient)
+				s.Require().True(ok)
+				grpcClient.config.addFlags(tc.flag)
+				defer grpcClient.config.resetFlags(tc.flag)
+
+				err := c.CreateCollection(ctx, sch, 1)
+				s.Error(err)
+				s.ErrorIs(err, ErrFeatureNotSupported)
+			})
+		}
+	})
+}
+
+func (s *CollectionSuite) TestNewCollection() {
+	c := s.client
+	ctx, cancel := context.WithCancel(context.Background())
+
+	defer cancel()
+	s.resetMock()
+
+	s.Run("all_default", func() {
+		defer s.resetMock()
+
+		created := false
+		s.mock.EXPECT().CreateCollection(mock.Anything, mock.AnythingOfType("*milvuspb.CreateCollectionRequest")).
+			Run(func(ctx context.Context, req *server.CreateCollectionRequest) {
+				s.Equal(testCollectionName, req.GetCollectionName())
+				sschema := &schema.CollectionSchema{}
+				s.Require().NoError(proto.Unmarshal(req.GetSchema(), sschema))
+				s.Require().Equal(2, len(sschema.Fields))
+				for _, field := range sschema.Fields {
+					if field.GetName() == "id" {
+						s.Equal(schema.DataType_Int64, field.GetDataType())
+					}
+					if field.GetName() == "vector" {
+						s.Equal(schema.DataType_FloatVector, field.GetDataType())
+					}
+				}
+
+				s.Equal(entity.DefaultShardNumber, req.GetShardsNum())
+				s.Equal(entity.DefaultConsistencyLevel.CommonConsistencyLevel(), req.GetConsistencyLevel())
+				created = true
+			}).
+			Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).Call.Return(func(_ context.Context, _ *server.HasCollectionRequest) *server.BoolResponse {
+			return &server.BoolResponse{Status: &common.Status{}, Value: created}
+		}, nil)
+		s.mock.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().Flush(mock.Anything, mock.Anything).Return(&server.FlushResponse{
+			Status:     &common.Status{ErrorCode: common.ErrorCode_Success},
+			CollSegIDs: map[string]*schema.LongArray{},
+		}, nil)
+		s.mock.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&server.DescribeIndexResponse{
+			Status: &common.Status{ErrorCode: common.ErrorCode_Success},
+			IndexDescriptions: []*server.IndexDescription{
+				{FieldName: "vector", State: common.IndexState_Finished},
+			},
+		}, nil)
+		s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).Return(&server.GetLoadingProgressResponse{
+			Status:   &common.Status{ErrorCode: common.ErrorCode_Success},
+			Progress: 100,
+		}, nil)
+		s.mock.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&server.DescribeCollectionResponse{
+			Status: &common.Status{ErrorCode: common.ErrorCode_Success},
+			Schema: &schema.CollectionSchema{
+				Fields: []*schema.FieldSchema{
+					{Name: "id", DataType: schema.DataType_VarChar},
+					{Name: "vector", DataType: schema.DataType_FloatVector},
+				},
+			},
+		}, nil)
+
+		err := c.NewCollection(ctx, testCollectionName, testVectorDim)
+		s.NoError(err)
+	})
+
+	s.Run("with_custom_set", func() {
+		defer s.resetMock()
+		created := false
+		s.mock.EXPECT().CreateCollection(mock.Anything, mock.AnythingOfType("*milvuspb.CreateCollectionRequest")).
+			Run(func(ctx context.Context, req *server.CreateCollectionRequest) {
+				s.Equal(testCollectionName, req.GetCollectionName())
+				sschema := &schema.CollectionSchema{}
+				s.Require().NoError(proto.Unmarshal(req.GetSchema(), sschema))
+				s.Require().Equal(2, len(sschema.Fields))
+				for _, field := range sschema.Fields {
+					if field.GetName() == "my_pk" {
+						s.Equal(schema.DataType_VarChar, field.GetDataType())
+					}
+					if field.GetName() == "embedding" {
+						s.Equal(schema.DataType_FloatVector, field.GetDataType())
+					}
+				}
+
+				s.Equal(entity.DefaultShardNumber, req.GetShardsNum())
+				s.Equal(entity.ClEventually.CommonConsistencyLevel(), req.GetConsistencyLevel())
+				created = true
+			}).
+			Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).Call.Return(func(_ context.Context, _ *server.HasCollectionRequest) *server.BoolResponse {
+			return &server.BoolResponse{Status: &common.Status{}, Value: created}
+		}, nil)
+		s.mock.EXPECT().CreateIndex(mock.Anything, mock.Anything).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().Flush(mock.Anything, mock.Anything).Return(&server.FlushResponse{
+			Status:     &common.Status{ErrorCode: common.ErrorCode_Success},
+			CollSegIDs: map[string]*schema.LongArray{},
+		}, nil)
+		s.mock.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(&server.DescribeIndexResponse{
+			Status: &common.Status{ErrorCode: common.ErrorCode_Success},
+			IndexDescriptions: []*server.IndexDescription{
+				{FieldName: "embedding", State: common.IndexState_Finished},
+			},
+		}, nil)
+		s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).Return(&server.GetLoadingProgressResponse{
+			Status:   &common.Status{ErrorCode: common.ErrorCode_Success},
+			Progress: 100,
+		}, nil)
+		s.mock.EXPECT().DescribeCollection(mock.Anything, mock.Anything).Return(&server.DescribeCollectionResponse{
+			Status: &common.Status{ErrorCode: common.ErrorCode_Success},
+			Schema: &schema.CollectionSchema{
+				Fields: []*schema.FieldSchema{
+					{Name: "my_pk", DataType: schema.DataType_VarChar},
+					{Name: "embedding", DataType: schema.DataType_FloatVector},
+				},
+			},
+		}, nil)
+
+		err := c.NewCollection(ctx, testCollectionName, testVectorDim, WithPKFieldName("my_pk"), WithPKFieldType(entity.FieldTypeVarChar), WithVectorFieldName("embedding"), WithConsistencyLevel(entity.ClEventually))
+		s.NoError(err)
+	})
+
+	s.Run("varchar_autoid", func() {
+		defer s.resetMock()
+
+		err := c.NewCollection(ctx, testCollectionName, testVectorDim, WithPKFieldType(entity.FieldTypeVarChar), WithAutoID(true))
+		s.Error(err)
+	})
 }
 
 func (s *CollectionSuite) TestRenameCollection() {
@@ -377,6 +501,132 @@ func (s *CollectionSuite) TestAlterCollection() {
 	})
 }
 
+func (s *CollectionSuite) TestLoadCollection() {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	c := s.client
+
+	s.Run("normal_run_async", func() {
+		defer s.resetMock()
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+
+		s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+
+		err := c.LoadCollection(ctx, testCollectionName, true)
+		s.NoError(err)
+	})
+
+	s.Run("normal_run_sync", func() {
+		defer s.resetMock()
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+
+		s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).
+			Return(&server.GetLoadingProgressResponse{
+				Status:   &common.Status{ErrorCode: common.ErrorCode_Success},
+				Progress: 100,
+			}, nil)
+
+		err := c.LoadCollection(ctx, testCollectionName, true)
+		s.NoError(err)
+	})
+
+	s.Run("load_default_replica", func() {
+		defer s.resetMock()
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+
+		s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.LoadCollectionRequest) {
+			s.Equal(testDefaultReplicaNumber, req.GetReplicaNumber())
+			s.Equal(testCollectionName, req.GetCollectionName())
+		}).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+
+		err := c.LoadCollection(ctx, testCollectionName, true)
+		s.NoError(err)
+	})
+
+	s.Run("load_multiple_replica", func() {
+		defer s.resetMock()
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+
+		s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.LoadCollectionRequest) {
+			s.Equal(testMultiReplicaNumber, req.GetReplicaNumber())
+			s.Equal(testCollectionName, req.GetCollectionName())
+		}).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+
+		err := c.LoadCollection(ctx, testCollectionName, true, WithReplicaNumber(testMultiReplicaNumber))
+		s.NoError(err)
+	})
+
+	s.Run("has_collection_failure", func() {
+		s.Run("return_false", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: false}, nil)
+
+			err := c.LoadCollection(ctx, testCollectionName, true)
+			s.Error(err)
+		})
+
+		s.Run("return_error", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(nil, errors.New("mock error"))
+
+			err := c.LoadCollection(ctx, testCollectionName, true)
+			s.Error(err)
+		})
+	})
+
+	s.Run("load_collection_failure", func() {
+		s.Run("failure_status", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+
+			s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).
+				Return(&common.Status{ErrorCode: common.ErrorCode_UnexpectedError}, nil)
+
+			err := c.LoadCollection(ctx, testCollectionName, true)
+			s.Error(err)
+		})
+
+		s.Run("return_error", func() {
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+
+			s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).
+				Return(nil, errors.New("mock error"))
+
+			err := c.LoadCollection(ctx, testCollectionName, true)
+			s.Error(err)
+		})
+	})
+
+	s.Run("get_loading_progress_failure", func() {
+		defer s.resetMock()
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+
+		s.mock.EXPECT().LoadCollection(mock.Anything, mock.Anything).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).
+			Return(nil, errors.New("mock error"))
+
+		err := c.LoadCollection(ctx, testCollectionName, false)
+		s.Error(err)
+	})
+
+	s.Run("service_not_ready", func() {
+		c := &GrpcClient{}
+		err := c.LoadCollection(ctx, testCollectionName, false)
+		s.ErrorIs(err, ErrClientNotReady)
+	})
+}
+
 func TestCollectionSuite(t *testing.T) {
 	suite.Run(t, new(CollectionSuite))
 }
@@ -421,90 +671,6 @@ func TestGrpcClientDropCollection(t *testing.T) {
 	})
 }
 
-func TestGrpcClientLoadCollection(t *testing.T) {
-	ctx := context.Background()
-	c := testClient(ctx, t)
-	mockServer.SetInjection(MHasCollection, hasCollectionDefault)
-	// injection check collection name equals
-	mockServer.SetInjection(MLoadCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
-		req, ok := raw.(*server.LoadCollectionRequest)
-		if !ok {
-			return BadRequestStatus()
-		}
-		assert.Equal(t, testCollectionName, req.GetCollectionName())
-		return SuccessStatus()
-	})
-	t.Run("Load collection normal async", func(t *testing.T) {
-		assert.Nil(t, c.LoadCollection(ctx, testCollectionName, true))
-	})
-	t.Run("Load collection sync", func(t *testing.T) {
-
-		loadTime := rand.Intn(500) + 500 // in milli seconds, 100~1000 milliseconds
-		passed := false                  // ### flag variable
-		start := time.Now()
-
-		mockServer.SetInjection(MShowCollections, func(_ context.Context, raw proto.Message) (proto.Message, error) {
-			req, ok := raw.(*server.ShowCollectionsRequest)
-			r := &server.ShowCollectionsResponse{}
-			if !ok || req == nil {
-				s, err := BadRequestStatus()
-				r.Status = s
-				return r, err
-			}
-			s, err := SuccessStatus()
-			r.Status = s
-			r.CollectionIds = []int64{1}
-			var perc int64
-			if time.Since(start) > time.Duration(loadTime)*time.Millisecond {
-				t.Log("passed")
-				perc = 100
-				passed = true
-			}
-			r.InMemoryPercentages = []int64{perc}
-			return r, err
-		})
-		assert.Nil(t, c.LoadCollection(ctx, testCollectionName, false))
-		assert.True(t, passed)
-
-		start = time.Now()
-		passed = false
-		quickCtx, cancel := context.WithTimeout(ctx, 50*time.Millisecond)
-		defer cancel()
-		assert.NotNil(t, c.LoadCollection(quickCtx, testCollectionName, false))
-
-		// remove injection
-		mockServer.DelInjection(MShowCollections)
-	})
-	t.Run("Load default replica", func(t *testing.T) {
-		mockServer.SetInjection(MLoadCollection, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
-			req, ok := raw.(*server.LoadCollectionRequest)
-			if !ok {
-				return BadRequestStatus()
-			}
-			assert.Equal(t, testDefaultReplicaNumber, req.GetReplicaNumber())
-			assert.Equal(t, testCollectionName, req.GetCollectionName())
-			return SuccessStatus()
-		})
-		defer mockServer.DelInjection(MLoadCollection)
-		assert.Nil(t, c.LoadCollection(ctx, testCollectionName, true))
-	})
-	t.Run("Load multiple replica", func(t *testing.T) {
-		mockServer.DelInjection(MLoadCollection)
-
-		mockServer.SetInjection(MLoadCollection, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
-			req, ok := raw.(*server.LoadCollectionRequest)
-			if !ok {
-				return BadRequestStatus()
-			}
-			assert.Equal(t, testMultiReplicaNumber, req.GetReplicaNumber())
-			assert.Equal(t, testCollectionName, req.GetCollectionName())
-			return SuccessStatus()
-		})
-		defer mockServer.DelInjection(MLoadCollection)
-		assert.Nil(t, c.LoadCollection(ctx, testCollectionName, true, WithReplicaNumber(testMultiReplicaNumber)))
-	})
-}
-
 func TestReleaseCollection(t *testing.T) {
 	ctx := context.Background()
 
@@ -764,7 +930,6 @@ func TestGrpcClientGetLoadingProgress(t *testing.T) {
 	progress, err := c.GetLoadingProgress(ctx, testCollectionName, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, int64(100), progress)
-
 }
 
 func TestGrpcClientGetLoadState(t *testing.T) {
@@ -793,5 +958,4 @@ func TestGrpcClientGetLoadState(t *testing.T) {
 	state, err := c.GetLoadState(ctx, testCollectionName, []string{})
 	assert.NoError(t, err)
 	assert.Equal(t, entity.LoadStateLoaded, state)
-
 }
diff --git a/client/data_test.go b/client/data_test.go
index 38475dd..ee69729 100644
--- a/client/data_test.go
+++ b/client/data_test.go
@@ -489,6 +489,32 @@ func (s *QuerySuite) SetupSuite() {
 		WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(testVectorDim))
 }
 
+func (s *QuerySuite) GetFail() {
+	c := s.client
+	ctx, cancel := context.WithCancel((context.Background()))
+	defer cancel()
+
+	idCol := entity.NewColumnInt64("ID", []int64{1})
+	s.Run("service_not_ready", func() {
+		_, err := (&GrpcClient{}).Get(ctx, testCollectionName, idCol)
+		s.Error(err)
+		s.ErrorIs(err, ErrClientNotReady)
+	})
+
+	s.Run("ids_len_0", func() {
+		_, err := c.Get(ctx, testCollectionName, entity.NewColumnInt64("ID", []int64{}), GetWithOutputFields("ID"))
+		s.Error(err)
+	})
+
+	s.Run("describe_failed", func() {
+		defer s.resetMock()
+		s.setupDescribeCollectionError(common.ErrorCode_Success, errors.New("mock error"))
+
+		_, err := c.Get(ctx, testCollectionName, idCol)
+		s.Error(err)
+	})
+}
+
 func (s *QuerySuite) TestQueryByPksFail() {
 	c := s.client
 	ctx, cancel := context.WithCancel(context.Background())
diff --git a/client/interceptor_test.go b/client/interceptor_test.go
index cdd9379..b4b9d88 100644
--- a/client/interceptor_test.go
+++ b/client/interceptor_test.go
@@ -37,7 +37,7 @@ func TestAPIInterceptor(t *testing.T) {
 	ctx := context.Background()
 	ctx = apiKeyInterceptor(ctx, "test-token")
 	md, ok := metadata.FromOutgoingContext(ctx)
-	value := crypto.Base64Encode(fmt.Sprintf("Bearer: %s", "test-token"))
+	value := crypto.Base64Encode("test-token")
 	assert.True(t, ok)
 	assert.Equal(t, 1, len(md[authorizationHeader]))
 	assert.Equal(t, value, md[authorizationHeader][0])
diff --git a/client/options_test.go b/client/options_test.go
index a70d04c..25a418f 100644
--- a/client/options_test.go
+++ b/client/options_test.go
@@ -15,7 +15,6 @@ import (
 	"math/rand"
 	"testing"
 
-	common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	server "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 	"github.com/stretchr/testify/assert"
@@ -24,38 +23,38 @@ import (
 func TestCreateCollectionWithConsistencyLevel(t *testing.T) {
 	opt := WithConsistencyLevel(entity.ConsistencyLevel(entity.ClBounded))
 	assert.NotNil(t, opt)
-	req := &server.CreateCollectionRequest{}
+	o := &createCollOpt{}
 
 	assert.NotPanics(t, func() {
-		opt(req)
+		opt(o)
 	})
 
-	assert.Equal(t, common.ConsistencyLevel_Bounded, req.GetConsistencyLevel())
+	assert.Equal(t, entity.ClBounded, o.ConsistencyLevel)
 }
 
-func TestLoadCollectionWithReplicaNumber(t *testing.T) {
-	opt := WithReplicaNumber(testMultiReplicaNumber)
+func TestCreateCollectionWithPartitionNum(t *testing.T) {
+	partitionNum := rand.Int63n(1000) + 1
+	opt := WithPartitionNum(partitionNum)
 	assert.NotNil(t, opt)
-	req := &server.LoadCollectionRequest{}
+	o := &createCollOpt{}
 
 	assert.NotPanics(t, func() {
-		opt(req)
+		opt(o)
 	})
 
-	assert.Equal(t, testMultiReplicaNumber, req.GetReplicaNumber())
+	assert.Equal(t, partitionNum, o.NumPartitions)
 }
 
-func TestCreateCollectionWithPartitionNum(t *testing.T) {
-	partitionNum := rand.Int63n(1000) + 1
-	opt := WithPartitionNum(partitionNum)
+func TestLoadCollectionWithReplicaNumber(t *testing.T) {
+	opt := WithReplicaNumber(testMultiReplicaNumber)
 	assert.NotNil(t, opt)
-	req := &server.CreateCollectionRequest{}
+	req := &server.LoadCollectionRequest{}
 
 	assert.NotPanics(t, func() {
 		opt(req)
 	})
 
-	assert.Equal(t, partitionNum, req.GetNumPartitions())
+	assert.Equal(t, testMultiReplicaNumber, req.GetReplicaNumber())
 }
 
 func TestMakeSearchQueryOption(t *testing.T) {
diff --git a/client/partition_test.go b/client/partition_test.go
index 5502308..b89a45e 100644
--- a/client/partition_test.go
+++ b/client/partition_test.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"math/rand"
 	"testing"
-	"time"
 
 	"github.com/cockroachdb/errors"
 
@@ -14,6 +13,8 @@ import (
 	server "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/suite"
 )
 
 // return injection asserts collection name matchs
@@ -167,168 +168,6 @@ func TestGrpcClientShowPartitions(t *testing.T) {
 	}
 }
 
-func TestGrpcClientLoadPartitions(t *testing.T) {
-	ctx := context.Background()
-	c := testClient(ctx, t)
-
-	mockServer.SetInjection(MHasCollection, hasCollectionDefault)
-	mockServer.SetInjection(MHasPartition, hasPartitionInjection(t, testCollectionName, true, "_part1", "_part2", "_part3", "_part4"))
-
-	type testCase struct {
-		collName      string
-		partitions    []*entity.Partition
-		shouldSuccess bool
-		loadNames     []string
-	}
-	cases := []testCase{
-		{
-			collName: testCollectionName,
-			partitions: []*entity.Partition{
-				{
-					ID:   1,
-					Name: "_part1",
-				},
-				{
-					ID:   2,
-					Name: "_part2",
-				},
-				{
-					ID:   3,
-					Name: "_part3",
-				},
-			},
-			loadNames:     []string{"_part1", "_part2"},
-			shouldSuccess: true,
-		},
-		{
-			collName: testCollectionName,
-			partitions: []*entity.Partition{
-				{
-					ID:   1,
-					Name: "_part1",
-				},
-				{
-					ID:   2,
-					Name: "_part2",
-				},
-				{
-					ID:   3,
-					Name: "_part3",
-				},
-			},
-			loadNames:     []string{"_part4", "_part1"},
-			shouldSuccess: false,
-		},
-	}
-	for _, tc := range cases {
-		// one segment per paritions
-		start := time.Now()
-		loadTime := rand.Intn(1500) + 100
-		loaded := false
-		mockServer.SetInjection(MShowPartitions, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
-			req, ok := raw.(*server.ShowPartitionsRequest)
-			resp := &server.ShowPartitionsResponse{}
-			if !ok {
-				s, err := BadRequestStatus()
-				resp.Status = s
-				return resp, err
-			}
-			assert.Equal(t, testCollectionName, req.GetCollectionName())
-			resp.PartitionIDs = make([]int64, 0, len(tc.loadNames))
-			resp.PartitionNames = make([]string, 0, len(tc.loadNames))
-			var perc int64
-			if time.Since(start) > time.Duration(loadTime*int(time.Millisecond)) {
-				perc = 100
-				loaded = true
-			}
-			for _, part := range tc.partitions {
-				resp.PartitionIDs = append(resp.PartitionIDs, part.ID)
-				resp.PartitionNames = append(resp.PartitionNames, part.Name)
-				resp.InMemoryPercentages = append(resp.InMemoryPercentages, perc)
-			}
-			s, err := SuccessStatus()
-			resp.Status = s
-			return resp, err
-		})
-
-		err := c.LoadPartitions(ctx, tc.collName, tc.loadNames, false)
-		if tc.shouldSuccess {
-			assert.Nil(t, err)
-			assert.True(t, loaded)
-		} else {
-			assert.NotNil(t, err)
-		}
-	}
-
-	t.Run("some Service failed", func(t *testing.T) {
-		tc := &testCase{
-			collName: testCollectionName,
-			partitions: []*entity.Partition{
-				{
-					ID:   1,
-					Name: "_part1",
-				},
-				{
-					ID:   2,
-					Name: "_part2",
-				},
-				{
-					ID:   3,
-					Name: "_part3",
-				},
-			},
-			loadNames: []string{"_part1", "_part2"},
-		}
-
-		// paritions will not be loaded
-		mockServer.SetInjection(MShowPartitions, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
-			req, ok := raw.(*server.ShowPartitionsRequest)
-			resp := &server.ShowPartitionsResponse{}
-			if !ok {
-				s, err := BadRequestStatus()
-				resp.Status = s
-				return resp, err
-			}
-			assert.Equal(t, testCollectionName, req.GetCollectionName())
-			resp.PartitionIDs = make([]int64, 0, len(tc.loadNames))
-			resp.PartitionNames = make([]string, 0, len(tc.loadNames))
-			var perc int64
-			for _, part := range tc.partitions {
-				resp.PartitionIDs = append(resp.PartitionIDs, part.ID)
-				resp.PartitionNames = append(resp.PartitionNames, part.Name)
-				resp.InMemoryPercentages = append(resp.InMemoryPercentages, perc)
-			}
-			s, err := SuccessStatus()
-			resp.Status = s
-			return resp, err
-		})
-		quickCtx, cancel := context.WithTimeout(ctx, 10*time.Millisecond)
-		defer cancel()
-		assert.NotNil(t, c.LoadPartitions(quickCtx, tc.collName, tc.loadNames, false))
-
-		mockServer.SetInjection(MShowPartitions, func(_ context.Context, raw proto.Message) (proto.Message, error) {
-			return &server.ShowPartitionsResponse{}, errors.New("always fail")
-		})
-		defer mockServer.DelInjection(MShowPartitions)
-
-		err := c.LoadPartitions(ctx, tc.collName, tc.loadNames, false)
-		assert.NotNil(t, err)
-
-		mockServer.SetInjection(MLoadPartitions, func(_ context.Context, raw proto.Message) (proto.Message, error) {
-			return &common.Status{}, errors.New("has partition failed")
-		})
-		err = c.LoadPartitions(ctx, tc.collName, tc.loadNames, false)
-		assert.NotNil(t, err)
-
-		mockServer.SetInjection(MHasPartition, func(_ context.Context, raw proto.Message) (proto.Message, error) {
-			return &server.BoolResponse{}, errors.New("has partition failed")
-		})
-		err = c.LoadPartitions(ctx, tc.collName, tc.loadNames, false)
-		assert.NotNil(t, err)
-
-	})
-}
-
 func TestGrpcClientReleasePartitions(t *testing.T) {
 	ctx := context.Background()
 
@@ -411,3 +250,166 @@ func TestGrpcShowPartitions(t *testing.T) {
 		assert.Error(t, err)
 	})
 }
+
+type PartitionSuite struct {
+	MockSuiteBase
+}
+
+func (s *PartitionSuite) TestLoadPartitions() {
+	c := s.client
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	partNames := []string{"part_1", "part_2"}
+	mPartNames := map[string]struct{}{"part_1": {}, "part_2": {}}
+
+	s.Run("normal_run_async", func() {
+		defer s.resetMock()
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+		s.mock.EXPECT().HasPartition(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.HasPartitionRequest) {
+			s.Equal(testCollectionName, req.GetCollectionName())
+			_, ok := mPartNames[req.GetPartitionName()]
+			s.True(ok)
+		}).Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+		s.mock.EXPECT().LoadPartitions(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.LoadPartitionsRequest) {
+			s.Equal(testCollectionName, req.GetCollectionName())
+			s.ElementsMatch(partNames, req.GetPartitionNames())
+		}).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+
+		err := c.LoadPartitions(ctx, testCollectionName, partNames, true)
+		s.NoError(err)
+	})
+
+	s.Run("normal_run_sync", func() {
+		defer s.resetMock()
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+		s.mock.EXPECT().HasPartition(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.HasPartitionRequest) {
+			s.Equal(testCollectionName, req.GetCollectionName())
+			_, ok := mPartNames[req.GetPartitionName()]
+			s.True(ok)
+		}).Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+		s.mock.EXPECT().LoadPartitions(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.LoadPartitionsRequest) {
+			s.Equal(testCollectionName, req.GetCollectionName())
+			s.ElementsMatch(partNames, req.GetPartitionNames())
+		}).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).
+			Return(&server.GetLoadingProgressResponse{Status: &common.Status{ErrorCode: common.ErrorCode_Success}, Progress: 100}, nil)
+
+		err := c.LoadPartitions(ctx, testCollectionName, partNames, false)
+		s.NoError(err)
+	})
+
+	s.Run("has_collection_failure", func() {
+		s.Run("return_false", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: false}, nil)
+
+			err := c.LoadPartitions(ctx, testCollectionName, partNames, false)
+			s.Error(err)
+		})
+
+		s.Run("return_error", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(nil, errors.New("mock error"))
+
+			err := c.LoadPartitions(ctx, testCollectionName, partNames, false)
+			s.Error(err)
+		})
+	})
+
+	s.Run("has_partition_failure", func() {
+		s.Run("return_false", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+			s.mock.EXPECT().HasPartition(mock.Anything, mock.Anything).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: false}, nil)
+
+			err := c.LoadPartitions(ctx, testCollectionName, partNames, false)
+			s.Error(err)
+		})
+
+		s.Run("return_error", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+			s.mock.EXPECT().HasPartition(mock.Anything, mock.Anything).
+				Return(nil, errors.New("mock"))
+
+			err := c.LoadPartitions(ctx, testCollectionName, partNames, false)
+			s.Error(err)
+		})
+	})
+
+	s.Run("load_partitions_failure", func() {
+		s.Run("fail_status_code", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+			s.mock.EXPECT().HasPartition(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.HasPartitionRequest) {
+				s.Equal(testCollectionName, req.GetCollectionName())
+				_, ok := mPartNames[req.GetPartitionName()]
+				s.True(ok)
+			}).Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+			s.mock.EXPECT().LoadPartitions(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.LoadPartitionsRequest) {
+				s.Equal(testCollectionName, req.GetCollectionName())
+				s.ElementsMatch(partNames, req.GetPartitionNames())
+			}).Return(&common.Status{ErrorCode: common.ErrorCode_UnexpectedError}, nil)
+
+			err := c.LoadPartitions(ctx, testCollectionName, partNames, true)
+			s.Error(err)
+		})
+
+		s.Run("return_error", func() {
+			defer s.resetMock()
+			s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+				Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+			s.mock.EXPECT().HasPartition(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.HasPartitionRequest) {
+				s.Equal(testCollectionName, req.GetCollectionName())
+				_, ok := mPartNames[req.GetPartitionName()]
+				s.True(ok)
+			}).Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+			s.mock.EXPECT().LoadPartitions(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.LoadPartitionsRequest) {
+				s.Equal(testCollectionName, req.GetCollectionName())
+				s.ElementsMatch(partNames, req.GetPartitionNames())
+			}).Return(&common.Status{ErrorCode: common.ErrorCode_UnexpectedError}, nil)
+
+			err := c.LoadPartitions(ctx, testCollectionName, partNames, true)
+			s.Error(err)
+		})
+	})
+
+	s.Run("get_loading_progress_failure", func() {
+		defer s.resetMock()
+		s.mock.EXPECT().HasCollection(mock.Anything, &server.HasCollectionRequest{CollectionName: testCollectionName}).
+			Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+		s.mock.EXPECT().HasPartition(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.HasPartitionRequest) {
+			s.Equal(testCollectionName, req.GetCollectionName())
+			_, ok := mPartNames[req.GetPartitionName()]
+			s.True(ok)
+		}).Return(&server.BoolResponse{Status: &common.Status{}, Value: true}, nil)
+		s.mock.EXPECT().LoadPartitions(mock.Anything, mock.Anything).Run(func(_ context.Context, req *server.LoadPartitionsRequest) {
+			s.Equal(testCollectionName, req.GetCollectionName())
+			s.ElementsMatch(partNames, req.GetPartitionNames())
+		}).Return(&common.Status{ErrorCode: common.ErrorCode_Success}, nil)
+		s.mock.EXPECT().GetLoadingProgress(mock.Anything, mock.Anything).
+			Return(nil, errors.New("mock error"))
+
+		err := c.LoadPartitions(ctx, testCollectionName, partNames, false)
+		s.Error(err)
+	})
+
+	s.Run("service_not_ready", func() {
+		c := &GrpcClient{}
+		err := c.LoadPartitions(ctx, testCollectionName, partNames, false)
+		s.ErrorIs(err, ErrClientNotReady)
+	})
+}
+
+func TestPartitionSuite(t *testing.T) {
+	suite.Run(t, new(PartitionSuite))
+}