github.com/cloudwego/hertz@v0.9.3/pkg/network/netpoll/transport_test.go (about)

     1  // Copyright 2023 CloudWeGo Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  
    16  //go:build !windows
    17  // +build !windows
    18  
    19  package netpoll
    20  
    21  import (
    22  	"context"
    23  	"net"
    24  	"sync/atomic"
    25  	"syscall"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/cloudwego/hertz/pkg/common/config"
    30  	"github.com/cloudwego/hertz/pkg/common/test/assert"
    31  	"github.com/cloudwego/hertz/pkg/network"
    32  	"golang.org/x/sys/unix"
    33  )
    34  
    35  func TestTransport(t *testing.T) {
    36  	const nw = "tcp"
    37  	const addr = "localhost:10103"
    38  	t.Run("TestDefault", func(t *testing.T) {
    39  		var onConnFlag, onAcceptFlag, onDataFlag int32
    40  		transporter := NewTransporter(&config.Options{
    41  			Addr:    addr,
    42  			Network: nw,
    43  			OnConnect: func(ctx context.Context, conn network.Conn) context.Context {
    44  				atomic.StoreInt32(&onConnFlag, 1)
    45  				return ctx
    46  			},
    47  			WriteTimeout: time.Second,
    48  			OnAccept: func(conn net.Conn) context.Context {
    49  				atomic.StoreInt32(&onAcceptFlag, 1)
    50  				return context.Background()
    51  			},
    52  		})
    53  		go transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
    54  			atomic.StoreInt32(&onDataFlag, 1)
    55  			return nil
    56  		})
    57  		defer transporter.Close()
    58  		time.Sleep(100 * time.Millisecond)
    59  
    60  		dial := NewDialer()
    61  		conn, err := dial.DialConnection(nw, addr, time.Second, nil)
    62  		assert.Nil(t, err)
    63  		_, err = conn.Write([]byte("123"))
    64  		assert.Nil(t, err)
    65  		time.Sleep(100 * time.Millisecond)
    66  
    67  		assert.Assert(t, atomic.LoadInt32(&onConnFlag) == 1)
    68  		assert.Assert(t, atomic.LoadInt32(&onAcceptFlag) == 1)
    69  		assert.Assert(t, atomic.LoadInt32(&onDataFlag) == 1)
    70  	})
    71  
    72  	t.Run("TestSenseClientDisconnection", func(t *testing.T) {
    73  		var onReqFlag int32
    74  		transporter := NewTransporter(&config.Options{
    75  			Addr:                     addr,
    76  			Network:                  nw,
    77  			SenseClientDisconnection: true,
    78  		})
    79  
    80  		go transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
    81  			atomic.StoreInt32(&onReqFlag, 1)
    82  			time.Sleep(100 * time.Millisecond)
    83  			assert.DeepEqual(t, context.Canceled, ctx.Err())
    84  			return nil
    85  		})
    86  		defer transporter.Close()
    87  		time.Sleep(100 * time.Millisecond)
    88  
    89  		dial := NewDialer()
    90  		conn, err := dial.DialConnection(nw, addr, time.Second, nil)
    91  		assert.Nil(t, err)
    92  		_, err = conn.Write([]byte("123"))
    93  		assert.Nil(t, err)
    94  		err = conn.Close()
    95  		assert.Nil(t, err)
    96  		time.Sleep(100 * time.Millisecond)
    97  
    98  		assert.Assert(t, atomic.LoadInt32(&onReqFlag) == 1)
    99  	})
   100  
   101  	t.Run("TestListenConfig", func(t *testing.T) {
   102  		listenCfg := &net.ListenConfig{Control: func(network, address string, c syscall.RawConn) error {
   103  			return c.Control(func(fd uintptr) {
   104  				syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEADDR, 1)
   105  				syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, unix.SO_REUSEPORT, 1)
   106  			})
   107  		}}
   108  		transporter := NewTransporter(&config.Options{
   109  			Addr:         addr,
   110  			Network:      nw,
   111  			ListenConfig: listenCfg,
   112  		})
   113  		go transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
   114  			return nil
   115  		})
   116  		defer transporter.Close()
   117  	})
   118  
   119  	t.Run("TestExceptionCase", func(t *testing.T) {
   120  		assert.Panic(t, func() { // listen err
   121  			transporter := NewTransporter(&config.Options{
   122  				Network: "unknow",
   123  			})
   124  			transporter.ListenAndServe(func(ctx context.Context, conn interface{}) error {
   125  				return nil
   126  			})
   127  		})
   128  	})
   129  }