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 }