mirror of
https://github.com/0ceanSlim/grain.git
synced 2024-11-22 08:37:13 +00:00
109 lines
3.4 KiB
Go
109 lines
3.4 KiB
Go
package tests
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"grain/config"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
func TestWebSocketRateLimit(t *testing.T) {
|
|
rateLimiter := config.NewRateLimiter(rate.Limit(1), 1, rate.Limit(100), 200, rate.Limit(100), 200)
|
|
|
|
// First message should be allowed
|
|
if allowed, _ := rateLimiter.AllowWs(); !allowed {
|
|
t.Error("First WebSocket message should be allowed")
|
|
}
|
|
|
|
// Second message should be rate-limited
|
|
if allowed, msg := rateLimiter.AllowWs(); allowed {
|
|
t.Error("Second WebSocket message should be rate-limited")
|
|
} else {
|
|
expectedMsg := "WebSocket message rate limit exceeded"
|
|
if msg != expectedMsg {
|
|
t.Errorf("Expected message: %s, got: %s", expectedMsg, msg)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestEventRateLimit(t *testing.T) {
|
|
rateLimiter := config.NewRateLimiter(rate.Limit(100), 200, rate.Limit(1), 1, rate.Limit(100), 200)
|
|
rateLimiter.AddKindLimit(1, rate.Limit(1), 1)
|
|
rateLimiter.AddCategoryLimit("regular", rate.Limit(1), 1)
|
|
|
|
// First event should be allowed
|
|
if allowed, _ := rateLimiter.AllowEvent(1, "regular"); !allowed {
|
|
t.Error("First event should be allowed")
|
|
}
|
|
|
|
// Second event should be rate-limited
|
|
if allowed, msg := rateLimiter.AllowEvent(1, "regular"); allowed {
|
|
t.Error("Second event should be rate-limited")
|
|
} else {
|
|
expectedMsg := "Global event rate limit exceeded"
|
|
if msg != expectedMsg {
|
|
t.Errorf("Expected message: %s, got: %s", expectedMsg, msg)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestReqRateLimit(t *testing.T) {
|
|
rateLimiter := config.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(1), 1)
|
|
|
|
// First REQ should be allowed
|
|
if allowed, _ := rateLimiter.AllowReq(); !allowed {
|
|
t.Error("First REQ message should be allowed")
|
|
}
|
|
|
|
// Second REQ should be rate-limited
|
|
if allowed, msg := rateLimiter.AllowReq(); allowed {
|
|
t.Error("Second REQ message should be rate-limited")
|
|
} else {
|
|
expectedMsg := "REQ rate limit exceeded"
|
|
if msg != expectedMsg {
|
|
t.Errorf("Expected message: %s, got: %s", expectedMsg, msg)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestKindRateLimit(t *testing.T) {
|
|
rateLimiter := config.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(100), 200)
|
|
rateLimiter.AddKindLimit(1, rate.Limit(1), 1)
|
|
|
|
// First event of kind 1 should be allowed
|
|
if allowed, _ := rateLimiter.AllowEvent(1, "regular"); !allowed {
|
|
t.Error("First event of kind 1 should be allowed")
|
|
}
|
|
|
|
// Second event of kind 1 should be rate-limited
|
|
if allowed, msg := rateLimiter.AllowEvent(1, "regular"); allowed {
|
|
t.Error("Second event of kind 1 should be rate-limited")
|
|
} else {
|
|
expectedMsg := "Rate limit exceeded for kind: 1"
|
|
if msg != expectedMsg {
|
|
t.Errorf("Expected message: %s, got: %s", expectedMsg, msg)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCategoryRateLimit(t *testing.T) {
|
|
rateLimiter := config.NewRateLimiter(rate.Limit(100), 200, rate.Limit(100), 200, rate.Limit(100), 200)
|
|
rateLimiter.AddCategoryLimit("regular", rate.Limit(1), 1)
|
|
|
|
// First event in category "regular" should be allowed
|
|
if allowed, _ := rateLimiter.AllowEvent(1, "regular"); !allowed {
|
|
t.Error("First event in category 'regular' should be allowed")
|
|
}
|
|
|
|
// Second event in category "regular" should be rate-limited
|
|
if allowed, msg := rateLimiter.AllowEvent(1, "regular"); allowed {
|
|
t.Error("Second event in category 'regular' should be rate-limited")
|
|
} else {
|
|
expectedMsg := "Rate limit exceeded for category: regular"
|
|
if msg != expectedMsg {
|
|
t.Errorf("Expected message: %s, got: %s", expectedMsg, msg)
|
|
}
|
|
}
|
|
}
|