add more tests

example/rpc/proxy/proxy.go

@@ -42,7 +42,7 @@ func main() {
 		ListenOn: *listen,
 	}, func(grpcServer *grpc.Server) {
 		unary.RegisterGreeterServer(grpcServer, &GreetServer{
-			RpcProxy: rpcx.NewRpcProxy(*server),
+			RpcProxy: rpcx.NewProxy(*server),
 		})
 	})
 	proxy.Start()

go.mod

@@ -3,6 +3,7 @@ module github.com/tal-tech/go-zero
 go 1.14
 
 require (
+	9fans.net/go v0.0.2 // indirect
 	github.com/DATA-DOG/go-sqlmock v1.4.1
 	github.com/alicebob/gopher-json v0.0.0-20180125190556-5a6b3ba71ee6 // indirect
 	github.com/alicebob/miniredis v2.5.0+incompatible

go.sum

@@ -1,3 +1,5 @@
+9fans.net/go v0.0.2 h1:RYM6lWITV8oADrwLfdzxmt8ucfW6UtP9v1jg4qAbqts=
+9fans.net/go v0.0.2/go.mod h1:lfPdxjq9v8pVQXUMBCx5EO5oLXWQFlKRQgs1kEkjoIM=
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=

rpcx/client_test.go

@@ -0,0 +1,101 @@
+package rpcx
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"net"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/logx"
+	"github.com/tal-tech/go-zero/rpcx/mock"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+func init() {
+	logx.Disable()
+}
+
+func dialer() func(context.Context, string) (net.Conn, error) {
+	listener := bufconn.Listen(1024 * 1024)
+	server := grpc.NewServer()
+	mock.RegisterDepositServiceServer(server, &mock.DepositServer{})
+
+	go func() {
+		if err := server.Serve(listener); err != nil {
+			log.Fatal(err)
+		}
+	}()
+
+	return func(context.Context, string) (net.Conn, error) {
+		return listener.Dial()
+	}
+}
+
+func TestDepositServer_Deposit(t *testing.T) {
+	tests := []struct {
+		name    string
+		amount  float32
+		res     *mock.DepositResponse
+		errCode codes.Code
+		errMsg  string
+	}{
+		{
+			"invalid request with negative amount",
+			-1.11,
+			nil,
+			codes.InvalidArgument,
+			fmt.Sprintf("cannot deposit %v", -1.11),
+		},
+		{
+			"valid request with non negative amount",
+			0.00,
+			&mock.DepositResponse{Ok: true},
+			codes.OK,
+			"",
+		},
+	}
+
+	directClient := MustNewClient(RpcClientConf{
+		Endpoints: []string{"foo"},
+		App:       "foo",
+		Token:     "bar",
+		Timeout:   1000,
+	}, WithDialOption(grpc.WithInsecure()), WithDialOption(grpc.WithContextDialer(dialer())))
+	targetClient, err := NewClientWithTarget("foo", WithDialOption(grpc.WithInsecure()),
+		WithDialOption(grpc.WithContextDialer(dialer())))
+	assert.Nil(t, err)
+	clients := []Client{
+		directClient,
+		targetClient,
+	}
+	for _, tt := range tests {
+		for _, client := range clients {
+			t.Run(tt.name, func(t *testing.T) {
+				cli := mock.NewDepositServiceClient(client.Conn())
+				request := &mock.DepositRequest{Amount: tt.amount}
+				response, err := cli.Deposit(context.Background(), request)
+				if response != nil {
+					assert.True(t, len(response.String()) > 0)
+					if response.GetOk() != tt.res.GetOk() {
+						t.Error("response: expected", tt.res.GetOk(), "received", response.GetOk())
+					}
+				}
+				if err != nil {
+					if e, ok := status.FromError(err); ok {
+						if e.Code() != tt.errCode {
+							t.Error("error code: expected", codes.InvalidArgument, "received", e.Code())
+						}
+						if e.Message() != tt.errMsg {
+							t.Error("error message: expected", tt.errMsg, "received", e.Message())
+						}
+					}
+				}
+			})
+		}
+	}
+}

rpcx/mock/deposit.pb.go

@@ -0,0 +1,159 @@
+// Code generated by protoc-gen-go.
+// source: deposit.proto
+// DO NOT EDIT!
+
+/*
+Package mock is a generated protocol buffer package.
+
+It is generated from these files:
+	deposit.proto
+
+It has these top-level messages:
+	DepositRequest
+	DepositResponse
+*/
+package mock
+
+import proto "github.com/golang/protobuf/proto"
+import fmt "fmt"
+import math "math"
+
+import (
+	context "golang.org/x/net/context"
+	grpc "google.golang.org/grpc"
+)
+
+// Reference imports to suppress errors if they are not otherwise used.
+var _ = proto.Marshal
+var _ = fmt.Errorf
+var _ = math.Inf
+
+// This is a compile-time assertion to ensure that this generated file
+// is compatible with the proto package it is being compiled against.
+// A compilation error at this line likely means your copy of the
+// proto package needs to be updated.
+const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package
+
+type DepositRequest struct {
+	Amount float32 `protobuf:"fixed32,1,opt,name=amount" json:"amount,omitempty"`
+}
+
+func (m *DepositRequest) Reset()                    { *m = DepositRequest{} }
+func (m *DepositRequest) String() string            { return proto.CompactTextString(m) }
+func (*DepositRequest) ProtoMessage()               {}
+func (*DepositRequest) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{0} }
+
+func (m *DepositRequest) GetAmount() float32 {
+	if m != nil {
+		return m.Amount
+	}
+	return 0
+}
+
+type DepositResponse struct {
+	Ok bool `protobuf:"varint,1,opt,name=ok" json:"ok,omitempty"`
+}
+
+func (m *DepositResponse) Reset()                    { *m = DepositResponse{} }
+func (m *DepositResponse) String() string            { return proto.CompactTextString(m) }
+func (*DepositResponse) ProtoMessage()               {}
+func (*DepositResponse) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{1} }
+
+func (m *DepositResponse) GetOk() bool {
+	if m != nil {
+		return m.Ok
+	}
+	return false
+}
+
+func init() {
+	proto.RegisterType((*DepositRequest)(nil), "mock.DepositRequest")
+	proto.RegisterType((*DepositResponse)(nil), "mock.DepositResponse")
+}
+
+// Reference imports to suppress errors if they are not otherwise used.
+var _ context.Context
+var _ grpc.ClientConn
+
+// This is a compile-time assertion to ensure that this generated file
+// is compatible with the grpc package it is being compiled against.
+const _ = grpc.SupportPackageIsVersion4
+
+// Client API for DepositService service
+
+type DepositServiceClient interface {
+	Deposit(ctx context.Context, in *DepositRequest, opts ...grpc.CallOption) (*DepositResponse, error)
+}
+
+type depositServiceClient struct {
+	cc *grpc.ClientConn
+}
+
+func NewDepositServiceClient(cc *grpc.ClientConn) DepositServiceClient {
+	return &depositServiceClient{cc}
+}
+
+func (c *depositServiceClient) Deposit(ctx context.Context, in *DepositRequest, opts ...grpc.CallOption) (*DepositResponse, error) {
+	out := new(DepositResponse)
+	err := grpc.Invoke(ctx, "/mock.DepositService/Deposit", in, out, c.cc, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+// Server API for DepositService service
+
+type DepositServiceServer interface {
+	Deposit(context.Context, *DepositRequest) (*DepositResponse, error)
+}
+
+func RegisterDepositServiceServer(s *grpc.Server, srv DepositServiceServer) {
+	s.RegisterService(&_DepositService_serviceDesc, srv)
+}
+
+func _DepositService_Deposit_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(DepositRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(DepositServiceServer).Deposit(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: "/mock.DepositService/Deposit",
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(DepositServiceServer).Deposit(ctx, req.(*DepositRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+var _DepositService_serviceDesc = grpc.ServiceDesc{
+	ServiceName: "mock.DepositService",
+	HandlerType: (*DepositServiceServer)(nil),
+	Methods: []grpc.MethodDesc{
+		{
+			MethodName: "Deposit",
+			Handler:    _DepositService_Deposit_Handler,
+		},
+	},
+	Streams:  []grpc.StreamDesc{},
+	Metadata: "deposit.proto",
+}
+
+func init() { proto.RegisterFile("deposit.proto", fileDescriptor0) }
+
+var fileDescriptor0 = []byte{
+	// 139 bytes of a gzipped FileDescriptorProto
+	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x02, 0xff, 0xe2, 0xe2, 0x4d, 0x49, 0x2d, 0xc8,
+	0x2f, 0xce, 0x2c, 0xd1, 0x2b, 0x28, 0xca, 0x2f, 0xc9, 0x17, 0x62, 0xc9, 0xcd, 0x4f, 0xce, 0x56,
+	0xd2, 0xe0, 0xe2, 0x73, 0x81, 0x08, 0x07, 0xa5, 0x16, 0x96, 0xa6, 0x16, 0x97, 0x08, 0x89, 0x71,
+	0xb1, 0x25, 0xe6, 0xe6, 0x97, 0xe6, 0x95, 0x48, 0x30, 0x2a, 0x30, 0x6a, 0x30, 0x05, 0x41, 0x79,
+	0x4a, 0x8a, 0x5c, 0xfc, 0x70, 0x95, 0xc5, 0x05, 0xf9, 0x79, 0xc5, 0xa9, 0x42, 0x7c, 0x5c, 0x4c,
+	0xf9, 0xd9, 0x60, 0x65, 0x1c, 0x41, 0x4c, 0xf9, 0xd9, 0x46, 0x1e, 0x70, 0xc3, 0x82, 0x53, 0x8b,
+	0xca, 0x32, 0x93, 0x53, 0x85, 0xcc, 0xb8, 0xd8, 0xa1, 0x22, 0x42, 0x22, 0x7a, 0x20, 0x0b, 0xf5,
+	0x50, 0x6d, 0x93, 0x12, 0x45, 0x13, 0x85, 0x98, 0x9c, 0xc4, 0x06, 0x76, 0xa3, 0x31, 0x20, 0x00,
+	0x00, 0xff, 0xff, 0x62, 0x37, 0xf2, 0x36, 0xb4, 0x00, 0x00, 0x00,
+}

rpcx/mock/deposit.proto

@@ -0,0 +1,15 @@
+syntax = "proto3";
+
+package mock;
+
+message DepositRequest {
+  float amount = 1;
+}
+
+message DepositResponse {
+  bool ok = 1;
+}
+
+service DepositService {
+  rpc Deposit(DepositRequest) returns (DepositResponse);
+}

rpcx/mock/depositserver.go

@@ -0,0 +1,19 @@
+package mock
+
+import (
+	"context"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type DepositServer struct {
+}
+
+func (*DepositServer) Deposit(ctx context.Context, req *DepositRequest) (*DepositResponse, error) {
+	if req.GetAmount() < 0 {
+		return nil, status.Errorf(codes.InvalidArgument, "cannot deposit %v", req.GetAmount())
+	}
+
+	return &DepositResponse{Ok: true}, nil
+}

rpcx/proxy.go

@@ -18,7 +18,7 @@ type RpcProxy struct {
 	lock        sync.Mutex
 }
 
-func NewRpcProxy(backend string, opts ...internal.ClientOption) *RpcProxy {
+func NewProxy(backend string, opts ...internal.ClientOption) *RpcProxy {
 	return &RpcProxy{
 		backend:     backend,
 		clients:     make(map[string]Client),
@@ -56,5 +56,5 @@ func (p *RpcProxy) TakeConn(ctx context.Context) (*grpc.ClientConn, error) {
 		return nil, err
 	}
 
-	return val.(*RpcClient).Conn(), nil
+	return val.(Client).Conn(), nil
 }

rpcx/proxy_test.go

@@ -0,0 +1,66 @@
+package rpcx
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/rpcx/mock"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func TestProxy(t *testing.T) {
+	tests := []struct {
+		name    string
+		amount  float32
+		res     *mock.DepositResponse
+		errCode codes.Code
+		errMsg  string
+	}{
+		{
+			"invalid request with negative amount",
+			-1.11,
+			nil,
+			codes.InvalidArgument,
+			fmt.Sprintf("cannot deposit %v", -1.11),
+		},
+		{
+			"valid request with non negative amount",
+			0.00,
+			&mock.DepositResponse{Ok: true},
+			codes.OK,
+			"",
+		},
+	}
+
+	proxy := NewProxy("foo", WithDialOption(grpc.WithInsecure()),
+		WithDialOption(grpc.WithContextDialer(dialer())))
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			conn, err := proxy.TakeConn(context.Background())
+			assert.Nil(t, err)
+			cli := mock.NewDepositServiceClient(conn)
+			request := &mock.DepositRequest{Amount: tt.amount}
+			response, err := cli.Deposit(context.Background(), request)
+			if response != nil {
+				assert.True(t, len(response.String()) > 0)
+				if response.GetOk() != tt.res.GetOk() {
+					t.Error("response: expected", tt.res.GetOk(), "received", response.GetOk())
+				}
+			}
+			if err != nil {
+				if e, ok := status.FromError(err); ok {
+					if e.Code() != tt.errCode {
+						t.Error("error code: expected", codes.InvalidArgument, "received", e.Code())
+					}
+					if e.Message() != tt.errMsg {
+						t.Error("error message: expected", tt.errMsg, "received", e.Message())
+					}
+				}
+			}
+		})
+	}
+}