78 lines
2.1 KiB
Go
78 lines
2.1 KiB
Go
|
package middleware
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
)
|
||
|
|
|
||
|
|
func setupRouterWithRateLimit(rps float64, burst int) *gin.Engine {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
r := gin.New()
|
||
|
|
r.Use(RateLimiter(rps, burst))
|
||
|
|
r.GET("/ping", func(c *gin.Context) {
|
||
|
|
c.JSON(http.StatusOK, gin.H{"message": "pong"})
|
||
|
|
})
|
||
|
|
return r
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRateLimiter_MengizinkanRequestSelamaTokenTersedia(t *testing.T) {
|
||
|
|
// burst 3 artinya 3 request pertama harus tetap lolos meski dikirim beruntun
|
||
|
|
r := setupRouterWithRateLimit(1, 3)
|
||
|
|
|
||
|
|
for i := 0; i < 3; i++ {
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/ping", nil)
|
||
|
|
r.ServeHTTP(w, req)
|
||
|
|
|
||
|
|
if w.Code != http.StatusOK {
|
||
|
|
t.Fatalf("request ke-%d: mau status 200, dapat %d", i+1, w.Code)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRateLimiter_MenolakSetelahBurstHabis(t *testing.T) {
|
||
|
|
// burst 2: request ke-3 yang dikirim beruntun (tanpa jeda) harus ditolak 429
|
||
|
|
r := setupRouterWithRateLimit(1, 2)
|
||
|
|
|
||
|
|
for i := 0; i < 2; i++ {
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/ping", nil)
|
||
|
|
r.ServeHTTP(w, req)
|
||
|
|
if w.Code != http.StatusOK {
|
||
|
|
t.Fatalf("request ke-%d seharusnya lolos, dapat status %d", i+1, w.Code)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
req := httptest.NewRequest(http.MethodGet, "/ping", nil)
|
||
|
|
r.ServeHTTP(w, req)
|
||
|
|
|
||
|
|
if w.Code != http.StatusTooManyRequests {
|
||
|
|
t.Fatalf("request ke-3 seharusnya ditolak 429, dapat status %d", w.Code)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRateLimiter_IPBerbedaPunyaJatahTerpisah(t *testing.T) {
|
||
|
|
// burst 1: IP A dan IP B masing-masing harus dapat 1 jatah request sendiri-sendiri
|
||
|
|
r := setupRouterWithRateLimit(1, 1)
|
||
|
|
|
||
|
|
reqA := httptest.NewRequest(http.MethodGet, "/ping", nil)
|
||
|
|
reqA.RemoteAddr = "1.1.1.1:1234"
|
||
|
|
wA := httptest.NewRecorder()
|
||
|
|
r.ServeHTTP(wA, reqA)
|
||
|
|
|
||
|
|
reqB := httptest.NewRequest(http.MethodGet, "/ping", nil)
|
||
|
|
reqB.RemoteAddr = "2.2.2.2:5678"
|
||
|
|
wB := httptest.NewRecorder()
|
||
|
|
r.ServeHTTP(wB, reqB)
|
||
|
|
|
||
|
|
if wA.Code != http.StatusOK {
|
||
|
|
t.Errorf("IP A request pertama seharusnya lolos, dapat status %d", wA.Code)
|
||
|
|
}
|
||
|
|
if wB.Code != http.StatusOK {
|
||
|
|
t.Errorf("IP B request pertama seharusnya lolos meski IP A sudah pakai jatahnya, dapat status %d", wB.Code)
|
||
|
|
}
|
||
|
|
}
|