status-go/rpc/chain/rpc_limiter_test.go

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)
}