Skip to content

Commit

Permalink
remove mannul mock code
Browse files Browse the repository at this point in the history 
Signed-off-by: Wei Liu <wei.liu@zilliz.com>
  • Loading branch information
@weiliu1031
weiliu1031 committed Dec 8, 2023
1 parent 7e81245 commit 320139f
Show file tree
Hide file tree
Showing 8 changed files with 127 additions and 230 deletions.
2 changes: 1 addition & 1 deletion internal/distributed/datacoord/service_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,9 @@ package grpcdatacoord

import (
"context"
"errors"
"testing"

"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/tikv/client-go/v2/txnkv"
Expand Down
94 changes: 48 additions & 46 deletions internal/distributed/datanode/service_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,16 +18,16 @@ package grpcdatanode

import (
"context"
"fmt"
"testing"

"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
clientv3 "go.etcd.io/etcd/client/v3"
"google.golang.org/grpc"

"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/types"
Expand Down Expand Up @@ -183,12 +183,15 @@ func (m *MockDataNode) DropImport(ctx context.Context, req *datapb.DropImportReq
}

// /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
type mockDataCoord struct {
types.DataCoordClient
}
func Test_NewServer(t *testing.T) {
paramtable.Init()
ctx := context.Background()
server, err := NewServer(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, server)

func (m *mockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
mockRootCoord := mocks.NewMockRootCoordClient(t)
mockRootCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
StateCode: commonpb.StateCode_Healthy,
},
Expand All @@ -198,20 +201,13 @@ func (m *mockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.Ge
StateCode: commonpb.StateCode_Healthy,
},
},
}, nil
}

func (m *mockDataCoord) Stop() error {
return fmt.Errorf("stop error")
}

// /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
type mockRootCoord struct {
types.RootCoordClient
}
}, nil)
server.newRootCoordClient = func() (types.RootCoordClient, error) {
return mockRootCoord, nil
}

func (m *mockRootCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
mockDataCoord := mocks.NewMockDataCoordClient(t)
mockDataCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
StateCode: commonpb.StateCode_Healthy,
},
Expand All @@ -221,27 +217,9 @@ func (m *mockRootCoord) GetComponentStates(ctx context.Context, req *milvuspb.Ge
StateCode: commonpb.StateCode_Healthy,
},
},
}, nil
}

func (m *mockRootCoord) Stop() error {
return fmt.Errorf("stop error")
}

// /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
func Test_NewServer(t *testing.T) {
paramtable.Init()
ctx := context.Background()
server, err := NewServer(ctx, nil)
assert.NoError(t, err)
assert.NotNil(t, server)

server.newRootCoordClient = func() (types.RootCoordClient, error) {
return &mockRootCoord{}, nil
}

}, nil)
server.newDataCoordClient = func() (types.DataCoordClient, error) {
return &mockDataCoord{}, nil
return mockDataCoord, nil
}

t.Run("Run", func(t *testing.T) {
Expand Down Expand Up @@ -371,16 +349,40 @@ func Test_Run(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, server)

server.datanode = &MockDataNode{
regErr: errors.New("error"),
}

mockRootCoord := mocks.NewMockRootCoordClient(t)
mockRootCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
StateCode: commonpb.StateCode_Healthy,
},
Status: merr.Success(),
SubcomponentStates: []*milvuspb.ComponentInfo{
{
StateCode: commonpb.StateCode_Healthy,
},
},
}, nil)
server.newRootCoordClient = func() (types.RootCoordClient, error) {
return &mockRootCoord{}, nil
return mockRootCoord, nil
}

mockDataCoord := mocks.NewMockDataCoordClient(t)
mockDataCoord.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
StateCode: commonpb.StateCode_Healthy,
},
Status: merr.Success(),
SubcomponentStates: []*milvuspb.ComponentInfo{
{
StateCode: commonpb.StateCode_Healthy,
},
},
}, nil)
server.newDataCoordClient = func() (types.DataCoordClient, error) {
return &mockDataCoord{}, nil
return mockDataCoord, nil
}

server.datanode = &MockDataNode{
regErr: errors.New("error"),
}

err = server.Run()
Expand Down
56 changes: 11 additions & 45 deletions internal/distributed/querycoord/service_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,55 +25,17 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/tikv/client-go/v2/txnkv"
"google.golang.org/grpc"

"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/proto/internalpb"
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/types"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/milvus-io/milvus/pkg/util/tikv"
)

// /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
type MockRootCoord struct {
types.RootCoordClient
stopErr error
stateErr commonpb.ErrorCode
}

func (m *MockRootCoord) Close() error {
return m.stopErr
}

func (m *MockRootCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opt ...grpc.CallOption) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
Status: &commonpb.Status{ErrorCode: m.stateErr},
}, nil
}

// /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
type MockDataCoord struct {
types.DataCoordClient
stopErr error
stateErr commonpb.ErrorCode
}

func (m *MockDataCoord) Close() error {
return m.stopErr
}

func (m *MockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
Status: &commonpb.Status{ErrorCode: m.stateErr},
}, nil
}

