github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/default_server_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 trans

import (
	"context"
	"errors"
	"net"
	"testing"

	"github.com/golang/mock/gomock"

	"github.com/cloudwego/kitex/internal/mocks"
	"github.com/cloudwego/kitex/internal/mocks/stats"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
)

var (
	svcInfo      = mocks.ServiceInfo()
	svcSearchMap = map[string]*serviceinfo.ServiceInfo{
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockMethod):          svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockExceptionMethod): svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockErrorMethod):     svcInfo,
		remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockOnewayMethod):    svcInfo,
		mocks.MockMethod:          svcInfo,
		mocks.MockExceptionMethod: svcInfo,
		mocks.MockErrorMethod:     svcInfo,
		mocks.MockOnewayMethod:    svcInfo,
	}
)

func TestDefaultSvrTransHandler(t *testing.T) {
	buf := remote.NewReaderWriterBuffer(1024)
	ext := &MockExtension{
		NewWriteByteBufferFunc: func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
			return buf
		},
		NewReadByteBufferFunc: func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
			return buf
		},
	}

	tagEncode, tagDecode := 0, 0
	opt := &remote.ServerOption{
		Codec: &MockCodec{
			EncodeFunc: func(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
				tagEncode++
				test.Assert(t, out == buf)
				return nil
			},
			DecodeFunc: func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
				tagDecode++
				test.Assert(t, in == buf)
				return nil
			},
		},
		SvcSearchMap:  svcSearchMap,
		TargetSvcInfo: svcInfo,
	}

	handler, err := NewDefaultSvrTransHandler(opt, ext)
	test.Assert(t, err == nil)

	ctx := context.Background()
	conn := &mocks.Conn{}
	msg := &MockMessage{
		RPCInfoFunc: func() rpcinfo.RPCInfo {
			return newMockRPCInfo()
		},
		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)
	test.Assert(t, tagEncode == 1, tagEncode)
	test.Assert(t, tagDecode == 0, tagDecode)

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

func TestSvrTransHandlerBizError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockTracer := stats.NewMockTracer(ctrl)
	mockTracer.EXPECT().Start(gomock.Any()).DoAndReturn(func(ctx context.Context) context.Context { return ctx }).AnyTimes()
	mockTracer.EXPECT().Finish(gomock.Any()).DoAndReturn(func(ctx context.Context) {
		err := rpcinfo.GetRPCInfo(ctx).Stats().Error()
		test.Assert(t, err != nil)
	}).AnyTimes()

	buf := remote.NewReaderWriterBuffer(1024)
	ext := &MockExtension{
		NewWriteByteBufferFunc: func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
			return buf
		},
		NewReadByteBufferFunc: func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
			return buf
		},
	}

	tracerCtl := &rpcinfo.TraceController{}
	tracerCtl.Append(mockTracer)
	opt := &remote.ServerOption{
		Codec: &MockCodec{
			EncodeFunc: func(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
				return nil
			},
			DecodeFunc: func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
				msg.SpecifyServiceInfo(mocks.MockServiceName, mocks.MockMethod)
				return nil
			},
		},
		SvcSearchMap:  svcSearchMap,
		TargetSvcInfo: svcInfo,
		TracerCtl:     tracerCtl,
		InitOrResetRPCInfoFunc: func(ri rpcinfo.RPCInfo, addr net.Addr) rpcinfo.RPCInfo {
			rpcinfo.AsMutableEndpointInfo(ri.From()).SetAddress(addr)
			return ri
		},
	}
	ri := rpcinfo.NewRPCInfo(rpcinfo.EmptyEndpointInfo(), rpcinfo.FromBasicInfo(&rpcinfo.EndpointBasicInfo{}),
		rpcinfo.NewInvocation("", mocks.MockMethod), nil, rpcinfo.NewRPCStats())
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)

	svrHandler, err := NewDefaultSvrTransHandler(opt, ext)
	pl := remote.NewTransPipeline(svrHandler)
	svrHandler.SetPipeline(pl)
	if setter, ok := svrHandler.(remote.InvokeHandleFuncSetter); ok {
		setter.SetInvokeHandleFunc(func(ctx context.Context, req, resp interface{}) (err error) {
			return kerrors.ErrBiz.WithCause(errors.New("mock"))
		})
	}
	test.Assert(t, err == nil)
	err = svrHandler.OnRead(ctx, &mocks.Conn{})
	test.Assert(t, err == nil)
}

