Repository: milvus-io/milvus-sdk-go
Pull request #117 (2021-07-09 05:18:40)

Add client unit tests
Add client unit tests

Related #114 

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

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/117#" 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":[]} -->
sunby: /lgtm
/approve
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/117#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/117#issuecomment-876953167" title="Approved">sunby</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,sunby]

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":[]} -->

Patch:
diff --git a/client/client.go b/client/client.go
index 4e01b20..0628006 100644
--- a/client/client.go
+++ b/client/client.go
@@ -63,7 +63,7 @@ type Client interface {
 	CreateIndex(ctx context.Context, collName string, fieldName string, idx entity.Index) error
 	// DescribeIndex describe index on collection
 	// currently index naming is not supported, so only one index on vector field is supported
-	DescribeIndex(ctx context.Context, collName string) ([]entity.Index, error)
+	DescribeIndex(ctx context.Context, collName string, fieldName string) ([]entity.Index, error)
 	// DropINdex drop index from collection with specified field name
 	DropIndex(ctx context.Context, collName string, fieldName string) error
 	// GetIndexState get index state with specified collection and field name
diff --git a/client/client_grpc_dd.go b/client/client_grpc_dd.go
index 2b2be65..059b0fb 100644
--- a/client/client_grpc_dd.go
+++ b/client/client_grpc_dd.go
@@ -430,12 +430,13 @@ func (c *grpcClient) CreateIndex(ctx context.Context, collName string, fieldName
 }
 
 // DescribeIndex describe index
-func (c *grpcClient) DescribeIndex(ctx context.Context, collName string) ([]entity.Index, error) {
+func (c *grpcClient) DescribeIndex(ctx context.Context, collName string, fieldName string) ([]entity.Index, error) {
 	if c.service == nil {
 		return []entity.Index{}, ErrClientNotReady
 	}
 	req := &server.DescribeIndexRequest{
 		CollectionName: collName,
+		FieldName:      fieldName,
 		IndexName:      "", // empty string stands for all index on collection
 	}
 	resp, err := c.service.DescribeIndex(ctx, req)


Test patch:
diff --git a/client/client_grpc_dd_test.go b/client/client_grpc_dd_test.go
index f6785ed..5ac825b 100644
--- a/client/client_grpc_dd_test.go
+++ b/client/client_grpc_dd_test.go
@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"errors"
+	"fmt"
 	"math/rand"
 	"reflect"
 	"testing"
@@ -22,6 +23,9 @@ func TestGrpcClientNil(t *testing.T) {
 	tp := reflect.TypeOf(c)
 	v := reflect.ValueOf(c)
 	ctx := context.Background()
+	c2 := testClient(t, ctx)
+	v2 := reflect.ValueOf(c2)
+
 	ctxDone, cancel := context.WithCancel(context.Background())
 	cancel() // cancel here, so the ctx is done already
 
@@ -38,11 +42,9 @@ func TestGrpcClientNil(t *testing.T) {
 			if j == 1 {
 				//non-general solution, hard code context!
 				ins = append(ins, reflect.ValueOf(ctx))
-				t.Log("skip ctx")
 				continue
 			}
 			inT := mt.In(j)
-			t.Log(inT.String())
 
 			switch inT.Kind() {
 			case reflect.String: // pass empty
@@ -51,6 +53,11 @@ func TestGrpcClientNil(t *testing.T) {
 				ins = append(ins, reflect.ValueOf(0))
 			case reflect.Bool:
 				ins = append(ins, reflect.ValueOf(false))
+			case reflect.Interface:
+				idxType := reflect.TypeOf((*entity.Index)(nil)).Elem()
+				if inT.Implements(idxType) {
+					ins = append(ins, reflect.ValueOf(entity.NewFlatIndex(entity.L2)))
+				}
 			default:
 				ins = append(ins, reflect.Zero(inT))
 			}
@@ -60,9 +67,10 @@ func TestGrpcClientNil(t *testing.T) {
 		assert.EqualValues(t, ErrClientNotReady, outs[len(outs)-1].Interface())
 
 		// ctx done
+
 		if len(ins) > 0 { // with context param
 			ins[0] = reflect.ValueOf(ctxDone)
-			outs := v.MethodByName(m.Name).Call(ins)
+			outs := v2.MethodByName(m.Name).Call(ins)
 			assert.True(t, len(outs) > 0)
 			assert.False(t, outs[len(outs)-1].IsNil())
 		}
@@ -93,7 +101,7 @@ func TestGrpcClientClose(t *testing.T) {
 		assert.Nil(t, c.Close())
 	})
 
-	t.Run("dboule close", func(t *testing.T) {
+	t.Run("double close", func(t *testing.T) {
 		c := testClient(t, ctx)
 		assert.Nil(t, c.Close())
 		assert.Nil(t, c.Close())
@@ -104,7 +112,61 @@ func TestGrpcClientListCollections(t *testing.T) {
 	ctx := context.Background()
 	c := testClient(t, ctx)
 
-	c.ListCollections(ctx)
+	type testCase struct {
+		ids     []int64
+		names   []string
+		collNum int
+	}
+	caseLen := 5
+	cases := make([]testCase, 0, caseLen)
+	for i := 0; i < caseLen; i++ {
+		collNum := rand.Intn(5) + 2
+		tc := testCase{
+			ids:     make([]int64, 0, collNum),
+			names:   make([]string, 0, collNum),
+			collNum: collNum,
+		}
+		base := rand.Intn(1000)
+		for j := 0; j < collNum; j++ {
+			base += rand.Intn(1000)
+			tc.ids = append(tc.ids, int64(base))
+			base += rand.Intn(500)
+			tc.names = append(tc.names, fmt.Sprintf("coll_%d", base))
+		}
+		cases = append(cases, tc)
+	}
+
+	for _, tc := range cases {
+		mock.setInjection(mListCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+			s, err := successStatus()
+			resp := &server.ShowCollectionsResponse{
+				Status:          s,
+				CollectionIds:   tc.ids,
+				CollectionNames: tc.names,
+			}
+			return resp, err
+		})
+		collections, err := c.ListCollections(ctx)
+		if assert.Nil(t, err) && assert.Equal(t, tc.collNum, len(collections)) {
+			// assert element match
+			rids := make([]int64, 0, len(collections))
+			rnames := make([]string, 0, len(collections))
+			for _, collection := range collections {
+				rids = append(rids, collection.ID)
+				rnames = append(rnames, collection.Name)
+			}
+			assert.ElementsMatch(t, tc.ids, rids)
+			assert.ElementsMatch(t, tc.names, rnames)
+			// assert id & name match
+			for idx, rid := range rids {
+				for jdx, id := range tc.ids {
+					if rid == id {
+						assert.Equal(t, tc.names[jdx], rnames[idx])
+					}
+				}
+			}
+		}
+	}
 }
 
 func TestGrpcClientCreateCollection(t *testing.T) {
@@ -290,6 +352,215 @@ func TestGrpcClientLoadCollection(t *testing.T) {
 	})
 }
 
+func TestReleaseCollection(t *testing.T) {
+	ctx := context.Background()
+
+	c := testClient(t, ctx)
+
+	mock.setInjection(mReleaseCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.ReleaseCollectionRequest)
+		if !ok {
+			return badRequestStatus()
+		}
+		assert.Equal(t, testCollectionName, req.GetCollectionName())
+		return successStatus()
+	})
+
+	c.ReleaseCollection(ctx, testCollectionName)
+}
+
+func TestGrpcClientHasCollection(t *testing.T) {
+	ctx := context.Background()
+
+	c := testClient(t, ctx)
+
+	mock.setInjection(mHasCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.HasCollectionRequest)
+		resp := &server.BoolResponse{}
+		if !ok {
+			s, err := badRequestStatus()
+			assert.Fail(t, err.Error())
+			resp.Status = s
+			return resp, err
+		}
+		assert.Equal(t, req.CollectionName, testCollectionName)
+
+		s, err := successStatus()
+		resp.Status, resp.Value = s, true
+		return resp, err
+	})
+
+	has, err := c.HasCollection(ctx, testCollectionName)
+	assert.Nil(t, err)
+	assert.True(t, has)
+}
+
+func TestGrpcClientDescribeCollection(t *testing.T) {
+	ctx := context.Background()
+
+	c := testClient(t, ctx)
+
+	collectionID := rand.Int63()
+
+	mock.setInjection(mDescribeCollection, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.DescribeCollectionRequest)
+		resp := &server.DescribeCollectionResponse{}
+		if !ok {
+			s, err := badRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+
+		assert.Equal(t, testCollectionName, req.GetCollectionName())
+
+		sch := defaultSchema()
+		resp.Schema = sch.ProtoMessage()
+		resp.CollectionID = collectionID
+
+		s, err := successStatus()
+		resp.Status = s
+
+		return resp, err
+	})
+
+	collection, err := c.DescribeCollection(ctx, testCollectionName)
+	assert.Nil(t, err)
+	if assert.NotNil(t, collection) {
+		assert.Equal(t, collectionID, collection.ID)
+	}
+}
+
+func TestGrpcClientGetCollectionStatistics(t *testing.T) {
+	ctx := context.Background()
+
+	c := testClient(t, ctx)
+
+	stat := make(map[string]string)
+	stat["row_count"] = "0"
+
+	mock.setInjection(mGetCollectionStatistics, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.GetCollectionStatisticsRequest)
+		resp := &server.GetCollectionStatisticsResponse{}
+		if !ok {
+			s, err := badRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+		assert.Equal(t, testCollectionName, req.GetCollectionName())
+		s, err := successStatus()
+		resp.Status, resp.Stats = s, entity.MapKvPairs(stat)
+		return resp, err
+	})
+
+	rStat, err := c.GetCollectionStatistics(ctx, testCollectionName)
+	assert.Nil(t, err)
+	if assert.NotNil(t, rStat) {
+		for k, v := range stat {
+			rv, has := rStat[k]
+			assert.True(t, has)
+			assert.Equal(t, v, rv)
+		}
+	}
+}
+
+func TestGrpcClientCreatePartition(t *testing.T) {
+	//TODO add detail asserts
+	ctx := context.Background()
+	c := testClient(t, ctx)
+	assert.Nil(t, c.CreatePartition(ctx, testCollectionName, "_default_part"))
+}
+
+func TestGrpcClientDropPartition(t *testing.T) {
+	//TODO add detail asserts
+	ctx := context.Background()
+	c := testClient(t, ctx)
+	assert.Nil(t, c.DropPartition(ctx, testCollectionName, "_default_part"))
+}
+
+func TestGrpcClientHasPartition(t *testing.T) {
+	//TODO add detail asserts
+	ctx := context.Background()
+	c := testClient(t, ctx)
+	r, err := c.HasPartition(ctx, testCollectionName, "_default_part")
+	assert.Nil(t, err)
+	assert.False(t, r)
+}
+
+func TestGrpcClientShowPartitions(t *testing.T) {
+	ctx := context.Background()
+	c := testClient(t, ctx)
+	r, err := c.ShowPartitions(ctx, testCollectionName)
+	assert.Nil(t, err)
+	assert.NotNil(t, r)
+}
+
+func TestGrpcClientLoadPartitions(t *testing.T) {
+	//TODO add implementation
+	//ctx := context.Background()
+	//c := testClient(t, ctx)
+	//assert.Nil(t, c.LoadParitions(ctx, testCollectionName, "_default_part"))
+}
+
+func TestGrpcClientReleasePartitions(t *testing.T) {
+	//TODO add implementation
+}
+
+func TestGrpcClientCreateIndex(t *testing.T) {
+	ctx := context.Background()
+	c := testClient(t, ctx)
+	idx := entity.NewGenericIndex("", entity.Flat, map[string]string{
+		"nlist":       "1024",
+		"metric_type": "IP",
+	})
+
+	assert.Nil(t, c.CreateIndex(ctx, testCollectionName, "vector", idx))
+}
+
+func TestGrpcClientDropIndex(t *testing.T) {
+	ctx := context.Background()
+	c := testClient(t, ctx)
+
+	assert.Nil(t, c.DropIndex(ctx, testCollectionName, "vector"))
+}
+
+func TestGrpcClientDescribeIndex(t *testing.T) {
+	ctx := context.Background()
+
+	c := testClient(t, ctx)
+
+	fieldName := "vector"
+
+	mock.setInjection(mDescribeIndex, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.DescribeIndexRequest)
+		resp := &server.DescribeIndexResponse{}
+		if !ok {
+			s, err := badRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+		assert.Equal(t, fieldName, req.GetFieldName())
+		assert.Equal(t, testCollectionName, req.GetCollectionName())
+		resp.IndexDescriptions = []*server.IndexDescription{
+			{
+				IndexName: "_default",
+				IndexID:   1,
+				FieldName: req.GetFieldName(),
+				Params: entity.MapKvPairs(map[string]string{
+					"nlist":       "1024",
+					"metric_type": "IP",
+				}),
+			},
+		}
+		s, err := successStatus()
+		resp.Status = s
+		return resp, err
+	})
+
+	idxes, err := c.DescribeIndex(ctx, testCollectionName, fieldName)
+	assert.Nil(t, err)
+	assert.NotNil(t, idxes)
+}
+
 func successStatus() (*common.Status, error) {
 	return &common.Status{ErrorCode: common.ErrorCode_Success}, nil
 }
diff --git a/client/client_grpc_dm_test.go b/client/client_grpc_dm_test.go
new file mode 100644
index 0000000..0ff0a16
--- /dev/null
+++ b/client/client_grpc_dm_test.go
@@ -0,0 +1,86 @@
+package client
+
+import (
+	"context"
+	"math/rand"
+	"testing"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/milvus-io/milvus-sdk-go/entity"
+	"github.com/milvus-io/milvus-sdk-go/internal/proto/schema"
+	"github.com/milvus-io/milvus-sdk-go/internal/proto/server"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGrpcClientInsert(t *testing.T) {
+	ctx := context.Background()
+
+	c := testClient(t, ctx)
+
+	vector := generateFloatVector(4096, testVectorDim)
+	mock.setInjection(mInsert, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.InsertRequest)
+		resp := &server.MutationResult{}
+		if !ok {
+			s, err := badRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+		assert.EqualValues(t, 4096, req.GetNumRows())
+		assert.Equal(t, testCollectionName, req.GetCollectionName())
+		intIds := &schema.IDs_IntId{
+			IntId: &schema.LongArray{
+				Data: make([]int64, 4096),
+			},
+		}
+		resp.IDs = &schema.IDs{
+			IdField: intIds,
+		}
+		s, err := successStatus()
+		resp.Status = s
+		return resp, err
+	})
+	_, err := c.Insert(ctx, testCollectionName, "", // use default partition
+		entity.NewColumnFloatVector(testVectorField, testVectorDim, vector))
+
+	assert.Nil(t, err)
+}
+
+func TestGrpcClientFlush(t *testing.T) {
+
+	ctx := context.Background()
+
+	c := testClient(t, ctx)
+
+	assert.Nil(t, c.Flush(ctx, testCollectionName, false))
+	assert.Nil(t, c.Flush(ctx, testCollectionName, true))
+}
+
+func TestGrpcSearch(t *testing.T) {
+
+	ctx := context.Background()
+
+	c := testClient(t, ctx)
+
+	vectors := generateFloatVector(4096, testVectorDim)
+
+	results, err := c.Search(context.Background(), testCollectionName, []string{}, "int64 > 0", []string{"int64"}, []entity.Vector{entity.FloatVector(vectors[0])},
+		testVectorField, entity.L2, 10, map[string]string{
+			"nprobe": "10",
+		})
+
+	assert.Nil(t, err)
+	assert.NotNil(t, results)
+}
+
+func generateFloatVector(num, dim int) [][]float32 {
+	r := make([][]float32, 0, num)
+	for i := 0; i < num; i++ {
+		v := make([]float32, 0, dim)
+		for j := 0; j < dim; j++ {
+			v = append(v, rand.Float32())
+		}
+		r = append(r, v)
+	}
+	return r
+}
diff --git a/client/client_mock_test.go b/client/client_mock_test.go
index 4798562..41684c9 100644
--- a/client/client_mock_test.go
+++ b/client/client_mock_test.go
@@ -101,13 +101,31 @@ var (
 type serviceMethod int
 
 const (
-	mCreateCollection   serviceMethod = 1
-	mDropCollection     serviceMethod = 2
-	mHasCollection      serviceMethod = 3
-	mLoadCollection     serviceMethod = 4
-	mReleaseCollection  serviceMethod = 5
-	mDescribeCollection serviceMethod = 6
-	mListCollection     serviceMethod = 7
+	mCreateCollection        serviceMethod = 1
+	mDropCollection          serviceMethod = 2
+	mHasCollection           serviceMethod = 3
+	mLoadCollection          serviceMethod = 4
+	mReleaseCollection       serviceMethod = 5
+	mDescribeCollection      serviceMethod = 6
+	mListCollection          serviceMethod = 7
+	mGetCollectionStatistics serviceMethod = 8
+
+	mCreatePartition   serviceMethod = 9
+	mDropPartition     serviceMethod = 10
+	mHasPartition      serviceMethod = 11
+	mLoadPartitions    serviceMethod = 12
+	mReleasePartitions serviceMethod = 13
+	mShowPartitions    serviceMethod = 14
+
+	mCreateIndex           serviceMethod = 20
+	mDropIndex             serviceMethod = 21
+	mDescribeIndex         serviceMethod = 22
+	mGetIndexState         serviceMethod = 23
+	mGetIndexBuildProgress serviceMethod = 24
+
+	mInsert serviceMethod = 30
+	mFlush  serviceMethod = 31
+	mSearch serviceMethod = 32
 
 	mGetPersistentSegmentInfo serviceMethod = 98
 	mGetQuerySegmentInfo      serviceMethod = 99
@@ -193,15 +211,30 @@ func (m *mockServer) ReleaseCollection(ctx context.Context, req *server.ReleaseC
 		return r.(*common.Status), err
 	}
 	return successStatus()
-
 }
 
-func (m *mockServer) DescribeCollection(_ context.Context, _ *server.DescribeCollectionRequest) (*server.DescribeCollectionResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) DescribeCollection(ctx context.Context, req *server.DescribeCollectionRequest) (*server.DescribeCollectionResponse, error) {
+	f := m.getInjection(mDescribeCollection)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.DescribeCollectionResponse), err
+	}
+	r := &server.DescribeCollectionResponse{}
+	s, err := successStatus()
+	r.Status = s
+	return r, err
 }
 
