github.com/cloudwego/kitex@v0.9.0/pkg/remote/remotecli/client_test.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package remotecli
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"net"
    23  	"testing"
    24  
    25  	"github.com/golang/mock/gomock"
    26  
    27  	"github.com/cloudwego/kitex/internal/mocks"
    28  	mocksnetpoll "github.com/cloudwego/kitex/internal/mocks/netpoll"
    29  	mocksremote "github.com/cloudwego/kitex/internal/mocks/remote"
    30  	"github.com/cloudwego/kitex/internal/test"
    31  	"github.com/cloudwego/kitex/pkg/discovery"
    32  	"github.com/cloudwego/kitex/pkg/kerrors"
    33  	"github.com/cloudwego/kitex/pkg/remote"
    34  	"github.com/cloudwego/kitex/pkg/remote/connpool"
    35  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    36  	"github.com/cloudwego/kitex/pkg/rpcinfo/remoteinfo"
    37  	"github.com/cloudwego/kitex/pkg/transmeta"
    38  	"github.com/cloudwego/kitex/pkg/utils"
    39  )
    40  
    41  // TestNewClientNoAddr test new client return err because no addr
    42  func TestNewClientNoAddr(t *testing.T) {
    43  	ctrl := gomock.NewController(t)
    44  	defer ctrl.Finish()
    45  
    46  	// 1. prepare mock data
    47  	ctx := context.Background()
    48  
    49  	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
    50  	conf := rpcinfo.NewRPCConfig()
    51  	ri := rpcinfo.NewRPCInfo(nil, to, rpcinfo.NewInvocation("", ""), conf, rpcinfo.NewRPCStats())
    52  
    53  	hdlr := mocksremote.NewMockClientTransHandler(ctrl)
    54  	connPool := connpool.NewLongPool("destService", poolCfg)
    55  	opt := &remote.ClientOption{
    56  		ConnPool: connPool,
    57  		Dialer:   mocksremote.NewMockDialer(ctrl),
    58  	}
    59  
    60  	// 2. test
    61  	_, err := NewClient(ctx, ri, hdlr, opt)
    62  	test.Assert(t, err != nil, err)
    63  	test.Assert(t, errors.Is(err, kerrors.ErrGetConnection))
    64  }
    65  
    66  // TestNewClient test new a client
    67  func TestNewClient(t *testing.T) {
    68  	ctrl := gomock.NewController(t)
    69  	defer ctrl.Finish()
    70  
    71  	addr := utils.NewNetAddr("tcp", "to")
    72  	ri := newMockRPCInfo(addr)
    73  	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
    74  
    75  	hdlr := mocksremote.NewMockClientTransHandler(ctrl)
    76  
    77  	conn := mocksnetpoll.NewMockConnection(ctrl)
    78  	conn.EXPECT().RemoteAddr().Return(addr).AnyTimes()
    79  
    80  	dialer := mocksremote.NewMockDialer(ctrl)
    81  	dialer.EXPECT().DialTimeout(addr.Network(), addr.String(), gomock.Any()).Return(conn, nil).Times(1)
    82  
    83  	connPool := connpool.NewLongPool("destService", poolCfg)
    84  
    85  	opt := &remote.ClientOption{
    86  		ConnPool: connPool,
    87  		Dialer:   dialer,
    88  		Option: remote.Option{
    89  			StreamingMetaHandlers: []remote.StreamingMetaHandler{
    90  				transmeta.ClientHTTP2Handler,
    91  			},
    92  		},
    93  	}
    94  
    95  	cli, err := NewClient(ctx, ri, hdlr, opt)
    96  	test.Assert(t, err == nil, err)
    97  	test.Assert(t, cli != nil, cli)
    98  }
    99  
   100  func newMockRPCInfo(addr net.Addr) rpcinfo.RPCInfo {
   101  	from := rpcinfo.NewEndpointInfo("from", "method", nil, nil)
   102  	to := remoteinfo.NewRemoteInfo(&rpcinfo.EndpointBasicInfo{}, "")
   103  	to.SetInstance(discovery.NewInstance(addr.Network(), addr.String(), discovery.DefaultWeight, nil))
   104  	conf := rpcinfo.NewRPCConfig()
   105  	ri := rpcinfo.NewRPCInfo(from, to, rpcinfo.NewInvocation("", ""), conf, rpcinfo.NewRPCStats())
   106  
   107  	return ri
   108  }
   109  
   110  func newMockOption(ctrl *gomock.Controller, addr net.Addr) *remote.ClientOption {
   111  	conn := mocksnetpoll.NewMockConnection(ctrl)
   112  	conn.EXPECT().RemoteAddr().Return(addr).AnyTimes()
   113  	conn.EXPECT().IsActive().Return(true).AnyTimes()
   114  	conn.EXPECT().Close().Return(nil).AnyTimes()
   115  
   116  	dialer := mocksremote.NewMockDialer(ctrl)
   117  	dialer.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).Return(conn, nil).AnyTimes()
   118  
   119  	connPool := connpool.NewLongPool("destService", poolCfg)
   120  
   121  	opt := &remote.ClientOption{
   122  		ConnPool: connPool,
   123  		Dialer:   dialer,
   124  	}
   125  
   126  	return opt
   127  }
   128  
   129  // TestSend test send msg
   130  func TestSend(t *testing.T) {
   131  	ctrl := gomock.NewController(t)
   132  	defer ctrl.Finish()
   133  
   134  	// 1. prepare mock data
   135  	ctx := context.Background()
   136  	addr := utils.NewNetAddr("tcp", "to")
   137  	ri := newMockRPCInfo(addr)
   138  
   139  	var sendMsg remote.Message
   140  	hdlr := mocksremote.NewMockClientTransHandler(ctrl)
   141  	hdlr.EXPECT().Write(gomock.Any(), gomock.Any(), sendMsg).Return(ctx, nil).MinTimes(1)
   142  
   143  	opt := newMockOption(ctrl, addr)
   144  
   145  	cli, err := NewClient(ctx, ri, hdlr, opt)
   146  	test.Assert(t, err == nil, err)
   147  	test.Assert(t, cli != nil, cli)
   148  
   149  	// 2. test
   150  	err = cli.Send(ctx, ri, sendMsg)
   151  	test.Assert(t, err == nil, err)
   152  }
   153  
   154  // TestSendErr test send return err because write fail
   155  func TestSendErr(t *testing.T) {
   156  	ctrl := gomock.NewController(t)
   157  	defer ctrl.Finish()
   158  
   159  	// 1. prepare mock data
   160  	addr := utils.NewNetAddr("tcp", "to")
   161  	ri := newMockRPCInfo(addr)
   162  	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)
   163  
   164  	var sendMsg remote.Message
   165  	errMsg := "mock test send err"
   166  
   167  	hdlr := mocksremote.NewMockClientTransHandler(ctrl)
   168  	hdlr.EXPECT().Write(gomock.Any(), gomock.Any(), sendMsg).Return(ctx, errors.New(errMsg)).MinTimes(1)
   169  
   170  	opt := newMockOption(ctrl, addr)
   171  
   172  	conn := mocksnetpoll.NewMockConnection(ctrl)
   173  	conn.EXPECT().RemoteAddr().Return(addr).AnyTimes()
   174  	conn.EXPECT().Close().Return(nil).Times(1)
   175  
   176  	dialer := mocksremote.NewMockDialer(ctrl)
   177  	dialer.EXPECT().DialTimeout(addr.Network(), addr.String(), gomock.Any()).Return(conn, nil).Times(1)
   178  	opt.Dialer = dialer
   179  
   180  	// 2. test
   181  	cli, err := NewClient(ctx, ri, hdlr, opt)
   182  	test.Assert(t, err == nil, err)
   183  	test.Assert(t, cli != nil, cli)
   184  
   185  	err = cli.Send(ctx, ri, sendMsg)
   186  	test.Assert(t, err != nil, err)
   187  	test.Assert(t, err.Error() == errMsg)
   188  }
   189  
   190  // TestRecv test recv msg
   191  func TestRecv(t *testing.T) {
   192  	ctrl := gomock.NewController(t)
   193  	defer ctrl.Finish()
   194  
   195  	ctx := context.Background()
   196  
   197  	addr := utils.NewNetAddr("tcp", "to")
   198  	ri := newMockRPCInfo(addr)
   199  
   200  	svcInfo := mocks.ServiceInfo()
   201  	var resp interface{}
   202  	msg := remote.NewMessage(resp, svcInfo, ri, remote.Call, remote.Client)
   203  
   204  	hdlr := mocksremote.NewMockClientTransHandler(ctrl)
   205  	hdlr.EXPECT().Read(gomock.Any(), gomock.Any(), msg).Return(ctx, nil).MinTimes(1)
   206  	hdlr.EXPECT().OnMessage(gomock.Any(), gomock.Any(), gomock.Any()).Return(ctx, nil).AnyTimes()
   207  
   208  	opt := newMockOption(ctrl, addr)
   209  
   210  	cli, err := NewClient(ctx, ri, hdlr, opt)
   211  	test.Assert(t, err == nil, err)
   212  	test.Assert(t, cli != nil, cli)
   213  
   214  	err = cli.Recv(ctx, ri, msg)
   215  	test.Assert(t, err == nil, err)
   216  
   217  	readErr := errors.New("read failed")
   218  	onMessageErr := errors.New("on message failed")
   219  
   220  	hdlr = mocksremote.NewMockClientTransHandler(ctrl)
   221  	hdlr.EXPECT().Read(gomock.Any(), gomock.Any(), gomock.Any()).Return(ctx, readErr).AnyTimes()
   222  	hdlr.EXPECT().OnMessage(gomock.Any(), gomock.Any(), gomock.Any()).Return(ctx, onMessageErr).AnyTimes()
   223  
   224  	cli, err = NewClient(ctx, ri, hdlr, opt)
   225  	test.Assert(t, err == nil, err)
   226  	test.Assert(t, cli != nil, cli)
   227  
   228  	err = cli.Recv(ctx, ri, msg)
   229  	test.Assert(t, err == readErr, err)
   230  }
   231  
   232  // TestRecvOneWay test recv onw way msg
   233  func TestRecvOneWay(t *testing.T) {
   234  	ctrl := gomock.NewController(t)
   235  	defer ctrl.Finish()
   236  
   237  	ctx := context.Background()
   238  
   239  	addr := utils.NewNetAddr("tcp", "to")
   240  	ri := newMockRPCInfo(addr)
   241  
   242  	hdlr := mocksremote.NewMockClientTransHandler(ctrl)
   243  
   244  	opt := newMockOption(ctrl, addr)
   245  	opt.ConnPool = nil
   246  
   247  	conn := mocksnetpoll.NewMockConnection(ctrl)
   248  	conn.EXPECT().RemoteAddr().Return(addr).AnyTimes()
   249  	conn.EXPECT().Close().Return(nil).MinTimes(1)
   250  	dialer := mocksremote.NewMockDialer(ctrl)
   251  	dialer.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).Return(conn, nil).AnyTimes()
   252  	opt.Dialer = dialer
   253  
   254  	cli, err := NewClient(ctx, ri, hdlr, opt)
   255  	test.Assert(t, err == nil, err)
   256  	test.Assert(t, cli != nil, cli)
   257  
   258  	err = cli.Recv(ctx, ri, nil)
   259  	test.Assert(t, err == nil, err)
   260  }
   261  
   262  // TestRecycle test recycle
   263  func TestRecycle(t *testing.T) {
   264  	ctrl := gomock.NewController(t)
   265  	defer ctrl.Finish()
   266  
   267  	ctx := context.Background()
   268  
   269  	addr := utils.NewNetAddr("tcp", "to")
   270  	ri := newMockRPCInfo(addr)
   271  
   272  	hdlr := mocksremote.NewMockClientTransHandler(ctrl)
   273  
   274  	opt := newMockOption(ctrl, addr)
   275  
   276  	cli, err := NewClient(ctx, ri, hdlr, opt)
   277  	test.Assert(t, err == nil, err)
   278  	test.Assert(t, cli != nil, cli)
   279  
   280  	cli.Recycle()
   281  	afterCli, ok := cli.(*client)
   282  	test.Assert(t, ok)
   283  	test.Assert(t, afterCli.transHdlr == nil)
   284  	test.Assert(t, afterCli.conn == nil)
   285  	test.Assert(t, afterCli.connManager == nil)
   286  }