limiter_test.go — fiber Source File

Source Code

package limiter

import (
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/valyala/fasthttp"

	"github.com/gofiber/fiber/v3"
	"github.com/gofiber/fiber/v3/internal/storage/memory"
)

type failingLimiterStorage struct {
	data map[string][]byte
	errs map[string]error
}

const testLimiterClientKey = "client-key"

func newFailingLimiterStorage() *failingLimiterStorage {
	return &failingLimiterStorage{
		data: make(map[string][]byte),
		errs: make(map[string]error),
	}
}

// countingFailStorage fails set operations after a specified number of successful calls
type countingFailStorage struct {
	*failingLimiterStorage
	setFailErr error
	setCount   int
	failAfterN int
}

func newCountingFailStorage(failAfterN int, err error) *countingFailStorage {
	return &countingFailStorage{
		failingLimiterStorage: newFailingLimiterStorage(),
		failAfterN:            failAfterN,
		setFailErr:            err,
	}
}

func (s *countingFailStorage) SetWithContext(ctx context.Context, key string, val []byte, exp time.Duration) error {
	s.setCount++
	if s.setCount > s.failAfterN {
		return s.setFailErr
	}
	return s.failingLimiterStorage.SetWithContext(ctx, key, val, exp)
}

type contextRecord struct {
// ... (1663 more lines)