-func (m *mockServer) GetCollectionStatistics(_ context.Context, _ *server.GetCollectionStatisticsRequest) (*server.GetCollectionStatisticsResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) GetCollectionStatistics(ctx context.Context, req *server.GetCollectionStatisticsRequest) (*server.GetCollectionStatisticsResponse, error) {
+	f := m.getInjection(mGetCollectionStatistics)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.GetCollectionStatisticsResponse), err
+	}
+	r := &server.GetCollectionStatisticsResponse{}
+	s, err := successStatus()
+	r.Status = s
+	return r, err
 }
 
 func (m *mockServer) ShowCollections(ctx context.Context, req *server.ShowCollectionsRequest) (*server.ShowCollectionsResponse, error) {
@@ -214,68 +247,149 @@ func (m *mockServer) ShowCollections(ctx context.Context, req *server.ShowCollec
 	return &server.ShowCollectionsResponse{Status: s}, err
 }
 
-func (m *mockServer) CreatePartition(_ context.Context, _ *server.CreatePartitionRequest) (*common.Status, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) CreatePartition(ctx context.Context, req *server.CreatePartitionRequest) (*common.Status, error) {
+	f := m.getInjection(mCreatePartition)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*common.Status), err
+	}
+	return successStatus()
+
 }
 
