github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpollmux/client_handler_test.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package netpollmux

import (
	"context"
	"strings"
	"testing"
	"time"

	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/codec"
	"github.com/cloudwego/kitex/pkg/remote/trans"
	np "github.com/cloudwego/kitex/pkg/remote/trans/netpoll"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
)

func newTestRemoteClientOptionWithInStr(s string) *remote.ClientOption {
	opt := &remote.ClientOption{
		Codec: &MockCodec{
			EncodeFunc: func(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
				r := mockHeader(msg.RPCInfo().Invocation().SeqID(), s)
				_, err := out.WriteBinary(r.Bytes())
				return err
			},
			DecodeFunc: func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
				in.Skip(3 * codec.Size32)
				_, err := in.ReadString(len(s))
				return err
			},
		},
		ConnPool: NewMuxConnPool(1),
	}
	return opt
}

func TestNewCliTransHandler(t *testing.T) {
	// check success
	handler, err := NewCliTransHandlerFactory().NewTransHandler(&remote.ClientOption{
		ConnPool: &MuxPool{},
	})
	test.Assert(t, err == nil, err)
	test.Assert(t, handler != nil)
	// check fail
	handler, err = NewCliTransHandlerFactory().NewTransHandler(&remote.ClientOption{})
	test.Assert(t, err != nil)
	test.Assert(t, handler == nil)
}

func TestCliTransHandler(t *testing.T) {
	s := "hello world"
	opt := &remote.ClientOption{
		Codec: &MockCodec{
			EncodeFunc: func(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
				r := mockHeader(msg.RPCInfo().Invocation().SeqID(), s)
				n, err := out.WriteBinary(r.Bytes())
				test.Assert(t, err == nil, err)
				test.Assert(t, n == r.Len(), n, r.Len())
				return err
			},
			DecodeFunc: func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
				l := in.ReadableLen()
				test.Assert(t, l == 3*codec.Size32+len(s))
				in.Skip(3 * codec.Size32)
				got, err := in.ReadString(len(s))
				test.Assert(t, err == nil, err)
				test.Assert(t, got == s, got, s)
				return err
			},
		},
		ConnPool: NewMuxConnPool(1),
	}

	handler, err := NewCliTransHandlerFactory().NewTransHandler(opt)
	test.Assert(t, err == nil)

	ctx := context.Background()
	buf := netpoll.NewLinkBuffer(1024)
	npconn := &MockNetpollConn{
		ReaderFunc: func() (r netpoll.Reader) {
			return buf
		},
		WriterFunc: func() (r netpoll.Writer) {
			return buf
		},
	}
	conn := newMuxCliConn(npconn)

	ri := newMockRPCInfo()
	msg := &MockMessage{
		RPCInfoFunc: func() rpcinfo.RPCInfo {
			return ri
		},
		ServiceInfoFunc: func() *serviceinfo.ServiceInfo {
			return &serviceinfo.ServiceInfo{
				Methods: map[string]serviceinfo.MethodInfo{
					"method": serviceinfo.NewMethodInfo(nil, nil, nil, false),
				},
			}
		},
	}
	ctx, err = handler.Write(ctx, conn, msg)
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err == nil, err)

	time.Sleep(5 * time.Millisecond)
	buf.Flush()
	test.Assert(t, buf.Len() > len(s), buf.Len())

	err = conn.OnRequest(ctx, npconn)
	test.Assert(t, err == nil, err)

	ctx, err = handler.Read(ctx, conn, msg)
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err == nil, err)
}

// TestWriteNoMethod test client_handler write return unknown method err
func TestWriteNoMethod(t *testing.T) {
	// 1. prepare mock data
	s := "hello world"
	opt := newTestRemoteClientOptionWithInStr(s)

	handler, err := NewCliTransHandlerFactory().NewTransHandler(opt)
	test.Assert(t, err == nil)

	ctx := context.Background()

	var isWriteBufFlushed bool

	buf := netpoll.NewLinkBuffer(1024)
	npconn := &MockNetpollConn{
		ReaderFunc: func() (r netpoll.Reader) {
			return buf
		},
		WriterFunc: func() (r netpoll.Writer) {
			isWriteBufFlushed = true
			return buf
		},
	}
	conn := newMuxCliConn(npconn)

	ri := newMockRPCInfo()
	msg := &MockMessage{
		RPCInfoFunc: func() rpcinfo.RPCInfo {
			return ri
		},
		ServiceInfoFunc: func() *serviceinfo.ServiceInfo {
			return &serviceinfo.ServiceInfo{}
		},
	}

	// 2. test
	ctx, err = handler.Write(ctx, conn, msg)
	// check ctx/err not nil
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err != nil, err)
	tErr, ok := err.(*remote.TransError)
	test.Assert(t, ok)
	test.Assert(t, tErr.TypeID() == remote.UnknownMethod)
	test.Assert(t, !isWriteBufFlushed)
}

