github.com/cloudwego/kitex@v0.9.0/pkg/remote/codec/protobuf/protobuf_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 protobuf

import (
	"context"
	"errors"
	"testing"

	"google.golang.org/protobuf/proto"

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

var (
	payloadCodec = &protobufCodec{}
	svcInfo      = mocks.ServiceInfo()
)

func init() {
	svcInfo.Methods["mock"] = serviceinfo.NewMethodInfo(nil, newMockReqArgs, nil, false)
}

func TestNormal(t *testing.T) {
	ctx := context.Background()

	// encode // client side
	sendMsg := initSendMsg(transport.TTHeader)
	out := remote.NewWriterBuffer(256)
	err := payloadCodec.Marshal(ctx, sendMsg, out)
	test.Assert(t, err == nil, err)

	// decode server side
	recvMsg := initRecvMsg()
	buf, err := out.Bytes()
	recvMsg.SetPayloadLen(len(buf))
	test.Assert(t, err == nil, err)
	in := remote.NewReaderBuffer(buf)
	err = payloadCodec.Unmarshal(ctx, recvMsg, in)
	test.Assert(t, err == nil, err)

	// compare Req Arg
	sendReq := (sendMsg.Data()).(*MockReqArgs).Req
	recvReq := (recvMsg.Data()).(*MockReqArgs).Req
	test.Assert(t, sendReq.Msg == recvReq.Msg)
	test.Assert(t, len(sendReq.StrList) == len(recvReq.StrList))
	test.Assert(t, len(sendReq.StrMap) == len(recvReq.StrMap))
	for i, item := range sendReq.StrList {
		test.Assert(t, item == recvReq.StrList[i])
	}
	for k := range sendReq.StrMap {
		test.Assert(t, sendReq.StrMap[k] == recvReq.StrMap[k])
	}
}

func TestException(t *testing.T) {
	ctx := context.Background()
	ink := rpcinfo.NewInvocation("", "mock")
	ri := rpcinfo.NewRPCInfo(nil, nil, ink, nil, nil)
	errInfo := "mock exception"
	transErr := remote.NewTransErrorWithMsg(remote.UnknownMethod, errInfo)
	// encode server side
	errMsg := initServerErrorMsg(transport.TTHeader, ri, transErr)
	out := remote.NewWriterBuffer(256)
	err := payloadCodec.Marshal(ctx, errMsg, out)
	test.Assert(t, err == nil, err)

	// decode client side
	recvMsg := initClientRecvMsg(ri)
	buf, err := out.Bytes()
	recvMsg.SetPayloadLen(len(buf))
	test.Assert(t, err == nil, err)
	in := remote.NewReaderBuffer(buf)
	err = payloadCodec.Unmarshal(ctx, recvMsg, in)
	test.Assert(t, err != nil)
	transErr, ok := err.(*remote.TransError)
	test.Assert(t, ok)
	test.Assert(t, err.Error() == errInfo)
	test.Assert(t, transErr.Error() == errInfo)
	test.Assert(t, transErr.TypeID() == remote.UnknownMethod)
}

func TestTransErrorUnwrap(t *testing.T) {
	errMsg := "mock err"
	transErr := remote.NewTransError(remote.InternalError, NewPbError(1000, errMsg))
	uwErr, ok := transErr.Unwrap().(PBError)
	test.Assert(t, ok)
	test.Assert(t, uwErr.TypeID() == 1000)
	test.Assert(t, transErr.Error() == errMsg)

	uwErr2, ok := errors.Unwrap(transErr).(PBError)
	test.Assert(t, ok)
	test.Assert(t, uwErr2.TypeID() == 1000)
	test.Assert(t, uwErr2.Error() == errMsg)
}

func BenchmarkNormalParallel(b *testing.B) {
	ctx := context.Background()

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			// encode // client side
			sendMsg := initSendMsg(transport.TTHeader)
			out := remote.NewWriterBuffer(256)
			err := payloadCodec.Marshal(ctx, sendMsg, out)
			test.Assert(b, err == nil, err)

			// decode server side
			recvMsg := initRecvMsg()
			buf, err := out.Bytes()
			recvMsg.SetPayloadLen(len(buf))
			test.Assert(b, err == nil, err)
			in := remote.NewReaderBuffer(buf)
			err = payloadCodec.Unmarshal(ctx, recvMsg, in)
			test.Assert(b, err == nil, err)

			// compare Req Arg
			sendReq := (sendMsg.Data()).(*MockReqArgs).Req
			recvReq := (recvMsg.Data()).(*MockReqArgs).Req
			test.Assert(b, sendReq.Msg == recvReq.Msg)
			test.Assert(b, len(sendReq.StrList) == len(recvReq.StrList))
			test.Assert(b, len(sendReq.StrMap) == len(recvReq.StrMap))
			for i, item := range sendReq.StrList {
				test.Assert(b, item == recvReq.StrList[i])
			}
			for k := range sendReq.StrMap {
				test.Assert(b, sendReq.StrMap[k] == recvReq.StrMap[k])
			}
		}
	})
}