-func (m *mockServer) DropPartition(_ context.Context, _ *server.DropPartitionRequest) (*common.Status, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) DropPartition(ctx context.Context, req *server.DropPartitionRequest) (*common.Status, error) {
+	f := m.getInjection(mDropPartition)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*common.Status), err
+	}
+	return successStatus()
 }
 
-func (m *mockServer) HasPartition(_ context.Context, _ *server.HasPartitionRequest) (*server.BoolResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) HasPartition(ctx context.Context, req *server.HasPartitionRequest) (*server.BoolResponse, error) {
+	f := m.getInjection(mHasPartition)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.BoolResponse), err
+	}
+	s, err := successStatus()
+	return &server.BoolResponse{Status: s, Value: false}, err
 }
 
-func (m *mockServer) LoadPartitions(_ context.Context, _ *server.LoadPartitionsRequest) (*common.Status, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) LoadPartitions(ctx context.Context, req *server.LoadPartitionsRequest) (*common.Status, error) {
+	f := m.getInjection(mLoadPartitions)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*common.Status), err
+	}
+	return successStatus()
 }
 
-func (m *mockServer) ReleasePartitions(_ context.Context, _ *server.ReleasePartitionsRequest) (*common.Status, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) ReleasePartitions(ctx context.Context, req *server.ReleasePartitionsRequest) (*common.Status, error) {
+	f := m.getInjection(mReleasePartitions)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*common.Status), err
+	}
+	return successStatus()
 }
 
 func (m *mockServer) GetPartitionStatistics(_ context.Context, _ *server.GetPartitionStatisticsRequest) (*server.GetPartitionStatisticsResponse, error) {
 	panic("not implemented") // TODO: Implement
 }
 