// TestWriteOneWayMethod test client_handler write oneway method
func TestWriteOneWayMethod(t *testing.T) {
	// 1. prepare mock data
	s := "hello world"
	opt := newTestRemoteClientOptionWithInStr(s)

	handler, err := NewCliTransHandlerFactory().NewTransHandler(opt)
	test.Assert(t, err == nil)

	ctx := context.Background()
	ri := newMockRPCInfo()

	buf := netpoll.NewLinkBuffer(2048)
	var isWrite bool
	var isRead bool
	npconn := &MockNetpollConn{
		ReaderFunc: func() (r netpoll.Reader) {
			isRead = true
			return buf
		},
		WriterFunc: func() (r netpoll.Writer) {
			isWrite = true
			return buf
		},
	}
	conn := newMuxCliConn(npconn)

	oneWayMsg := &MockMessage{
		RPCInfoFunc: func() rpcinfo.RPCInfo {
			return ri
		},
		ServiceInfoFunc: func() *serviceinfo.ServiceInfo {
			return &serviceinfo.ServiceInfo{
				Methods: map[string]serviceinfo.MethodInfo{
					"method": serviceinfo.NewMethodInfo(nil, nil, nil, true),
				},
			}
		},
	}
	ctx, err = handler.Write(ctx, conn, oneWayMsg)
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err == nil, err)

	time.Sleep(5 * time.Millisecond)
	buf.Flush()
	test.Assert(t, isWrite)
	test.Assert(t, buf.Len() > len(s), buf.Len())

	err = conn.OnRequest(ctx, npconn)
	test.Assert(t, err == nil, err)

	test.Assert(t, isRead)
}

// TestReadTimeout test client_handler read timeout because of no execute OnRequest
func TestReadTimeout(t *testing.T) {
	s := "hello world"
	opt := newTestRemoteClientOptionWithInStr(s)

	handler, err := NewCliTransHandlerFactory().NewTransHandler(opt)
	test.Assert(t, err == nil, err)

	ctx := context.Background()
	buf := netpoll.NewLinkBuffer(1024)

	var isRead, isWrite bool

	npconn := &MockNetpollConn{
		ReaderFunc: func() (r netpoll.Reader) {
			isRead = true
			return buf
		},
		WriterFunc: func() (r netpoll.Writer) {
			isWrite = true
			return buf
		},
	}
	conn := newMuxCliConn(npconn)

	cfg := rpcinfo.NewRPCConfig()
	rwTimeout := 1 * time.Second
	rpcinfo.AsMutableRPCConfig(cfg).SetRPCTimeout(rwTimeout)

	method := "method"
	c := rpcinfo.NewEndpointInfo("", method, nil, nil)
	to := rpcinfo.NewEndpointInfo("", method, nil, nil)
	ri := rpcinfo.NewRPCInfo(c, to, rpcinfo.NewInvocation("", method), cfg, rpcinfo.NewRPCStats())
	msg := &MockMessage{
		RPCInfoFunc: func() rpcinfo.RPCInfo {
			return ri
		},
		ServiceInfoFunc: func() *serviceinfo.ServiceInfo {
			return &serviceinfo.ServiceInfo{
				Methods: map[string]serviceinfo.MethodInfo{
					"method": serviceinfo.NewMethodInfo(nil, nil, nil, false),
				},
			}
		},
	}

	ctx, err = handler.Write(ctx, conn, msg)
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err == nil, err)

	time.Sleep(5 * time.Millisecond)
	buf.Flush()
	test.Assert(t, isWrite)
	test.Assert(t, buf.Len() > len(s), buf.Len())

	ctx, err = handler.Read(ctx, conn, msg)
	test.Assert(t, rwTimeout <= trans.GetReadTimeout(ri.Config()))
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err != nil, err)
	test.Assert(t, strings.Contains(err.Error(), "recv wait timeout"))
	test.Assert(t, !isRead)

	handler.OnError(ctx, err, conn)
}

// TestCallbackClose test asyncCallback Close
func TestCallbackClose(t *testing.T) {
	buf := netpoll.NewLinkBuffer()
	bufWriter := np.NewWriterByteBuffer(buf)
	callback := newAsyncCallback(buf, bufWriter)
	callback.Close()
}