github.com/cloudwego/kitex@v0.9.0/pkg/generic/binarythrift_codec_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 generic

import (
	"context"
	"testing"

	"github.com/apache/thrift/lib/go/thrift"

	kt "github.com/cloudwego/kitex/internal/mocks/thrift"
	"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/pkg/utils"
)

func TestBinaryThriftCodec(t *testing.T) {
	req := kt.NewMockReq()
	args := kt.NewMockTestArgs()
	args.Req = req
	// encode
	rc := utils.NewThriftMessageCodec()
	buf, err := rc.Encode("mock", thrift.CALL, 100, args)
	test.Assert(t, err == nil, err)

	btc := &binaryThriftCodec{thriftCodec}
	cliMsg := &mockMessage{
		RPCInfoFunc: func() rpcinfo.RPCInfo {
			return newMockRPCInfo()
		},
		RPCRoleFunc: func() remote.RPCRole {
			return remote.Client
		},
		DataFunc: func() interface{} {
			return &Args{
				Request: buf,
				Method:  "mock",
			}
		},
	}
	seqID, err := GetSeqID(cliMsg.Data().(*Args).Request.(binaryReqType))
	test.Assert(t, err == nil, err)
	test.Assert(t, seqID == 100, seqID)

	rwbuf := remote.NewReaderWriterBuffer(1024)
	// change seqID to 1
	err = btc.Marshal(context.Background(), cliMsg, rwbuf)
	test.Assert(t, err == nil, err)
	seqID, err = GetSeqID(cliMsg.Data().(*Args).Request.(binaryReqType))
	test.Assert(t, err == nil, err)
	test.Assert(t, seqID == 1, seqID)

	// server side
	arg := &Args{}
	svrMsg := &mockMessage{
		RPCInfoFunc: func() rpcinfo.RPCInfo {
			return newMockRPCInfo()
		},
		RPCRoleFunc: func() remote.RPCRole {
			return remote.Server
		},
		DataFunc: func() interface{} {
			return arg
		},
		PayloadLenFunc: func() int {
			return rwbuf.ReadableLen()
		},
		ServiceInfoFunc: func() *serviceinfo.ServiceInfo {
			return ServiceInfo(serviceinfo.Thrift)
		},
	}
	err = btc.Unmarshal(context.Background(), svrMsg, rwbuf)
	test.Assert(t, err == nil, err)
	reqBuf := svrMsg.Data().(*Args).Request.(binaryReqType)
	seqID, err = GetSeqID(reqBuf)
	test.Assert(t, err == nil, err)
	test.Assert(t, seqID == 1, seqID)

	var req2 kt.MockTestArgs
	method, seqID2, err2 := rc.Decode(reqBuf, &req2)
	test.Assert(t, err2 == nil, err)
	test.Assert(t, seqID2 == 1, seqID)
	test.Assert(t, method == "mock", method)
}

func TestBinaryThriftCodecExceptionError(t *testing.T) {
	ctx := context.Background()
	btc := &binaryThriftCodec{thriftCodec}
	cliMsg := &mockMessage{
		RPCInfoFunc: func() rpcinfo.RPCInfo {
			return newEmptyMethodRPCInfo()
		},
		RPCRoleFunc: func() remote.RPCRole {
			return remote.Server
		},
		MessageTypeFunc: func() remote.MessageType {
			return remote.Exception
		},
	}

	rwbuf := remote.NewReaderWriterBuffer(1024)
	// test data is empty
	err := btc.Marshal(ctx, cliMsg, rwbuf)
	test.Assert(t, err.Error() == "invalid marshal data in rawThriftBinaryCodec: nil")
	cliMsg.DataFunc = func() interface{} {
		return &remote.TransError{}
	}

	// empty method
	err = btc.Marshal(ctx, cliMsg, rwbuf)
	test.Assert(t, err.Error() == "rawThriftBinaryCodec Marshal exception failed, err: empty methodName in thrift Marshal")

	cliMsg.RPCInfoFunc = func() rpcinfo.RPCInfo {
		return newMockRPCInfo()
	}
	err = btc.Marshal(ctx, cliMsg, rwbuf)
	test.Assert(t, err == nil)
	err = btc.Unmarshal(ctx, cliMsg, rwbuf)
	test.Assert(t, err.Error() == "unknown application exception")

	// test server role
	cliMsg.MessageTypeFunc = func() remote.MessageType {
		return remote.Call
	}
	cliMsg.DataFunc = func() interface{} {
		return &Result{
			Success: binaryReqType{},
		}
	}
	err = btc.Marshal(ctx, cliMsg, rwbuf)
	test.Assert(t, err == nil)
}