func TestSvrTransHandlerReadErr(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockTracer := stats.NewMockTracer(ctrl)
	mockTracer.EXPECT().Start(gomock.Any()).DoAndReturn(func(ctx context.Context) context.Context { return ctx }).AnyTimes()
	mockTracer.EXPECT().Finish(gomock.Any()).DoAndReturn(func(ctx context.Context) {
		err := rpcinfo.GetRPCInfo(ctx).Stats().Error()
		test.Assert(t, err != nil)
	}).AnyTimes()

	buf := remote.NewReaderWriterBuffer(1024)
	ext := &MockExtension{
		NewWriteByteBufferFunc: func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
			return buf
		},
		NewReadByteBufferFunc: func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
			return buf
		},
	}

	mockErr := errors.New("mock")
	tracerCtl := &rpcinfo.TraceController{}
	tracerCtl.Append(mockTracer)
	opt := &remote.ServerOption{
		Codec: &MockCodec{
			EncodeFunc: func(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
				return nil
			},
			DecodeFunc: func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
				msg.SpecifyServiceInfo(mocks.MockServiceName, mocks.MockMethod)
				return mockErr
			},
		},
		SvcSearchMap:  svcSearchMap,
		TargetSvcInfo: svcInfo,
		TracerCtl:     tracerCtl,
		InitOrResetRPCInfoFunc: func(ri rpcinfo.RPCInfo, addr net.Addr) rpcinfo.RPCInfo {
			rpcinfo.AsMutableEndpointInfo(ri.From()).SetAddress(addr)
			return ri
		},
	}
	ri := rpcinfo.NewRPCInfo(rpcinfo.EmptyEndpointInfo(), rpcinfo.FromBasicInfo(&rpcinfo.EndpointBasicInfo{}),
		rpcinfo.NewInvocation("", mocks.MockMethod), nil, rpcinfo.NewRPCStats())
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)

	svrHandler, err := NewDefaultSvrTransHandler(opt, ext)
	test.Assert(t, err == nil)
	pl := remote.NewTransPipeline(svrHandler)
	svrHandler.SetPipeline(pl)
	err = svrHandler.OnRead(ctx, &mocks.Conn{})
	test.Assert(t, err != nil)
	test.Assert(t, errors.Is(err, mockErr))
}

func TestSvrTransHandlerOnReadHeartbeat(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockTracer := stats.NewMockTracer(ctrl)
	mockTracer.EXPECT().Start(gomock.Any()).DoAndReturn(func(ctx context.Context) context.Context { return ctx }).AnyTimes()
	mockTracer.EXPECT().Finish(gomock.Any()).DoAndReturn(func(ctx context.Context) {
		err := rpcinfo.GetRPCInfo(ctx).Stats().Error()
		test.Assert(t, err == nil)
	}).AnyTimes()

	buf := remote.NewReaderWriterBuffer(1024)
	ext := &MockExtension{
		NewWriteByteBufferFunc: func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
			return buf
		},
		NewReadByteBufferFunc: func(ctx context.Context, conn net.Conn, msg remote.Message) remote.ByteBuffer {
			return buf
		},
	}

	tracerCtl := &rpcinfo.TraceController{}
	tracerCtl.Append(mockTracer)
	opt := &remote.ServerOption{
		Codec: &MockCodec{
			EncodeFunc: func(ctx context.Context, msg remote.Message, out remote.ByteBuffer) error {
				if msg.MessageType() != remote.Heartbeat {
					return errors.New("response is not of MessageType Heartbeat")
				}
				return nil
			},
			DecodeFunc: func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
				msg.SetMessageType(remote.Heartbeat)
				msg.SpecifyServiceInfo(mocks.MockServiceName, mocks.MockMethod)
				return nil
			},
		},
		SvcSearchMap:  svcSearchMap,
		TargetSvcInfo: svcInfo,
		TracerCtl:     tracerCtl,
		InitOrResetRPCInfoFunc: func(ri rpcinfo.RPCInfo, addr net.Addr) rpcinfo.RPCInfo {
			rpcinfo.AsMutableEndpointInfo(ri.From()).SetAddress(addr)
			return ri
		},
	}
	ri := rpcinfo.NewRPCInfo(rpcinfo.EmptyEndpointInfo(), rpcinfo.FromBasicInfo(&rpcinfo.EndpointBasicInfo{}),
		rpcinfo.NewInvocation("", mocks.MockMethod), nil, rpcinfo.NewRPCStats())
	ctx := rpcinfo.NewCtxWithRPCInfo(context.Background(), ri)

	svrHandler, err := NewDefaultSvrTransHandler(opt, ext)
	test.Assert(t, err == nil)
	pl := remote.NewTransPipeline(svrHandler)
	svrHandler.SetPipeline(pl)
	err = svrHandler.OnRead(ctx, &mocks.Conn{})
	test.Assert(t, err == nil)
}