func initSendMsg(tp transport.Protocol) remote.Message {
	var _args MockReqArgs
	_args.Req = prepareReq()
	ink := rpcinfo.NewInvocation("", "mock")
	ri := rpcinfo.NewRPCInfo(nil, nil, ink, nil, nil)
	msg := remote.NewMessage(&_args, svcInfo, ri, remote.Call, remote.Client)
	msg.SetProtocolInfo(remote.NewProtocolInfo(tp, svcInfo.PayloadCodec))
	return msg
}

func initRecvMsg() remote.Message {
	var _args MockReqArgs
	ink := rpcinfo.NewInvocation("", "mock")
	ri := rpcinfo.NewRPCInfo(nil, nil, ink, nil, nil)
	msg := remote.NewMessage(&_args, svcInfo, ri, remote.Call, remote.Server)
	return msg
}

func initServerErrorMsg(tp transport.Protocol, ri rpcinfo.RPCInfo, transErr *remote.TransError) remote.Message {
	errMsg := remote.NewMessage(transErr, svcInfo, ri, remote.Exception, remote.Server)
	errMsg.SetProtocolInfo(remote.NewProtocolInfo(tp, svcInfo.PayloadCodec))
	return errMsg
}

func initClientRecvMsg(ri rpcinfo.RPCInfo) remote.Message {
	var resp interface{}
	clientRecvMsg := remote.NewMessage(resp, svcInfo, ri, remote.Reply, remote.Client)
	return clientRecvMsg
}

func prepareReq() *MockReq {
	strMap := make(map[string]string)
	strMap["key1"] = "val1"
	strMap["key2"] = "val2"
	strList := []string{"str1", "str2"}
	req := &MockReq{
		Msg:     "MockReq",
		StrMap:  strMap,
		StrList: strList,
	}
	return req
}

func newMockReqArgs() interface{} {
	return &MockReqArgs{}
}

type MockReqArgs struct {
	Req *MockReq
}

func (p *MockReqArgs) Marshal(out []byte) ([]byte, error) {
	if !p.IsSetReq() {
		return out, nil
	}
	return proto.Marshal(p.Req)
}

func (p *MockReqArgs) Unmarshal(in []byte) error {
	msg := new(MockReq)
	if err := proto.Unmarshal(in, msg); err != nil {
		return err
	}
	p.Req = msg
	return nil
}

var STServiceTestObjReqArgsReqDEFAULT *MockReq

func (p *MockReqArgs) GetReq() *MockReq {
	if !p.IsSetReq() {
		return STServiceTestObjReqArgsReqDEFAULT
	}
	return p.Req
}

func (p *MockReqArgs) IsSetReq() bool {
	return p.Req != nil
}