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

	"github.com/bytedance/gopkg/cloud/metainfo"

	"github.com/cloudwego/kitex/internal/mocks"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/logid"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/metadata"
	"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 TestClientReadMetainfo(t *testing.T) {
	ctx := context.Background()
	ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""), nil, rpcinfo.NewRPCStats())
	msg := remote.NewMessage(nil, mocks.ServiceInfo(), ri, remote.Call, remote.Client)
	hd := map[string]string{
		"hello": "world",
	}
	msg.TransInfo().PutTransStrInfo(hd)

	var err error
	ctx, err = MetainfoClientHandler.ReadMeta(ctx, msg)
	test.Assert(t, err == nil)

	kvs := metainfo.RecvAllBackwardValues(ctx)
	test.Assert(t, len(kvs) == 0)

	ctx = metainfo.WithBackwardValues(context.Background())
	ctx, err = MetainfoClientHandler.ReadMeta(ctx, msg)
	test.Assert(t, err == nil)

	kvs = metainfo.RecvAllBackwardValues(ctx)
	test.Assert(t, len(kvs) == 1)
	test.Assert(t, kvs["hello"] == "world")
}

func TestClientWriteMetainfo(t *testing.T) {
	ri := rpcinfo.NewRPCInfo(nil, nil, rpcinfo.NewInvocation("", ""), nil, rpcinfo.NewRPCStats())
	ctx := context.Background()
	ctx = metainfo.WithValue(ctx, "tk", "tv")
	ctx = metainfo.WithPersistentValue(ctx, "pk", "pv")
	msg := remote.NewMessage(nil, mocks.ServiceInfo(), ri, remote.Call, remote.Client)

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.PurePayload, serviceinfo.Thrift))
	ctx, err := MetainfoClientHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)

	kvs := msg.TransInfo().TransStrInfo()
	test.Assert(t, len(kvs) == 2, kvs)
	test.Assert(t, kvs[metainfo.PrefixTransient+"tk"] == "tv")
	test.Assert(t, kvs[metainfo.PrefixPersistent+"pk"] == "pv")

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.TTHeader, serviceinfo.Thrift))
	_, err = MetainfoClientHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)

	kvs = msg.TransInfo().TransStrInfo()
	test.Assert(t, len(kvs) == 2, kvs)
	test.Assert(t, kvs[metainfo.PrefixTransient+"tk"] == "tv")
	test.Assert(t, kvs[metainfo.PrefixPersistent+"pk"] == "pv")
}

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

	hd := map[string]string{
		"hello":                          "world",
		metainfo.PrefixTransient + "tk":  "tv",
		metainfo.PrefixPersistent + "pk": "pv",
	}
	msg.TransInfo().PutTransStrInfo(hd)

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.PurePayload, serviceinfo.Thrift))
	ctx, err := MetainfoServerHandler.ReadMeta(ctx0, msg)
	tvs := metainfo.GetAllValues(ctx)
	pvs := metainfo.GetAllPersistentValues(ctx)
	test.Assert(t, err == nil)
	test.Assert(t, len(tvs) == 1 && tvs["tk"] == "tv", tvs)
	test.Assert(t, len(pvs) == 1 && pvs["pk"] == "pv")

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.TTHeader, serviceinfo.Thrift))
	ctx, err = MetainfoServerHandler.ReadMeta(ctx0, msg)
	ctx = metainfo.TransferForward(ctx)
	tvs = metainfo.GetAllValues(ctx)
	pvs = metainfo.GetAllPersistentValues(ctx)
	test.Assert(t, err == nil)
	test.Assert(t, len(tvs) == 1 && tvs["tk"] == "tv", tvs)
	test.Assert(t, len(pvs) == 1 && pvs["pk"] == "pv")

	ctx = metainfo.TransferForward(ctx)
	tvs = metainfo.GetAllValues(ctx)
	pvs = metainfo.GetAllPersistentValues(ctx)
	test.Assert(t, len(tvs) == 0, len(tvs))
	test.Assert(t, len(pvs) == 1 && pvs["pk"] == "pv")
}

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

	ctx := context.Background()
	ctx = metainfo.WithBackwardValuesToSend(ctx)
	ctx = metainfo.WithValue(ctx, "tk", "tv")
	ctx = metainfo.WithPersistentValue(ctx, "pk", "pv")
	ok := metainfo.SendBackwardValue(ctx, "bk", "bv")
	test.Assert(t, ok)

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.PurePayload, serviceinfo.Thrift))
	ctx, err := MetainfoServerHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)
	kvs := msg.TransInfo().TransStrInfo()
	test.Assert(t, len(kvs) == 1 && kvs["bk"] == "bv", kvs)

	msg.SetProtocolInfo(remote.NewProtocolInfo(transport.TTHeader, serviceinfo.Thrift))
	_, err = MetainfoServerHandler.WriteMeta(ctx, msg)
	test.Assert(t, err == nil)
	kvs = msg.TransInfo().TransStrInfo()
	test.Assert(t, len(kvs) == 1 && kvs["bk"] == "bv", kvs)
}

func Test_addStreamID(t *testing.T) {
	t.Run("without-stream-log-id", func(t *testing.T) {
		md := metadata.MD{
			transmeta.HTTPStreamLogID: nil,
		}
		ctx := context.Background()
		ctx = addStreamIDToContext(ctx, md)
		logID := logid.GetStreamLogID(ctx)
		test.Assert(t, logID == "", logID) // won't generate a new one
	})

	t.Run("with-stream-log-id", func(t *testing.T) {
		md := metadata.MD{
			transmeta.HTTPStreamLogID: []string{"test"},
		}
		ctx := context.Background()
		ctx = addStreamIDToContext(ctx, md)
		logID := logid.GetStreamLogID(ctx)
		test.Assert(t, logID == "test", logID)
	})
}