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

import (
	"context"
	"strconv"
	"testing"
	"time"

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

func TestIsTTHeader(t *testing.T) {
	t.Run("with ttheader", func(t *testing.T) {
		ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""), nil, rpcinfo.NewRPCStats())
		msg := remote.NewMessage(nil, mocks.ServiceInfo(), ri, remote.Call, remote.Server)
		msg.SetProtocolInfo(remote.NewProtocolInfo(transport.TTHeader, serviceinfo.Thrift))
		test.Assert(t, isTTHeader(msg))
	})
	t.Run("with ttheader framed", func(t *testing.T) {
		ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""), nil, rpcinfo.NewRPCStats())
		msg := remote.NewMessage(nil, mocks.ServiceInfo(), ri, remote.Call, remote.Server)
		msg.SetProtocolInfo(remote.NewProtocolInfo(transport.TTHeaderFramed, serviceinfo.Thrift))
		test.Assert(t, isTTHeader(msg))
	})
}

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

	cfg := rpcinfo.NewRPCConfig()
	cfgMutable := rpcinfo.AsMutableRPCConfig(cfg)
	cfgMutable.SetRPCTimeout(time.Millisecond * 100)
	cfgMutable.LockConfig(rpcinfo.BitRPCTimeout)

	fromInfo := rpcinfo.NewEndpointInfo("fromServiceName", "fromMethod", nil, nil)
	toInfo := rpcinfo.NewEndpointInfo("toServiceName", "toMethod", nil, nil)
	ri := rpcinfo.NewRPCInfo(fromInfo, toInfo, rpcinfo.NewInvocation("", ""), cfg, rpcinfo.NewRPCStats())
	msg := remote.NewMessage(nil, mocks.ServiceInfo(), ri, remote.Call, remote.Client)

	// pure paylod, no effect
	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.PurePayload, serviceinfo.Thrift))
	ctx, err := ClientTTHeaderHandler.WriteMeta(ctx, msg)
	kvs := msg.TransInfo().TransIntInfo()
	test.Assert(t, err == nil)
	test.Assert(t, len(kvs) == 0)
	strKvs := msg.TransInfo().TransStrInfo()
	test.Assert(t, len(strKvs) == 0)

	// ttheader
	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.TTHeader, serviceinfo.Thrift))
	_, err = ClientTTHeaderHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)
	kvs = msg.TransInfo().TransIntInfo()
	test.Assert(t, len(kvs) > 0)
	test.Assert(t, kvs[transmeta.FromService] == fromInfo.ServiceName())
	test.Assert(t, kvs[transmeta.FromMethod] == fromInfo.Method())
	test.Assert(t, kvs[transmeta.ToService] == toInfo.ServiceName())
	test.Assert(t, kvs[transmeta.ToMethod] == toInfo.Method())
	test.Assert(t, kvs[transmeta.MsgType] == strconv.Itoa(int(remote.Call)))
	test.Assert(t, kvs[transmeta.TransportType] == unframedTransportType)
	test.Assert(t, kvs[transmeta.RPCTimeout] == "100")
	strKvs = msg.TransInfo().TransStrInfo()
	test.Assert(t, len(strKvs) == 1)
	test.Assert(t, strKvs[transmeta.HeaderIDLServiceName] == "")
}

func TestTTHeaderServerReadMetainfo(t *testing.T) {
	ctx := context.Background()
	ri := rpcinfo.NewRPCInfo(rpcinfo.EmptyEndpointInfo(), nil, rpcinfo.NewInvocation("", ""),
		rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
	msg := remote.NewMessage(nil, mocks.ServiceInfo(), ri, remote.Call, remote.Client)

	hd := map[uint16]string{
		transmeta.FromService: "fromService",
		transmeta.FromMethod:  "fromMethod",
		transmeta.RPCTimeout:  "100",
	}
	msg.TransInfo().PutTransIntInfo(hd)

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.PurePayload, serviceinfo.Thrift))
	_, err := ServerTTHeaderHandler.ReadMeta(ctx, msg)
	test.Assert(t, err == nil)
	fromEndPoint := msg.RPCInfo().From()
	test.Assert(t, fromEndPoint.ServiceName() == "")
	test.Assert(t, fromEndPoint.Method() == "")

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.TTHeader, serviceinfo.Thrift))
	_, err = ServerTTHeaderHandler.ReadMeta(ctx, msg)
	test.Assert(t, err == nil)
	test.Assert(t, msg.RPCInfo().From().ServiceName() == hd[transmeta.FromService])
	test.Assert(t, msg.RPCInfo().From().Method() == hd[transmeta.FromMethod])
	test.Assert(t, ri.Config().RPCTimeout() == 100*time.Millisecond)
}

func TestTTHeaderServerWriteMetainfo(t *testing.T) {
	ctx := context.Background()
	ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""),
		rpcinfo.NewRPCConfig(), rpcinfo.NewRPCStats())
	msg := remote.NewMessage(nil, mocks.ServiceInfo(), ri, remote.Call, remote.Client)

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.PurePayload, serviceinfo.Thrift))
	_, err := ServerTTHeaderHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)
	kvs := msg.TransInfo().TransIntInfo()
	test.Assert(t, len(kvs) == 0)

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.TTHeader, serviceinfo.Thrift))
	_, err = ServerTTHeaderHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)
	kvs = msg.TransInfo().TransIntInfo()
	test.Assert(t, kvs[transmeta.MsgType] == strconv.Itoa(int(remote.Call)))
}