// /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
func TestMain(m *testing.M) {
paramtable.Init()
Expand All @@ -96,13 +58,17 @@ func Test_NewServer(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, server)

mdc := &MockDataCoord{
stateErr: commonpb.ErrorCode_Success,
}

mrc := &MockRootCoord{
stateErr: commonpb.ErrorCode_Success,
}
mdc := mocks.NewMockDataCoordClient(t)
mdc.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil)

mrc := mocks.NewMockRootCoordClient(t)
mrc.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
}, nil)

mqc := getQueryCoord()
successStatus := merr.Success()
Expand Down
66 changes: 21 additions & 45 deletions internal/distributed/rootcoord/service_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,20 +24,18 @@ import (
"testing"
"time"

"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/tikv/client-go/v2/txnkv"
clientv3 "go.etcd.io/etcd/client/v3"
"google.golang.org/grpc"

"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/internal/mocks"
"github.com/milvus-io/milvus/internal/rootcoord"
"github.com/milvus-io/milvus/internal/types"
kvfactory "github.com/milvus-io/milvus/internal/util/dependency/kv"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/util/etcd"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/milvus-io/milvus/pkg/util/tikv"
)
Expand Down Expand Up @@ -109,42 +107,6 @@ func (m *mockCore) Stop() error {
return fmt.Errorf("stop error")
}

type mockDataCoord struct {
types.DataCoordClient
}

func (m *mockDataCoord) Close() error {
return nil
}

func (m *mockDataCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest, opts ...grpc.CallOption) (*milvuspb.ComponentStates, error) {
return &milvuspb.ComponentStates{
State: &milvuspb.ComponentInfo{
StateCode: commonpb.StateCode_Healthy,
},
Status: merr.Success(),
SubcomponentStates: []*milvuspb.ComponentInfo{
{
StateCode: commonpb.StateCode_Healthy,
},
},
}, nil
}

func (m *mockDataCoord) Stop() error {
return fmt.Errorf("stop error")
}

type mockQueryCoord struct {
types.QueryCoordClient
initErr error
startErr error
}

func (m *mockQueryCoord) Close() error {
return fmt.Errorf("stop error")
}

func TestRun(t *testing.T) {
paramtable.Init()
parameters := []string{"tikv", "etcd"}
Expand All @@ -169,11 +131,16 @@ func TestRun(t *testing.T) {
assert.Error(t, err)
assert.EqualError(t, err, "listen tcp: address 1000000: invalid port")

mockDataCoord := mocks.NewMockDataCoordClient(t)
mockDataCoord.EXPECT().Close().Return(nil)
svr.newDataCoordClient = func() types.DataCoordClient {
return &mockDataCoord{}
return mockDataCoord
}

mockQueryCoord := mocks.NewMockQueryCoordClient(t)
mockQueryCoord.EXPECT().Close().Return(nil)
svr.newQueryCoordClient = func() types.QueryCoordClient {
return &mockQueryCoord{}
return mockQueryCoord
}

paramtable.Get().Save(rcServerConfig.Port.Key, fmt.Sprintf("%d", rand.Int()%100+10000))
Expand Down Expand Up @@ -251,8 +218,10 @@ func TestServerRun_DataCoordClientInitErr(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, server)

mockDataCoord := mocks.NewMockDataCoordClient(t)
mockDataCoord.EXPECT().Close().Return(nil)
server.newDataCoordClient = func() types.DataCoordClient {
return &mockDataCoord{}
return mockDataCoord
}
assert.Panics(t, func() { server.Run() })

Expand All @@ -277,8 +246,10 @@ func TestServerRun_DataCoordClientStartErr(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, server)

mockDataCoord := mocks.NewMockDataCoordClient(t)
mockDataCoord.EXPECT().Close().Return(nil)
server.newDataCoordClient = func() types.DataCoordClient {
return &mockDataCoord{}
return mockDataCoord
}
assert.Panics(t, func() { server.Run() })

Expand All @@ -303,9 +274,12 @@ func TestServerRun_QueryCoordClientInitErr(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, server)

mockQueryCoord := mocks.NewMockQueryCoordClient(t)
mockQueryCoord.EXPECT().Close().Return(nil)
server.newQueryCoordClient = func() types.QueryCoordClient {
return &mockQueryCoord{initErr: errors.New("mock querycoord init error")}
return mockQueryCoord
}

assert.Panics(t, func() { server.Run() })

err = server.Stop()
Expand All @@ -329,8 +303,10 @@ func TestServer_QueryCoordClientStartErr(t *testing.T) {
assert.NoError(t, err)
assert.NotNil(t, server)

mockQueryCoord := mocks.NewMockQueryCoordClient(t)
mockQueryCoord.EXPECT().Close().Return(nil)
server.newQueryCoordClient = func() types.QueryCoordClient {
return &mockQueryCoord{startErr: errors.New("mock querycoord start error")}
return mockQueryCoord
}
assert.Panics(t, func() { server.Run() })

Expand Down
Loading

0 comments on commit 320139f

Please sign in to comment.