github.com/status-im/status-go@v1.1.0/rpc/chain/rpc_limiter_test.go (about)

     1  package chain
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/stretchr/testify/require"
     8  )
     9  
    10  func setupTest() (*InMemRequestsMapStorage, RequestLimiter) {
    11  	storage := NewInMemRequestsMapStorage()
    12  	rl := NewRequestLimiter(storage)
    13  	return storage, rl
    14  }
    15  
    16  func TestSetLimit(t *testing.T) {
    17  	storage, rl := setupTest()
    18  
    19  	// Define test inputs
    20  	tag := "testTag"
    21  	maxRequests := 10
    22  	interval := time.Second
    23  
    24  	// Call the SetLimit method
    25  	err := rl.SetLimit(tag, maxRequests, interval)
    26  	require.NoError(t, err)
    27  
    28  	// Verify that the data was saved to storage correctly
    29  	data, err := storage.Get(tag)
    30  	require.NoError(t, err)
    31  	require.Equal(t, tag, data.Tag)
    32  	require.Equal(t, interval, data.Period)
    33  	require.Equal(t, maxRequests, data.MaxReqs)
    34  	require.Equal(t, 0, data.NumReqs)
    35  }
    36  
    37  func TestGetLimit(t *testing.T) {
    38  	storage, rl := setupTest()
    39  
    40  	// Define test inputs
    41  	data := &LimitData{
    42  		Tag:     "testTag",
    43  		Period:  time.Second,
    44  		MaxReqs: 10,
    45  		NumReqs: 1,
    46  	}
    47  	err := storage.Set(data)
    48  	require.NoError(t, err)
    49  
    50  	// Call the GetLimit method
    51  	ret, err := rl.GetLimit(data.Tag)
    52  	require.NoError(t, err)
    53  
    54  	// Verify the returned data
    55  	require.Equal(t, data, ret)
    56  }
    57  
    58  func TestDeleteLimit(t *testing.T) {
    59  	storage, rl := setupTest()
    60  
    61  	// Define test inputs
    62  	tag := "testTag"
    63  	data := &LimitData{
    64  		Tag:     tag,
    65  		Period:  time.Second,
    66  		MaxReqs: 10,
    67  		NumReqs: 1,
    68  	}
    69  	err := storage.Set(data)
    70  	require.NoError(t, err)
    71  
    72  	// Call the DeleteLimit method
    73  	err = rl.DeleteLimit(tag)
    74  	require.NoError(t, err)
    75  
    76  	// Verify that the data was deleted from storage
    77  	limit, _ := storage.Get(tag)
    78  	require.Nil(t, limit)
    79  
    80  	// Test double delete
    81  	err = rl.DeleteLimit(tag)
    82  	require.NoError(t, err)
    83  }
    84  
    85  func TestAllowWithinPeriod(t *testing.T) {
    86  	storage, rl := setupTest()
    87  
    88  	// Define test inputs
    89  	tag := "testTag"
    90  	maxRequests := 10
    91  	interval := time.Second
    92  
    93  	// Set up the storage with test data
    94  	data := &LimitData{
    95  		Tag:       tag,
    96  		Period:    interval,
    97  		CreatedAt: time.Now(),
    98  		MaxReqs:   maxRequests,
    99  	}
   100  	err := storage.Set(data)
   101  	require.NoError(t, err)
   102  
   103  	// Call the Allow method
   104  	for i := 0; i < maxRequests; i++ {
   105  		allow, err := rl.Allow(tag)
   106  		require.NoError(t, err)
   107  
   108  		// Verify the result
   109  		require.True(t, allow)
   110  	}
   111  
   112  	// Call the Allow method again
   113  	allow, err := rl.Allow(tag)
   114  	require.ErrorIs(t, err, ErrRequestsOverLimit)
   115  	require.False(t, allow)
   116  }
   117  
   118  func TestAllowWhenPeriodPassed(t *testing.T) {
   119  	storage, rl := setupTest()
   120  
   121  	// Define test inputs
   122  	tag := "testTag"
   123  	maxRequests := 10
   124  	interval := time.Second
   125  
   126  	// Set up the storage with test data
   127  	data := &LimitData{
   128  		Tag:       tag,
   129  		Period:    interval,
   130  		CreatedAt: time.Now().Add(-interval),
   131  		MaxReqs:   maxRequests,
   132  		NumReqs:   maxRequests,
   133  	}
   134  	err := storage.Set(data)
   135  	require.NoError(t, err)
   136  
   137  	// Call the Allow method
   138  	allow, err := rl.Allow(tag)
   139  	require.NoError(t, err)
   140  
   141  	// Verify the result
   142  	require.True(t, allow)
   143  }
   144  
   145  func TestAllowRestrictInfinitelyWhenLimitReached(t *testing.T) {
   146  	storage, rl := setupTest()
   147  
   148  	// Define test inputs
   149  	tag := "testTag"
   150  	maxRequests := 10
   151  
   152  	// Set up the storage with test data
   153  	data := &LimitData{
   154  		Tag:       tag,
   155  		Period:    LimitInfinitely,
   156  		CreatedAt: time.Now(),
   157  		MaxReqs:   maxRequests,
   158  		NumReqs:   maxRequests,
   159  	}
   160  	err := storage.Set(data)
   161  	require.NoError(t, err)
   162  
   163  	// Call the Allow method
   164  	allow, err := rl.Allow(tag)
   165  	require.ErrorIs(t, err, ErrRequestsOverLimit)
   166  
   167  	// Verify the result
   168  	require.False(t, allow)
   169  }
   170  
   171  func TestAllowWhenLimitNotReachedForInfinitePeriod(t *testing.T) {
   172  	storage, rl := setupTest()
   173  
   174  	// Define test inputs
   175  	tag := "testTag"
   176  	maxRequests := 10
   177  
   178  	// Set up the storage with test data
   179  	data := &LimitData{
   180  		Tag:       tag,
   181  		Period:    LimitInfinitely,
   182  		CreatedAt: time.Now(),
   183  		MaxReqs:   maxRequests,
   184  		NumReqs:   maxRequests - 1,
   185  	}
   186  	err := storage.Set(data)
   187  	require.NoError(t, err)
   188  
   189  	// Call the Allow method
   190  	allow, err := rl.Allow(tag)
   191  	require.NoError(t, err)
   192  
   193  	// Verify the result
   194  	require.True(t, allow)
   195  }