-func (m *mockServer) ShowPartitions(_ context.Context, _ *server.ShowPartitionsRequest) (*server.ShowPartitionsResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) ShowPartitions(ctx context.Context, req *server.ShowPartitionsRequest) (*server.ShowPartitionsResponse, error) {
+	f := m.getInjection(mShowPartitions)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.ShowPartitionsResponse), err
+	}
+	s, err := successStatus()
+	return &server.ShowPartitionsResponse{Status: s}, err
 }
 
-func (m *mockServer) CreateIndex(_ context.Context, _ *server.CreateIndexRequest) (*common.Status, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) CreateIndex(ctx context.Context, req *server.CreateIndexRequest) (*common.Status, error) {
+	f := m.getInjection(mCreateIndex)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*common.Status), err
+	}
+	return successStatus()
 }
 
-func (m *mockServer) DescribeIndex(_ context.Context, _ *server.DescribeIndexRequest) (*server.DescribeIndexResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) DescribeIndex(ctx context.Context, req *server.DescribeIndexRequest) (*server.DescribeIndexResponse, error) {
+	f := m.getInjection(mDescribeIndex)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.DescribeIndexResponse), err
+	}
+	s, err := successStatus()
+	return &server.DescribeIndexResponse{Status: s}, err
 }
 
