status-go/rpc/chain/rpc_limiter_test.go

169 lines
3.3 KiB
Go
Raw Normal View History

package chain
import (
"testing"
"time"
"github.com/stretchr/testify/require"
)
func setupTest() (*InMemRequestsMapStorage, RequestLimiter) {
storage := NewInMemRequestsMapStorage()
rl := NewRequestLimiter(storage)
return storage, rl
}
func TestSetLimit(t *testing.T) {
storage, rl := setupTest()
// Define test inputs
tag := "testTag"
maxRequests := 10
interval := time.Second
// Call the SetLimit method
err := rl.SetLimit(tag, maxRequests, interval)
require.NoError(t, err)
// Verify that the data was saved to storage correctly
data, err := storage.Get(tag)
require.NoError(t, err)
require.Equal(t, tag, data.Tag)
require.Equal(t, interval, data.Period)
require.Equal(t, maxRequests, data.MaxReqs)
require.Equal(t, 0, data.NumReqs)
}
func TestGetLimit(t *testing.T) {
storage, rl := setupTest()
// Define test inputs
data := &LimitData{
Tag: "testTag",
Period: time.Second,
MaxReqs: 10,
NumReqs: 1,
}
err := storage.Set(data)
require.NoError(t, err)
// Call the GetLimit method
ret, err := rl.GetLimit(data.Tag)
require.NoError(t, err)
// Verify the returned data
require.Equal(t, data, ret)
}
func TestAllowWithinPeriod(t *testing.T) {
storage, rl := setupTest()
// Define test inputs
tag := "testTag"
maxRequests := 10
interval := time.Second
// Set up the storage with test data
data := &LimitData{
Tag: tag,
Period: interval,
CreatedAt: time.Now(),
MaxReqs: maxRequests,
}
err := storage.Set(data)
require.NoError(t, err)
// Call the Allow method
for i := 0; i < maxRequests; i++ {
allow, err := rl.Allow(tag)
require.NoError(t, err)
// Verify the result
require.True(t, allow)
}
// Call the Allow method again
allow, err := rl.Allow(tag)
require.NoError(t, err)
require.False(t, allow)
}
func TestAllowWhenPeriodPassed(t *testing.T) {
storage, rl := setupTest()
// Define test inputs
tag := "testTag"
maxRequests := 10
interval := time.Second
// Set up the storage with test data
data := &LimitData{
Tag: tag,
Period: interval,
CreatedAt: time.Now().Add(-interval),
MaxReqs: maxRequests,
NumReqs: maxRequests,
}
err := storage.Set(data)
require.NoError(t, err)
// Call the Allow method
allow, err := rl.Allow(tag)
require.NoError(t, err)
// Verify the result
require.True(t, allow)
}
func TestAllowRestrictInfinitelyWhenLimitReached(t *testing.T) {
storage, rl := setupTest()
// Define test inputs
tag := "testTag"
maxRequests := 10
// Set up the storage with test data
data := &LimitData{
Tag: tag,
Period: LimitInfinitely,
CreatedAt: time.Now(),
MaxReqs: maxRequests,
NumReqs: maxRequests,
}
err := storage.Set(data)
require.NoError(t, err)
// Call the Allow method
allow, err := rl.Allow(tag)
require.NoError(t, err)
// Verify the result
require.False(t, allow)
}
func TestAllowWhenLimitNotReachedForInfinitePeriod(t *testing.T) {
storage, rl := setupTest()
// Define test inputs
tag := "testTag"
maxRequests := 10
// Set up the storage with test data
data := &LimitData{
Tag: tag,
Period: LimitInfinitely,
CreatedAt: time.Now(),
MaxReqs: maxRequests,
NumReqs: maxRequests - 1,
}
err := storage.Set(data)
require.NoError(t, err)
// Call the Allow method
allow, err := rl.Allow(tag)
require.NoError(t, err)
// Verify the result
require.True(t, allow)
}