github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/nphttp2/server_handler_test.go

/*
 * Copyright 2022 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 nphttp2

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/grpc"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/serviceinfo"
	"github.com/cloudwego/kitex/pkg/streaming"
	"github.com/cloudwego/kitex/transport"
)

func TestServerHandler(t *testing.T) {
	// init
	opt := newMockServerOption()
	msg := newMockNewMessage()
	msg.ProtocolInfoFunc = func() remote.ProtocolInfo {
		return remote.NewProtocolInfo(transport.PurePayload, serviceinfo.Protobuf)
	}
	msg.RPCInfoFunc = func() rpcinfo.RPCInfo {
		return newMockRPCInfo()
	}
	npConn := newMockNpConn(mockAddr0)
	npConn.mockSettingFrame()
	tr, err := newMockServerTransport(npConn)
	test.Assert(t, err == nil, err)
	s := grpc.CreateStream(1, func(i int) {})
	serverConn := newServerConn(tr, s)
	defer serverConn.Close()

	// test NewTransHandler()
	handler, err := NewSvrTransHandlerFactory().NewTransHandler(opt)
	test.Assert(t, err == nil, err)

	// test Read()
	// mock grpc encoded msg data into stream recv buffer
	newMockStreamRecvHelloRequest(s)
	ctx, err := handler.Read(context.Background(), serverConn, msg)
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err == nil, err)

	// test write()
	ctx, err = handler.Write(context.Background(), serverConn, msg)
	test.Assert(t, ctx != nil, ctx)
	test.Assert(t, err == nil, err)

	// test SetInvokeHandleFunc()
	svrHdl := handler.(*svrTransHandler)
	svrHdl.SetInvokeHandleFunc(func(ctx context.Context, req, resp interface{}) (err error) {
		return nil
	})

	// mock a setting frame to pass the server side preface check
	npConn.mockSettingFrame()
	// mock a headerFrame so onRead() can start working
	npConn.mockMetaHeaderFrame()
	go func() {
		// test OnActive()
		ctx, err := handler.OnActive(newMockCtxWithRPCInfo(), npConn)
		test.Assert(t, err == nil, err)

		handler.OnRead(ctx, npConn)

		// test OnInactive()
		handler.OnInactive(ctx, npConn)
		test.Assert(t, err == nil, err)
	}()

	// sleep 50 mills so server can handle metaHeader frame
	time.Sleep(time.Millisecond * 50)

	// test OnError()
	handler.OnError(context.Background(), context.Canceled, npConn)

	// test SetPipeline()
	handler.SetPipeline(nil)
}

type mockStream struct {
	streaming.Stream
	recv func(msg interface{}) error
	send func(msg interface{}) error
}

func (s *mockStream) RecvMsg(m interface{}) error {
	return s.recv(m)
}

func (s *mockStream) SendMsg(m interface{}) error {
	return s.send(m)
}

func Test_invokeStreamUnaryHandler(t *testing.T) {
	t.Run("recv err", func(t *testing.T) {
		expectedErr := errors.New("mock err")
		s := &mockStream{
			recv: func(msg interface{}) error {
				return expectedErr
			},
		}
		var newArgsCalled, newResultCalled bool
		mi := serviceinfo.NewMethodInfo(nil,
			func() interface{} {
				newArgsCalled = true
				return nil
			},
			func() interface{} {
				newResultCalled = true
				return nil
			},
			false,
		)
		hdl := func(ctx context.Context, req, resp interface{}) (err error) {
			return nil
		}

		err := invokeStreamUnaryHandler(context.Background(), s, mi, hdl, nil)

		test.Assert(t, err == expectedErr, err)
		test.Assert(t, newArgsCalled)
		test.Assert(t, newResultCalled)
	})
	t.Run("handler err", func(t *testing.T) {
		expectedErr := errors.New("mock err")
		var newArgsCalled, newResultCalled, handlerCalled, recvCalled bool
		s := &mockStream{
			recv: func(msg interface{}) error {
				recvCalled = true
				return nil
			},
		}
		mi := serviceinfo.NewMethodInfo(nil,
			func() interface{} {
				newArgsCalled = true
				return nil
			},
			func() interface{} {
				newResultCalled = true
				return nil
			},
			false,
		)
		hdl := func(ctx context.Context, req, resp interface{}) (err error) {
			handlerCalled = true
			return expectedErr
		}

		err := invokeStreamUnaryHandler(context.Background(), s, mi, hdl, nil)

		test.Assert(t, err == expectedErr, err)
		test.Assert(t, recvCalled)
		test.Assert(t, newArgsCalled)
		test.Assert(t, newResultCalled)
		test.Assert(t, handlerCalled)
	})

	t.Run("biz err", func(t *testing.T) {
		expectedErr := kerrors.NewBizStatusError(100, "mock biz error")
		var newArgsCalled, newResultCalled, handlerCalled, recvCalled, sendCalled bool
		s := &mockStream{
			recv: func(msg interface{}) error {
				recvCalled = true
				return nil
			},
			send: func(msg interface{}) error {
				sendCalled = true
				return nil
			},
		}
		mi := serviceinfo.NewMethodInfo(nil,
			func() interface{} {
				newArgsCalled = true
				return nil
			},
			func() interface{} {
				newResultCalled = true
				return nil
			},
			false,
		)
		hdl := func(ctx context.Context, req, resp interface{}) (err error) {
			handlerCalled = true
			return nil
		}

		ivk := rpcinfo.NewInvocation("test", "test")
		ivk.SetBizStatusErr(expectedErr)
		ri := rpcinfo.NewRPCInfo(nil, nil, ivk, nil, nil)

		err := invokeStreamUnaryHandler(context.Background(), s, mi, hdl, ri)

		test.Assert(t, err == nil, err)
		test.Assert(t, recvCalled)
		test.Assert(t, newArgsCalled)
		test.Assert(t, newResultCalled)
		test.Assert(t, handlerCalled)
		test.Assert(t, !sendCalled)
	})

	t.Run("send err", func(t *testing.T) {
		expectedErr := errors.New("mock err")
		var newArgsCalled, newResultCalled, handlerCalled, recvCalled, sendCalled bool
		s := &mockStream{
			recv: func(msg interface{}) error {
				recvCalled = true
				return nil
			},
			send: func(msg interface{}) error {
				sendCalled = true
				return expectedErr
			},
		}
		mi := serviceinfo.NewMethodInfo(nil,
			func() interface{} {
				newArgsCalled = true
				return nil
			},
			func() interface{} {
				newResultCalled = true
				return nil
			},
			false,
		)
		hdl := func(ctx context.Context, req, resp interface{}) (err error) {
			handlerCalled = true
			return nil
		}

		ivk := rpcinfo.NewInvocation("test", "test")
		ri := rpcinfo.NewRPCInfo(nil, nil, ivk, nil, nil)

		err := invokeStreamUnaryHandler(context.Background(), s, mi, hdl, ri)

		test.Assert(t, err == expectedErr, err)
		test.Assert(t, recvCalled)
		test.Assert(t, newArgsCalled)
		test.Assert(t, newResultCalled)
		test.Assert(t, handlerCalled)
		test.Assert(t, sendCalled)
	})
}