-func (m *mockServer) GetIndexState(_ context.Context, _ *server.GetIndexStateRequest) (*server.GetIndexStateResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) GetIndexState(ctx context.Context, req *server.GetIndexStateRequest) (*server.GetIndexStateResponse, error) {
+	f := m.getInjection(mGetIndexState)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.GetIndexStateResponse), err
+	}
+	s, err := successStatus()
+	return &server.GetIndexStateResponse{Status: s}, err
+
 }
 
-func (m *mockServer) GetIndexBuildProgress(_ context.Context, _ *server.GetIndexBuildProgressRequest) (*server.GetIndexBuildProgressResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) GetIndexBuildProgress(ctx context.Context, req *server.GetIndexBuildProgressRequest) (*server.GetIndexBuildProgressResponse, error) {
+	f := m.getInjection(mGetIndexBuildProgress)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.GetIndexBuildProgressResponse), err
+	}
+	s, err := successStatus()
+	return &server.GetIndexBuildProgressResponse{Status: s}, err
+
 }
 
-func (m *mockServer) DropIndex(_ context.Context, _ *server.DropIndexRequest) (*common.Status, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) DropIndex(ctx context.Context, req *server.DropIndexRequest) (*common.Status, error) {
+	f := m.getInjection(mDropIndex)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*common.Status), err
+	}
+	return successStatus()
 }
 
