Function TestRateLimiterWithConfig

middleware/rate_limiter_test.go:74–127

(t *testing.T)

}

func TestRateLimiterWithConfig(t *testing.T) {
	var inMemoryStore = NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{Rate: 1, Burst: 3})

	e := echo.New()

	handler := func(c *echo.Context) error {
		return c.String(http.StatusOK, "test")
	}

	mw, err := RateLimiterConfig{
		IdentifierExtractor: func(c *echo.Context) (string, error) {
			id := c.Request().Header.Get(echo.HeaderXRealIP)
			if id == "" {
				return "", errors.New("invalid identifier")
			}
			return id, nil
		},
		DenyHandler: func(ctx *echo.Context, identifier string, err error) error {
			return ctx.JSON(http.StatusForbidden, nil)
		},
		ErrorHandler: func(ctx *echo.Context, err error) error {
			return ctx.JSON(http.StatusBadRequest, nil)
		},
		Store: inMemoryStore,
	}.ToMiddleware()
	assert.NoError(t, err)

	testCases := []struct {
		id   string
		code int
	}{
		{"127.0.0.1", http.StatusOK},
		{"127.0.0.1", http.StatusOK},
		{"127.0.0.1", http.StatusOK},
		{"127.0.0.1", http.StatusForbidden},
		{"", http.StatusBadRequest},
		{"127.0.0.1", http.StatusForbidden},
		{"127.0.0.1", http.StatusForbidden},
	}

	for _, tc := range testCases {
		req := httptest.NewRequest(http.MethodGet, "/", nil)
		req.Header.Add(echo.HeaderXRealIP, tc.id)

		rec := httptest.NewRecorder()

		c := e.NewContext(req, rec)

		err := mw(handler)(c)

		assert.NoError(t, err)
		assert.Equal(t, tc.code, rec.Code)
	}
}

func TestRateLimiterWithConfig_defaultDenyHandler(t *testing.T) {
	var inMemoryStore = NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{Rate: 1, Burst: 3})