github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpollmux/mux_conn_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 netpollmux
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  	"time"
    23  
    24  	"github.com/cloudwego/netpoll"
    25  
    26  	"github.com/cloudwego/kitex/internal/test"
    27  	"github.com/cloudwego/kitex/pkg/remote"
    28  	"github.com/cloudwego/kitex/pkg/remote/codec"
    29  	"github.com/cloudwego/kitex/pkg/rpcinfo"
    30  	"github.com/cloudwego/kitex/pkg/serviceinfo"
    31  )
    32  
    33  // TestOnRequest test muxSvrConn OnRequest return Err
    34  func TestOnRequestErr(t *testing.T) {
    35  	ctx := context.Background()
    36  
    37  	var isRead bool
    38  
    39  	buf := netpoll.NewLinkBuffer(3)
    40  	npconn := &MockNetpollConn{
    41  		ReaderFunc: func() (r netpoll.Reader) {
    42  			isRead = true
    43  			return buf
    44  		},
    45  	}
    46  	conn := newMuxCliConn(npconn)
    47  	err := conn.OnRequest(ctx, conn)
    48  	test.Assert(t, err != nil, err)
    49  	test.Assert(t, isRead)
    50  }
    51  
    52  // TestOnRequestErr test muxSvrConn OnRequest success
    53  func TestOnRequest(t *testing.T) {
    54  	s := "hello world"
    55  	ctx := context.Background()
    56  
    57  	opt := &remote.ClientOption{
    58  		Codec: &MockCodec{
    59  			EncodeFunc: func(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
    60  				r := mockHeader(msg.RPCInfo().Invocation().SeqID(), s)
    61  				n, err := out.WriteBinary(r.Bytes())
    62  				test.Assert(t, err == nil, err)
    63  				test.Assert(t, n == r.Len(), n, r.Len())
    64  				return err
    65  			},
    66  			DecodeFunc: func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
    67  				l := in.ReadableLen()
    68  				test.Assert(t, l == 3*codec.Size32+len(s))
    69  				in.Skip(3 * codec.Size32)
    70  				got, err := in.ReadString(len(s))
    71  				test.Assert(t, err == nil, err)
    72  				test.Assert(t, got == s, got, s)
    73  				return err
    74  			},
    75  		},
    76  		ConnPool: NewMuxConnPool(1),
    77  	}
    78  
    79  	buf := netpoll.NewLinkBuffer(1024)
    80  	conn := newMuxCliConn(&MockNetpollConn{
    81  		ReaderFunc: func() (r netpoll.Reader) {
    82  			return buf
    83  		},
    84  		WriterFunc: func() (r netpoll.Writer) {
    85  			return buf
    86  		},
    87  	})
    88  
    89  	handler, err := NewCliTransHandlerFactory().NewTransHandler(opt)
    90  	test.Assert(t, err == nil)
    91  
    92  	ri := newMockRPCInfo()
    93  	msg := &MockMessage{
    94  		RPCInfoFunc: func() rpcinfo.RPCInfo {
    95  			return ri
    96  		},
    97  		ServiceInfoFunc: func() *serviceinfo.ServiceInfo {
    98  			return &serviceinfo.ServiceInfo{
    99  				Methods: map[string]serviceinfo.MethodInfo{
   100  					"method": serviceinfo.NewMethodInfo(nil, nil, nil, false),
   101  				},
   102  			}
   103  		},
   104  	}
   105  	ctx, err = handler.Write(ctx, conn, msg)
   106  	test.Assert(t, ctx != nil, ctx)
   107  	test.Assert(t, err == nil, err)
   108  
   109  	time.Sleep(5 * time.Millisecond)
   110  	buf.Flush()
   111  
   112  	err = conn.OnRequest(ctx, conn)
   113  	test.Assert(t, err == nil, err)
   114  }