103 lines
2.5 KiB
Go
103 lines
2.5 KiB
Go
|
package ratelimit
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
)
|
||
|
|
|
||
|
|
// performRequest — минимальный прогон запроса через gin-роутер.
|
||
|
|
func performRequest(r *gin.Engine, method, path string) *httptest.ResponseRecorder {
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
req := httptest.NewRequest(method, path, nil)
|
||
|
|
r.ServeHTTP(w, req)
|
||
|
|
return w
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestLimiterAllow(t *testing.T) {
|
||
|
|
l := New(3, time.Minute)
|
||
|
|
|
||
|
|
for i := 1; i <= 3; i++ {
|
||
|
|
if !l.Allow("ip-1") {
|
||
|
|
t.Fatalf("attempt %d: want allowed", i)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if l.Allow("ip-1") {
|
||
|
|
t.Error("4th attempt: want denied")
|
||
|
|
}
|
||
|
|
// Ключи независимы: чужой бюджет не задет.
|
||
|
|
if !l.Allow("ip-2") {
|
||
|
|
t.Error("other key: want allowed")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestLimiterWindowReset(t *testing.T) {
|
||
|
|
l := New(1, 30*time.Millisecond)
|
||
|
|
|
||
|
|
if !l.Allow("k") {
|
||
|
|
t.Fatal("first attempt: want allowed")
|
||
|
|
}
|
||
|
|
if l.Allow("k") {
|
||
|
|
t.Error("second attempt: want denied")
|
||
|
|
}
|
||
|
|
|
||
|
|
time.Sleep(40 * time.Millisecond)
|
||
|
|
|
||
|
|
if !l.Allow("k") {
|
||
|
|
t.Error("after window: want allowed")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestLimiterCleanup — после чистки карта не тащит протухшие ключи,
|
||
|
|
// иначе под потоком разных адресов она росла бы без предела.
|
||
|
|
func TestLimiterCleanup(t *testing.T) {
|
||
|
|
l := New(5, 20*time.Millisecond)
|
||
|
|
|
||
|
|
l.Allow("old")
|
||
|
|
time.Sleep(30 * time.Millisecond) // ключ протух, но чистки ещё не было
|
||
|
|
|
||
|
|
l.Allow("fresh") // триггерит cleanupLocked
|
||
|
|
if _, exists := l.buckets["old"]; exists {
|
||
|
|
t.Error("stale key survived cleanup")
|
||
|
|
}
|
||
|
|
if _, exists := l.buckets["fresh"]; !exists {
|
||
|
|
t.Error("fresh key missing")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestMiddlewareAborts(t *testing.T) {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
|
||
|
|
l := New(1, time.Minute)
|
||
|
|
h := Middleware(l, ByIP, "too many attempts")
|
||
|
|
|
||
|
|
var passed int
|
||
|
|
r := gin.New()
|
||
|
|
r.Use(func(c *gin.Context) { // абстрактный «хендлер», который просто считает проходы
|
||
|
|
h(c)
|
||
|
|
if c.IsAborted() {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
passed++
|
||
|
|
c.JSON(200, gin.H{"ok": true})
|
||
|
|
})
|
||
|
|
|
||
|
|
w := performRequest(r, "POST", "/x")
|
||
|
|
if w.Code != http.StatusOK {
|
||
|
|
t.Fatalf("first request: status = %d, want 200", w.Code)
|
||
|
|
}
|
||
|
|
w = performRequest(r, "POST", "/x")
|
||
|
|
if w.Code != http.StatusTooManyRequests {
|
||
|
|
t.Fatalf("second request: status = %d, want 429", w.Code)
|
||
|
|
}
|
||
|
|
if body := w.Body.String(); body != `{"error":"too many attempts"}` {
|
||
|
|
t.Errorf("body = %s", body)
|
||
|
|
}
|
||
|
|
if passed != 1 {
|
||
|
|
t.Errorf("handler ran %d times, want 1", passed)
|
||
|
|
}
|
||
|
|
}
|