github.com/ava-labs/avalanchego@v1.11.11/network/throttling/inbound_conn_throttler_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package throttling
     5  
     6  import (
     7  	"context"
     8  	"net"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  var _ net.Listener = (*MockListener)(nil)
    15  
    16  type MockListener struct {
    17  	t         *testing.T
    18  	OnAcceptF func() (net.Conn, error)
    19  	OnCloseF  func() error
    20  	OnAddrF   func() net.Addr
    21  }
    22  
    23  func (ml *MockListener) Accept() (net.Conn, error) {
    24  	if ml.OnAcceptF == nil {
    25  		require.FailNow(ml.t, "unexpectedly called Accept")
    26  		return nil, nil
    27  	}
    28  	return ml.OnAcceptF()
    29  }
    30  
    31  func (ml *MockListener) Close() error {
    32  	if ml.OnCloseF == nil {
    33  		require.FailNow(ml.t, "unexpectedly called Close")
    34  		return nil
    35  	}
    36  	return ml.OnCloseF()
    37  }
    38  
    39  func (ml *MockListener) Addr() net.Addr {
    40  	if ml.OnAddrF == nil {
    41  		require.FailNow(ml.t, "unexpectedly called Addr")
    42  		return nil
    43  	}
    44  	return ml.OnAddrF()
    45  }
    46  
    47  func TestInboundConnThrottlerClose(t *testing.T) {
    48  	require := require.New(t)
    49  
    50  	closed := false
    51  	l := &MockListener{
    52  		t: t,
    53  		OnCloseF: func() error {
    54  			closed = true
    55  			return nil
    56  		},
    57  	}
    58  	wrappedL := NewThrottledListener(l, 1)
    59  	require.NoError(wrappedL.Close())
    60  	require.True(closed)
    61  
    62  	select {
    63  	case <-wrappedL.(*throttledListener).ctx.Done():
    64  	default:
    65  		require.FailNow("should have closed context")
    66  	}
    67  
    68  	// Accept() should return an error because the context is cancelled
    69  	_, err := wrappedL.Accept()
    70  	require.ErrorIs(err, context.Canceled)
    71  }
    72  
    73  func TestInboundConnThrottlerAddr(t *testing.T) {
    74  	addrCalled := false
    75  	l := &MockListener{
    76  		t: t,
    77  		OnAddrF: func() net.Addr {
    78  			addrCalled = true
    79  			return nil
    80  		},
    81  	}
    82  	wrappedL := NewThrottledListener(l, 1)
    83  	_ = wrappedL.Addr()
    84  	require.True(t, addrCalled)
    85  }
    86  
    87  func TestInboundConnThrottlerAccept(t *testing.T) {
    88  	require := require.New(t)
    89  
    90  	acceptCalled := false
    91  	l := &MockListener{
    92  		t: t,
    93  		OnAcceptF: func() (net.Conn, error) {
    94  			acceptCalled = true
    95  			return nil, nil
    96  		},
    97  	}
    98  	wrappedL := NewThrottledListener(l, 1)
    99  	_, err := wrappedL.Accept()
   100  	require.NoError(err)
   101  	require.True(acceptCalled)
   102  }