75 lines
1.9 KiB
Go
75 lines
1.9 KiB
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
)
|
||
|
|
|
||
|
|
func newTestEngine(mw gin.HandlerFunc) *gin.Engine {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
engine := gin.New()
|
||
|
|
engine.GET("/x", mw, func(c *gin.Context) {
|
||
|
|
c.JSON(http.StatusOK, gin.H{"ok": true})
|
||
|
|
})
|
||
|
|
return engine
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRateLimiter_AllowsWithinLimit(t *testing.T) {
|
||
|
|
engine := newTestEngine(RateLimiter(3, 1))
|
||
|
|
|
||
|
|
for i := 0; i < 3; i++ {
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||
|
|
req.RemoteAddr = "1.2.3.4:1111"
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
engine.ServeHTTP(w, req)
|
||
|
|
if w.Code != http.StatusOK {
|
||
|
|
t.Fatalf("request %d: expected 200, got %d", i+1, w.Code)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRateLimiter_BlocksOverLimit(t *testing.T) {
|
||
|
|
engine := newTestEngine(RateLimiter(2, 0)) // بدون شارژ دوباره، فقط ۲ تا مجازه
|
||
|
|
|
||
|
|
for i := 0; i < 2; i++ {
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||
|
|
req.RemoteAddr = "5.6.7.8:2222"
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
engine.ServeHTTP(w, req)
|
||
|
|
if w.Code != http.StatusOK {
|
||
|
|
t.Fatalf("request %d: expected 200, got %d", i+1, w.Code)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||
|
|
req.RemoteAddr = "5.6.7.8:2222"
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
engine.ServeHTTP(w, req)
|
||
|
|
if w.Code != http.StatusTooManyRequests {
|
||
|
|
t.Fatalf("expected 429 on 3rd request, got %d", w.Code)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRateLimiter_SeparateBucketsPerIP(t *testing.T) {
|
||
|
|
engine := newTestEngine(RateLimiter(1, 0))
|
||
|
|
|
||
|
|
req1 := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||
|
|
req1.RemoteAddr = "9.9.9.9:1"
|
||
|
|
w1 := httptest.NewRecorder()
|
||
|
|
engine.ServeHTTP(w1, req1)
|
||
|
|
if w1.Code != http.StatusOK {
|
||
|
|
t.Fatalf("expected 200 for first IP, got %d", w1.Code)
|
||
|
|
}
|
||
|
|
|
||
|
|
req2 := httptest.NewRequest(http.MethodGet, "/x", nil)
|
||
|
|
req2.RemoteAddr = "8.8.8.8:1"
|
||
|
|
w2 := httptest.NewRecorder()
|
||
|
|
engine.ServeHTTP(w2, req2)
|
||
|
|
if w2.Code != http.StatusOK {
|
||
|
|
t.Fatalf("expected 200 for a different IP (separate bucket), got %d", w2.Code)
|
||
|
|
}
|
||
|
|
}
|