-func (m *mockServer) Insert(_ context.Context, _ *server.InsertRequest) (*server.MutationResult, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) Insert(ctx context.Context, req *server.InsertRequest) (*server.MutationResult, error) {
+	f := m.getInjection(mInsert)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.MutationResult), err
+	}
+	s, err := successStatus()
+	return &server.MutationResult{Status: s}, err
 }
 
-func (m *mockServer) Search(_ context.Context, _ *server.SearchRequest) (*server.SearchResults, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) Search(ctx context.Context, req *server.SearchRequest) (*server.SearchResults, error) {
+	f := m.getInjection(mSearch)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.SearchResults), err
+	}
+	s, err := successStatus()
+	return &server.SearchResults{Status: s}, err
 }
 
 func (m *mockServer) Retrieve(_ context.Context, _ *server.RetrieveRequest) (*server.RetrieveResults, error) {
 	panic("not implemented") // TODO: Implement
 }
 
-func (m *mockServer) Flush(_ context.Context, _ *server.FlushRequest) (*server.FlushResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) Flush(ctx context.Context, req *server.FlushRequest) (*server.FlushResponse, error) {
+	f := m.getInjection(mFlush)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.FlushResponse), err
+	}
+	s, err := successStatus()
+	return &server.FlushResponse{Status: s}, err
 }
 
 func (m *mockServer) Query(_ context.Context, _ *server.QueryRequest) (*server.QueryResults, error) {