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

	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/internal/mocks"
	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/remote"
)

// TestOnActive test server_handler OnActive success
func TestOnActive(t *testing.T) {
	// 1. prepare mock data
	conn := &MockNetpollConn{
		Conn: mocks.Conn{
			RemoteAddrFunc: func() (r net.Addr) {
				return addr
			},
		},
		ReaderFunc: func() (r netpoll.Reader) {
			reader := &MockNetpollReader{
				ReleaseFunc: func() (err error) {
					return nil
				},
			}
			return reader
		},
		WriterFunc: func() (r netpoll.Writer) {
			writer := &MockNetpollWriter{
				FlushFunc: func() (err error) {
					return nil
				},
			}
			return writer
		},
	}
	ctx := context.Background()
	remote.NewTransPipeline(svrTransHdlr)
	_, err := svrTransHdlr.OnActive(ctx, conn)
	test.Assert(t, err == nil, err)
}

// TestOnRead test server_handler OnRead success
func TestOnRead(t *testing.T) {
	// 1. prepare mock data
	var isWriteBufFlushed bool
	var isReaderBufReleased bool
	var isInvoked bool
	conn := &MockNetpollConn{
		Conn: mocks.Conn{
			RemoteAddrFunc: func() (r net.Addr) {
				return addr
			},
		},
		ReaderFunc: func() (r netpoll.Reader) {
			reader := &MockNetpollReader{
				ReleaseFunc: func() (err error) {
					isReaderBufReleased = true
					return nil
				},
			}
			return reader
		},
		WriterFunc: func() (r netpoll.Writer) {
			writer := &MockNetpollWriter{
				FlushFunc: func() (err error) {
					isWriteBufFlushed = true
					return nil
				},
			}
			return writer
		},
	}

	if setter, ok := svrTransHdlr.(remote.InvokeHandleFuncSetter); ok {
		setter.SetInvokeHandleFunc(func(ctx context.Context, req, resp interface{}) (err error) {
			isInvoked = true
			return nil
		})
	}
	remote.NewTransPipeline(svrTransHdlr)

	// 2. test
	ctx := context.Background()
	ctx, err := svrTransHdlr.OnActive(ctx, conn)
	test.Assert(t, err == nil, err)

	err = svrTransHdlr.OnRead(ctx, conn)
	test.Assert(t, err == nil, err)
	test.Assert(t, isReaderBufReleased)
	test.Assert(t, isWriteBufFlushed)
	test.Assert(t, isInvoked)
}

// TestInvokeErr test server_handler invoke err
func TestInvokeErr(t *testing.T) {
	// 1. prepare mock data
	var isWriteBufFlushed bool
	var isReaderBufReleased bool
	var isInvoked bool
	conn := &MockNetpollConn{
		Conn: mocks.Conn{
			RemoteAddrFunc: func() (r net.Addr) {
				return addr
			},
			CloseFunc: func() (e error) {
				return nil
			},
		},
		ReaderFunc: func() (r netpoll.Reader) {
			reader := &MockNetpollReader{
				ReleaseFunc: func() (err error) {
					isReaderBufReleased = true
					return nil
				},
			}
			return reader
		},
		WriterFunc: func() (r netpoll.Writer) {
			writer := &MockNetpollWriter{
				FlushFunc: func() (err error) {
					isWriteBufFlushed = true
					return nil
				},
			}
			return writer
		},
		IsActiveFunc: func() (r bool) {
			return true
		},
	}
	remote.NewTransPipeline(svrTransHdlr)

	// mock invoke err
	if setter, ok := svrTransHdlr.(remote.InvokeHandleFuncSetter); ok {
		setter.SetInvokeHandleFunc(func(ctx context.Context, req, resp interface{}) (err error) {
			isInvoked = true
			return errors.New("mock invoke test")
		})
	}

	// 2. test
	ctx := context.Background()
	ctx, err := svrTransHdlr.OnActive(ctx, conn)
	test.Assert(t, err == nil, err)

	err = svrTransHdlr.OnRead(ctx, conn)
	test.Assert(t, err == nil, err)
	test.Assert(t, isReaderBufReleased)
	test.Assert(t, isWriteBufFlushed)
	test.Assert(t, isInvoked)
}

// TestPanicAfterRead test server_handler not panic after read
func TestPanicAfterRead(t *testing.T) {
	// 1. prepare mock data
	var isWriteBufFlushed bool
	var isReaderBufReleased bool
	var isInvoked bool
	var isClosed bool
	conn := &MockNetpollConn{
		Conn: mocks.Conn{
			RemoteAddrFunc: func() (r net.Addr) {
				return addr
			},
			CloseFunc: func() (e error) {
				isClosed = true
				return nil
			},
		},
		ReaderFunc: func() (r netpoll.Reader) {
			reader := &MockNetpollReader{
				ReleaseFunc: func() (err error) {
					isReaderBufReleased = true
					return nil
				},
			}
			return reader
		},
		WriterFunc: func() (r netpoll.Writer) {
			writer := &MockNetpollWriter{
				FlushFunc: func() (err error) {
					isWriteBufFlushed = true
					return nil
				},
			}
			return writer
		},
		IsActiveFunc: func() (r bool) {
			return true
		},
	}
	// pipeline nil panic
	svrTransHdlr.SetPipeline(nil)

	// 2. test
	ctx := context.Background()
	ctx, err := svrTransHdlr.OnActive(ctx, conn)
	test.Assert(t, err == nil, err)

	err = svrTransHdlr.OnRead(ctx, conn)
	test.Assert(t, err != nil, err)
	test.Assert(t, !isReaderBufReleased)
	test.Assert(t, !isWriteBufFlushed)
	test.Assert(t, !isInvoked)
	test.Assert(t, !isClosed)
}

// TestNoMethodInfo test server_handler without method info success
func TestNoMethodInfo(t *testing.T) {
	// 1. prepare mock data
	var isWriteBufFlushed bool
	var isReaderBufReleased bool
	var isClosed bool
	conn := &MockNetpollConn{
		Conn: mocks.Conn{
			RemoteAddrFunc: func() (r net.Addr) {
				return addr
			},
			CloseFunc: func() (e error) {
				isClosed = true
				return nil
			},
		},
		ReaderFunc: func() (r netpoll.Reader) {
			reader := &MockNetpollReader{
				ReleaseFunc: func() (err error) {
					isReaderBufReleased = true
					return nil
				},
			}
			return reader
		},
		WriterFunc: func() (r netpoll.Writer) {
			writer := &MockNetpollWriter{
				FlushFunc: func() (err error) {
					isWriteBufFlushed = true
					return nil
				},
			}
			return writer
		},
		IsActiveFunc: func() (r bool) {
			return true
		},
	}
	remote.NewTransPipeline(svrTransHdlr)
	svcInfo := svrOpt.TargetSvcInfo
	delete(svcInfo.Methods, method)

	// 2. test
	ctx := context.Background()
	ctx, err := svrTransHdlr.OnActive(ctx, conn)
	test.Assert(t, err == nil, err)

	err = svrTransHdlr.OnRead(ctx, conn)
	test.Assert(t, err != nil, err)
	test.Assert(t, isReaderBufReleased)
	test.Assert(t, isWriteBufFlushed)
	test.Assert(t, !isClosed)
}