github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpoll/trans_server_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 netpoll

import (
	"context"
	"errors"
	"net"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/golang/mock/gomock"

	"github.com/cloudwego/kitex/internal/mocks"
	mocksremote "github.com/cloudwego/kitex/internal/mocks/remote"
	"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"
)

var (
	svrTransHdlr remote.ServerTransHandler
	rwTimeout    = time.Second
	addrStr      = "test addr"
	addr         = utils.NewNetAddr("tcp", addrStr)
	method       = "mock"
	transSvr     *transServer
	svrOpt       *remote.ServerOption
)

func TestMain(m *testing.M) {
	svcInfo := mocks.ServiceInfo()
	svrOpt = &remote.ServerOption{
		InitOrResetRPCInfoFunc: func(ri rpcinfo.RPCInfo, addr net.Addr) rpcinfo.RPCInfo {
			fromInfo := rpcinfo.EmptyEndpointInfo()
			rpcCfg := rpcinfo.NewRPCConfig()
			mCfg := rpcinfo.AsMutableRPCConfig(rpcCfg)
			mCfg.SetReadWriteTimeout(rwTimeout)
			ink := rpcinfo.NewInvocation("", method)
			rpcStat := rpcinfo.NewRPCStats()
			nri := rpcinfo.NewRPCInfo(fromInfo, nil, ink, rpcCfg, rpcStat)
			rpcinfo.AsMutableEndpointInfo(nri.From()).SetAddress(addr)
			return nri
		},
		Codec: &MockCodec{
			EncodeFunc: nil,
			DecodeFunc: func(ctx context.Context, msg remote.Message, in remote.ByteBuffer) error {
				msg.SpecifyServiceInfo(mocks.MockServiceName, mocks.MockMethod)
				return nil
			},
		},
		SvcSearchMap: map[string]*serviceinfo.ServiceInfo{
			remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockMethod):          svcInfo,
			remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockExceptionMethod): svcInfo,
			remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockErrorMethod):     svcInfo,
			remote.BuildMultiServiceKey(mocks.MockServiceName, mocks.MockOnewayMethod):    svcInfo,
			mocks.MockMethod:          svcInfo,
			mocks.MockExceptionMethod: svcInfo,
			mocks.MockErrorMethod:     svcInfo,
			mocks.MockOnewayMethod:    svcInfo,
		},
		TargetSvcInfo: svcInfo,
		TracerCtl:     &rpcinfo.TraceController{},
	}
	svrTransHdlr, _ = newSvrTransHandler(svrOpt)
	transSvr = NewTransServerFactory().NewTransServer(svrOpt, svrTransHdlr).(*transServer)

	os.Exit(m.Run())
}

// TestCreateListener test trans_server CreateListener success
func TestCreateListener(t *testing.T) {
	// tcp init
	addrStr := "127.0.0.1:9091"
	addr = utils.NewNetAddr("tcp", addrStr)

	// test
	ln, err := transSvr.CreateListener(addr)
	test.Assert(t, err == nil, err)
	test.Assert(t, ln.Addr().String() == addrStr)
	ln.Close()

	// uds init
	addrStr = "server.addr"
	addr, err = net.ResolveUnixAddr("unix", addrStr)
	test.Assert(t, err == nil, err)

	// test
	ln, err = transSvr.CreateListener(addr)
	test.Assert(t, err == nil, err)
	test.Assert(t, ln.Addr().String() == addrStr)
	ln.Close()
}

// TestBootStrap test trans_server BootstrapServer success
func TestBootStrap(t *testing.T) {
	// tcp init
	addrStr := "127.0.0.1:9092"
	addr = utils.NewNetAddr("tcp", addrStr)

	// test
	ln, err := transSvr.CreateListener(addr)
	test.Assert(t, err == nil, err)
	test.Assert(t, ln.Addr().String() == addrStr)

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		err = transSvr.BootstrapServer(ln)
		test.Assert(t, err == nil, err)
		wg.Done()
	}()
	time.Sleep(10 * time.Millisecond)

	transSvr.Shutdown()
	wg.Wait()
}

// TestOnConnActive test trans_server onConnActive success
func TestConnOnActive(t *testing.T) {
	// 1. prepare mock data
	conn := &MockNetpollConn{
		SetReadTimeoutFunc: func(timeout time.Duration) (e error) {
			return nil
		},
		Conn: mocks.Conn{
			RemoteAddrFunc: func() (r net.Addr) {
				return addr
			},
		},
	}

	// 2. test
	connCount := 100
	for i := 0; i < connCount; i++ {
		transSvr.onConnActive(conn)
	}
	ctx := context.Background()

	currConnCount := transSvr.ConnCount()
	test.Assert(t, currConnCount.Value() == connCount)

	for i := 0; i < connCount; i++ {
		transSvr.onConnInactive(ctx, conn)
	}

	currConnCount = transSvr.ConnCount()
	test.Assert(t, currConnCount.Value() == 0)
}

// TestOnConnActivePanic test panic recover when panic happen in OnActive
func TestConnOnActiveAndOnInactivePanic(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer func() {
		ctrl.Finish()
	}()

	inboundHandler := mocksremote.NewMockInboundHandler(ctrl)
	transPl := remote.NewTransPipeline(svrTransHdlr)
	transPl.AddInboundHandler(inboundHandler)
	transSvrWithPl := NewTransServerFactory().NewTransServer(svrOpt, transPl).(*transServer)
	conn := &MockNetpollConn{
		Conn: mocks.Conn{
			RemoteAddrFunc: func() (r net.Addr) {
				return addr
			},
		},
	}

	// test1: recover OnActive panic
	inboundHandler.EXPECT().OnActive(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, conn net.Conn) (context.Context, error) {
		panic("mock panic")
	})
	transSvrWithPl.onConnActive(conn)

	// test2: recover OnInactive panic
	inboundHandler.EXPECT().OnInactive(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, conn net.Conn) (context.Context, error) {
		panic("mock panic")
	})
	transSvrWithPl.onConnInactive(context.Background(), conn)
}

// TestOnConnRead test trans_server onConnRead success
func TestConnOnRead(t *testing.T) {
	// 1. prepare mock data
	var isClosed bool
	conn := &MockNetpollConn{
		Conn: mocks.Conn{
			RemoteAddrFunc: func() (r net.Addr) {
				return addr
			},
			CloseFunc: func() (e error) {
				isClosed = true
				return nil
			},
		},
	}
	mockErr := errors.New("mock error")
	transSvr.transHdlr = &mocks.MockSvrTransHandler{
		OnReadFunc: func(ctx context.Context, conn net.Conn) error {
			return mockErr
		},
		Opt: transSvr.opt,
	}

	// 2. test
	err := transSvr.onConnRead(context.Background(), conn)
	test.Assert(t, err == nil, err)
	test.Assert(t, isClosed)
}