github.com/cloudwego/kitex@v0.9.0/pkg/utils/kitexutil/kitexutil_test.go

/*
 * Copyright 2023 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 kitexutil

import (
	"context"
	"net"
	"os"
	"reflect"
	"testing"

	mocks "github.com/cloudwego/kitex/internal/mocks/thrift/fast"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/utils"
	"github.com/cloudwego/kitex/transport"
)

var (
	testRi         rpcinfo.RPCInfo
	testCtx        context.Context
	panicCtx       context.Context
	caller         = "kitexutil.from.service"
	callee         = "kitexutil.to.service"
	idlServiceName = "MockService"
	fromAddr       = utils.NewNetAddr("test", "127.0.0.1:12345")
	fromMethod     = "from_method"
	method         = "method"
	tp             = transport.TTHeader
)

func TestMain(m *testing.M) {
	testRi = buildRPCInfo()

	testCtx = context.Background()
	testCtx = rpcinfo.NewCtxWithRPCInfo(testCtx, testRi)
	panicCtx = rpcinfo.NewCtxWithRPCInfo(context.Background(), &panicRPCInfo{})

	os.Exit(m.Run())
}

func TestGetCaller(t *testing.T) {
	type args struct {
		ctx context.Context
	}
	tests := []struct {
		name  string
		args  args
		want  string
		want1 bool
	}{
		{name: "Success", args: args{testCtx}, want: caller, want1: true},
		{name: "Failure", args: args{context.Background()}, want: "", want1: false},
		{name: "Panic recovered", args: args{panicCtx}, want: "", want1: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, got1 := GetCaller(tt.args.ctx)
			if got != tt.want {
				t.Errorf("GetCaller() got = %v, want %v", got, tt.want)
			}
			if got1 != tt.want1 {
				t.Errorf("GetCaller() got1 = %v, want %v", got1, tt.want1)
			}
		})
	}
}

func TestGetCallerAddr(t *testing.T) {
	type args struct {
		ctx context.Context
	}
	tests := []struct {
		name  string
		args  args
		want  net.Addr
		want1 bool
	}{
		{name: "Success", args: args{testCtx}, want: fromAddr, want1: true},
		{name: "Failure", args: args{context.Background()}, want: nil, want1: false},
		{name: "Panic recovered", args: args{panicCtx}, want: nil, want1: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, got1 := GetCallerAddr(tt.args.ctx)
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetCallerAddr() got = %v, want %v", got, tt.want)
			}
			if got1 != tt.want1 {
				t.Errorf("GetCallerAddr() got1 = %v, want %v", got1, tt.want1)
			}
		})
	}
}

func TestGetCallerIP(t *testing.T) {
	ip, ok := GetCallerIP(testCtx)
	test.Assert(t, ok)
	test.Assert(t, ip == "127.0.0.1", ip)

	ri := buildRPCInfo()
	rpcinfo.AsMutableEndpointInfo(ri.From()).SetAddress(utils.NewNetAddr("test", "127.0.0.1"))
	ip, ok = GetCallerIP(rpcinfo.NewCtxWithRPCInfo(context.Background(), ri))
	test.Assert(t, ok)
	test.Assert(t, ip == "127.0.0.1", ip)

	ip, ok = GetCallerIP(context.Background())
	test.Assert(t, !ok)
	test.Assert(t, ip == "", ip)

	ip, ok = GetCallerIP(panicCtx)
	test.Assert(t, !ok)
	test.Assert(t, ip == "", ip)

	rpcinfo.AsMutableEndpointInfo(ri.From()).SetAddress(utils.NewNetAddr("test", ""))
	ip, ok = GetCallerIP(rpcinfo.NewCtxWithRPCInfo(context.Background(), ri))
	test.Assert(t, !ok)
	test.Assert(t, ip == "", ip)
}

func TestGetMethod(t *testing.T) {
	type args struct {
		ctx context.Context
	}
	tests := []struct {
		name  string
		args  args
		want  string
		want1 bool
	}{
		{name: "Success", args: args{testCtx}, want: method, want1: true},
		{name: "Failure", args: args{context.Background()}, want: "", want1: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, got1 := GetMethod(tt.args.ctx)
			if got != tt.want {
				t.Errorf("GetMethod() got = %v, want %v", got, tt.want)
			}
			if got1 != tt.want1 {
				t.Errorf("GetMethod() got1 = %v, want %v", got1, tt.want1)
			}
		})
	}
}

func TestGetCallerHandlerMethod(t *testing.T) {
	type args struct {
		ctx context.Context
	}
	tests := []struct {
		name  string
		args  args
		want  string
		want1 bool
	}{
		{name: "Success", args: args{testCtx}, want: fromMethod, want1: true},
		{name: "Failure", args: args{context.Background()}, want: "", want1: false},
		{name: "Panic recovered", args: args{panicCtx}, want: "", want1: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, got1 := GetCallerHandlerMethod(tt.args.ctx)
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetCallerHandlerMethod() got = %v, want %v", got, tt.want)
			}
			if got1 != tt.want1 {
				t.Errorf("GetCallerHandlerMethod() got1 = %v, want %v", got1, tt.want1)
			}
		})
	}
}

func TestGetIDLServiceName(t *testing.T) {
	type args struct {
		ctx context.Context
	}
	tests := []struct {
		name  string
		args  args
		want  string
		want1 bool
	}{
		{name: "Success", args: args{testCtx}, want: idlServiceName, want1: true},
		{name: "Failure", args: args{context.Background()}, want: "", want1: false},
		{name: "Panic recovered", args: args{panicCtx}, want: "", want1: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, got1 := GetIDLServiceName(tt.args.ctx)
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetCallerHandlerMethod() got = %v, want %v", got, tt.want)
			}
			if got1 != tt.want1 {
				t.Errorf("GetCallerHandlerMethod() got1 = %v, want %v", got1, tt.want1)
			}
		})
	}
}

func TestGetRPCInfo(t *testing.T) {
	type args struct {
		ctx context.Context
	}
	tests := []struct {
		name  string
		args  args
		want  rpcinfo.RPCInfo
		want1 bool
	}{
		{name: "Success", args: args{testCtx}, want: testRi, want1: true},
		{name: "Failure", args: args{context.Background()}, want: nil, want1: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, got1 := GetRPCInfo(tt.args.ctx)
			if !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetRPCInfo() got = %v, want %v", got, tt.want)
			}
			if got1 != tt.want1 {
				t.Errorf("GetRPCInfo() got1 = %v, want %v", got1, tt.want1)
			}
		})
	}
}

func TestGetCtxTransportProtocol(t *testing.T) {
	type args struct {
		ctx context.Context
	}
	tests := []struct {
		name  string
		args  args
		want  string
		want1 bool
	}{
		{name: "Success", args: args{testCtx}, want: tp.String(), want1: true},
		{name: "Failure", args: args{context.Background()}, want: "", want1: false},
		{name: "Panic recovered", args: args{panicCtx}, want: "", want1: false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, got1 := GetTransportProtocol(tt.args.ctx)
			if got != tt.want {
				t.Errorf("GetTransportProtocol() got = %v, want %v", got, tt.want)
			}
			if got1 != tt.want1 {
				t.Errorf("GetTransportProtocol() got1 = %v, want %v", got1, tt.want1)
			}
		})
	}
}

func TestGetRealRequest(t *testing.T) {
	req := &mocks.MockReq{}
	arg := &mocks.MockTestArgs{Req: req}
	type args struct {
		req interface{}
	}
	tests := []struct {
		name string
		args args
		want interface{}
	}{
		{name: "success", args: args{arg}, want: req},
		{name: "nil input", args: args{nil}, want: nil},
		{name: "wrong interface", args: args{req}, want: nil},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := GetRealReqFromKitexArgs(tt.args.req); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetRealReqFromKitexArgs() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestGetRealResponse(t *testing.T) {
	success := "success"
	result := &mocks.MockTestResult{Success: &success}
	type args struct {
		resp interface{}
	}
	tests := []struct {
		name string
		args args
		want interface{}
	}{
		{name: "success", args: args{result}, want: &success},
		{name: "nil input", args: args{nil}, want: nil},
		{name: "wrong interface", args: args{success}, want: nil},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := GetRealRespFromKitexResult(tt.args.resp); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("GetRealRespFromKitexResult() = %v, want %v", got, tt.want)
			}
		})
	}
}

func buildRPCInfo() rpcinfo.RPCInfo {
	from := rpcinfo.NewEndpointInfo(caller, fromMethod, fromAddr, nil)
	to := rpcinfo.NewEndpointInfo(callee, method, nil, nil)
	ink := rpcinfo.NewInvocation(idlServiceName, method)
	config := rpcinfo.NewRPCConfig()
	config.(rpcinfo.MutableRPCConfig).SetTransportProtocol(tp)

	stats := rpcinfo.NewRPCStats()
	ri := rpcinfo.NewRPCInfo(from, to, ink, config, stats)
	return ri
}

type panicRPCInfo struct{}

func (m *panicRPCInfo) From() rpcinfo.EndpointInfo     { panic("Panic when invoke From") }
func (m *panicRPCInfo) To() rpcinfo.EndpointInfo       { panic("Panic when invoke To") }
func (m *panicRPCInfo) Invocation() rpcinfo.Invocation { panic("Panic when invoke Invocation") }
func (m *panicRPCInfo) Config() rpcinfo.RPCConfig      { panic("Panic when invoke Config") }
func (m *panicRPCInfo) Stats() rpcinfo.RPCStats        { panic("Panic when invoke Stats") }