github.com/cloudwego/kitex@v0.9.0/pkg/remote/connpool/short_pool_test.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package connpool
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"net"
    23  	"testing"
    24  	"time"
    25  
    26  	"github.com/golang/mock/gomock"
    27  
    28  	mocksnet "github.com/cloudwego/kitex/internal/mocks/net"
    29  
    30  	mocksremote "github.com/cloudwego/kitex/internal/mocks/remote"
    31  	"github.com/cloudwego/kitex/internal/test"
    32  	"github.com/cloudwego/kitex/pkg/remote"
    33  )
    34  
    35  func TestShortPool(t *testing.T) {
    36  	ctrl := gomock.NewController(t)
    37  	defer ctrl.Finish()
    38  
    39  	d := mocksremote.NewMockDialer(ctrl)
    40  	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
    41  		cost := time.Millisecond * 10
    42  		if timeout < cost {
    43  			return nil, errors.New("connection timeout")
    44  		}
    45  		return nil, nil
    46  	}).AnyTimes()
    47  	p := NewShortPool("test")
    48  
    49  	_, err := p.Get(context.TODO(), "tcp", "127.0.0.1:8080", remote.ConnOption{Dialer: d, ConnectTimeout: time.Second})
    50  	test.Assert(t, err == nil)
    51  
    52  	_, err = p.Get(context.TODO(), "tcp", "127.0.0.1:8888", remote.ConnOption{Dialer: d, ConnectTimeout: time.Millisecond})
    53  	test.Assert(t, err != nil)
    54  }
    55  
    56  func TestShortPoolPut(t *testing.T) {
    57  	ctrl := gomock.NewController(t)
    58  	defer ctrl.Finish()
    59  
    60  	conn := mocksnet.NewMockConn(ctrl)
    61  	conn.EXPECT().Close().AnyTimes()
    62  	d := mocksremote.NewMockDialer(ctrl)
    63  	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
    64  		return conn, nil
    65  	}).AnyTimes()
    66  	p := NewShortPool("test")
    67  
    68  	x, err := p.Get(context.TODO(), "tcp", "127.0.0.1:8080", remote.ConnOption{Dialer: d, ConnectTimeout: time.Second})
    69  	test.Assert(t, err == nil)
    70  	y, ok := x.(*shortConn)
    71  	test.Assert(t, ok)
    72  	test.Assert(t, y.Conn == conn)
    73  	test.Assert(t, !y.closed)
    74  
    75  	err = p.Put(x)
    76  	test.Assert(t, err == nil)
    77  	test.Assert(t, y.closed)
    78  }
    79  
    80  func TestShortPoolPutClosed(t *testing.T) {
    81  	ctrl := gomock.NewController(t)
    82  	defer ctrl.Finish()
    83  
    84  	p := NewShortPool("test")
    85  
    86  	var closed bool
    87  	cmto := errors.New("closed more than once")
    88  
    89  	conn := mocksnet.NewMockConn(ctrl)
    90  	conn.EXPECT().Close().DoAndReturn(func() (err error) {
    91  		if closed {
    92  			err = cmto
    93  		}
    94  		closed = true
    95  		return
    96  	})
    97  	d := mocksremote.NewMockDialer(ctrl)
    98  	d.EXPECT().DialTimeout(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(func(network, address string, timeout time.Duration) (net.Conn, error) {
    99  		return conn, nil
   100  	}).AnyTimes()
   101  
   102  	x, err := p.Get(context.TODO(), "tcp", "127.0.0.1:8080", remote.ConnOption{Dialer: d, ConnectTimeout: time.Second})
   103  	test.Assert(t, err == nil)
   104  	y, ok := x.(*shortConn)
   105  	test.Assert(t, ok)
   106  	test.Assert(t, y.Conn == conn)
   107  	test.Assert(t, !y.closed)
   108  	test.Assert(t, !closed)
   109  
   110  	err = y.Close()
   111  	test.Assert(t, err == nil)
   112  	test.Assert(t, y.closed)
   113  	test.Assert(t, closed)
   114  
   115  	err = y.Close()
   116  	test.Assert(t, err == nil)
   117  
   118  	err = p.Put(x)
   119  	test.Assert(t, err == nil)
   120  }