func newMockRPCInfo() rpcinfo.RPCInfo {
	c := rpcinfo.NewEndpointInfo("", "", nil, nil)
	s := rpcinfo.NewEndpointInfo("", "", nil, nil)
	ink := rpcinfo.NewInvocation("", "mock")
	ri := rpcinfo.NewRPCInfo(c, s, ink, nil, rpcinfo.NewRPCStats())
	return ri
}

func newEmptyMethodRPCInfo() rpcinfo.RPCInfo {
	c := rpcinfo.NewEndpointInfo("", "", nil, nil)
	s := rpcinfo.NewEndpointInfo("", "", nil, nil)
	ink := rpcinfo.NewInvocation("", "")
	ri := rpcinfo.NewRPCInfo(c, s, ink, nil, nil)
	return ri
}

var _ remote.Message = &mockMessage{}

type mockMessage struct {
	RPCInfoFunc         func() rpcinfo.RPCInfo
	ServiceInfoFunc     func() *serviceinfo.ServiceInfo
	SetServiceInfoFunc  func(svcName, methodName string) (*serviceinfo.ServiceInfo, error)
	DataFunc            func() interface{}
	NewDataFunc         func(method string) (ok bool)
	MessageTypeFunc     func() remote.MessageType
	SetMessageTypeFunc  func(remote.MessageType)
	RPCRoleFunc         func() remote.RPCRole
	PayloadLenFunc      func() int
	SetPayloadLenFunc   func(size int)
	TransInfoFunc       func() remote.TransInfo
	TagsFunc            func() map[string]interface{}
	ProtocolInfoFunc    func() remote.ProtocolInfo
	SetProtocolInfoFunc func(remote.ProtocolInfo)
	PayloadCodecFunc    func() remote.PayloadCodec
	SetPayloadCodecFunc func(pc remote.PayloadCodec)
	RecycleFunc         func()
}

func (m *mockMessage) RPCInfo() rpcinfo.RPCInfo {
	if m.RPCInfoFunc != nil {
		return m.RPCInfoFunc()
	}
	return nil
}

func (m *mockMessage) ServiceInfo() (si *serviceinfo.ServiceInfo) {
	if m.ServiceInfoFunc != nil {
		return m.ServiceInfoFunc()
	}
	return
}

func (m *mockMessage) SpecifyServiceInfo(svcName, methodName string) (si *serviceinfo.ServiceInfo, err error) {
	if m.SetServiceInfoFunc != nil {
		return m.SetServiceInfoFunc(svcName, methodName)
	}
	return nil, nil
}

func (m *mockMessage) Data() interface{} {
	if m.DataFunc != nil {
		return m.DataFunc()
	}
	return nil
}

func (m *mockMessage) NewData(method string) (ok bool) {
	if m.NewDataFunc != nil {
		return m.NewDataFunc(method)
	}
	return false
}

func (m *mockMessage) MessageType() (mt remote.MessageType) {
	if m.MessageTypeFunc != nil {
		return m.MessageTypeFunc()
	}
	return
}

func (m *mockMessage) SetMessageType(mt remote.MessageType) {
	if m.SetMessageTypeFunc != nil {
		m.SetMessageTypeFunc(mt)
	}
}

func (m *mockMessage) RPCRole() (r remote.RPCRole) {
	if m.RPCRoleFunc != nil {
		return m.RPCRoleFunc()
	}
	return
}

func (m *mockMessage) PayloadLen() int {
	if m.PayloadLenFunc != nil {
		return m.PayloadLenFunc()
	}
	return 0
}

func (m *mockMessage) SetPayloadLen(size int) {
	if m.SetPayloadLenFunc != nil {
		m.SetPayloadLenFunc(size)
	}
}

func (m *mockMessage) TransInfo() remote.TransInfo {
	if m.TransInfoFunc != nil {
		return m.TransInfoFunc()
	}
	return nil
}

func (m *mockMessage) Tags() map[string]interface{} {
	if m.TagsFunc != nil {
		return m.TagsFunc()
	}
	return nil
}

func (m *mockMessage) ProtocolInfo() (pi remote.ProtocolInfo) {
	if m.ProtocolInfoFunc != nil {
		return m.ProtocolInfoFunc()
	}
	return
}

func (m *mockMessage) SetProtocolInfo(pi remote.ProtocolInfo) {
	if m.SetProtocolInfoFunc != nil {
		m.SetProtocolInfoFunc(pi)
	}
}

func (m *mockMessage) PayloadCodec() remote.PayloadCodec {
	if m.PayloadCodecFunc != nil {
		return m.PayloadCodecFunc()
	}
	return nil
}

func (m *mockMessage) SetPayloadCodec(pc remote.PayloadCodec) {
	if m.SetPayloadCodecFunc != nil {
		m.SetPayloadCodecFunc(pc)
	}
}

func (m *mockMessage) Recycle() {
	if m.RecycleFunc != nil {
		m.RecycleFunc()
	}
}