github.com/MetalBlockchain/metalgo@v1.11.9/network/p2p/throttler_handler_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 p2p
     5  
     6  import (
     7  	"context"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/MetalBlockchain/metalgo/ids"
    14  	"github.com/MetalBlockchain/metalgo/utils/logging"
    15  )
    16  
    17  var _ Handler = (*TestHandler)(nil)
    18  
    19  func TestThrottlerHandlerAppGossip(t *testing.T) {
    20  	tests := []struct {
    21  		name      string
    22  		Throttler Throttler
    23  		expected  bool
    24  	}{
    25  		{
    26  			name:      "not throttled",
    27  			Throttler: NewSlidingWindowThrottler(time.Second, 1),
    28  			expected:  true,
    29  		},
    30  		{
    31  			name:      "throttled",
    32  			Throttler: NewSlidingWindowThrottler(time.Second, 0),
    33  		},
    34  	}
    35  	for _, tt := range tests {
    36  		t.Run(tt.name, func(t *testing.T) {
    37  			require := require.New(t)
    38  
    39  			called := false
    40  			handler := NewThrottlerHandler(
    41  				TestHandler{
    42  					AppGossipF: func(context.Context, ids.NodeID, []byte) {
    43  						called = true
    44  					},
    45  				},
    46  				tt.Throttler,
    47  				logging.NoLog{},
    48  			)
    49  
    50  			handler.AppGossip(context.Background(), ids.GenerateTestNodeID(), []byte("foobar"))
    51  			require.Equal(tt.expected, called)
    52  		})
    53  	}
    54  }
    55  
    56  func TestThrottlerHandlerAppRequest(t *testing.T) {
    57  	tests := []struct {
    58  		name        string
    59  		Throttler   Throttler
    60  		expectedErr error
    61  	}{
    62  		{
    63  			name:      "not throttled",
    64  			Throttler: NewSlidingWindowThrottler(time.Second, 1),
    65  		},
    66  		{
    67  			name:        "throttled",
    68  			Throttler:   NewSlidingWindowThrottler(time.Second, 0),
    69  			expectedErr: ErrThrottled,
    70  		},
    71  	}
    72  	for _, tt := range tests {
    73  		t.Run(tt.name, func(t *testing.T) {
    74  			require := require.New(t)
    75  
    76  			handler := NewThrottlerHandler(
    77  				NoOpHandler{},
    78  				tt.Throttler,
    79  				logging.NoLog{},
    80  			)
    81  			_, err := handler.AppRequest(context.Background(), ids.GenerateTestNodeID(), time.Time{}, []byte("foobar"))
    82  			require.ErrorIs(err, tt.expectedErr)
    83  		})
    84  	}
    85  }