rate limit and config tests added

This commit is contained in:
0ceanSlim 2024-07-26 16:02:43 -04:00
parent 4c5375a429
commit 4181ff5f94
2 changed files with 185 additions and 0 deletions

77
tests/config_test.go Normal file
View File

@ -0,0 +1,77 @@
package tests
import (
"testing"
"grain/relay/utils"
)
func TestConfigValidity(t *testing.T) {
config, err := utils.LoadConfig("../config.yml")
if err != nil {
t.Fatalf("Error loading config: %v", err)
}
// Check MongoDB settings
if config.MongoDB.URI == "" {
t.Error("MongoDB URI is required")
}
if config.MongoDB.Database == "" {
t.Error("MongoDB database name is required")
}
// Check Server settings
if config.Server.Port == "" {
t.Error("Server port is required")
}
// Check Rate Limit settings
if config.RateLimit.WsLimit == 0 {
t.Error("WebSocket limit is required")
}
if config.RateLimit.WsBurst == 0 {
t.Error("WebSocket burst is required")
}
if config.RateLimit.EventLimit == 0 {
t.Error("Event limit is required")
}
if config.RateLimit.EventBurst == 0 {
t.Error("Event burst is required")
}
if config.RateLimit.ReqLimit == 0 {
t.Error("REQ limit is required")
}
if config.RateLimit.ReqBurst == 0 {
t.Error("REQ burst is required")
}
// Check Category Limits
if len(config.RateLimit.CategoryLimits) == 0 {
t.Log("Warning: No category limits set")
}
// Check Kind Limits
if len(config.RateLimit.KindLimits) == 0 {
t.Log("Warning: No kind limits set")
}
// Validate individual category limits
for category, limits := range config.RateLimit.CategoryLimits {
if limits.Limit == 0 {
t.Errorf("Limit is required for category: %s", category)
}
if limits.Burst == 0 {
t.Errorf("Burst is required for category: %s", category)
}
}
// Validate individual kind limits
for _, kindLimit := range config.RateLimit.KindLimits {
if kindLimit.Limit == 0 {
t.Errorf("Limit is required for kind: %d", kindLimit.Kind)
}
if kindLimit.Burst == 0 {
t.Errorf("Burst is required for kind: %d", kindLimit.Kind)
}
}
}

108
tests/rateLimits_test.go Normal file
View File

@ -0,0 +1,108 @@
package tests
import (
"testing"
"grain/relay/utils"
"golang.org/x/time/rate"
)
func TestWebSocketRateLimit(t *testing.T) {
rateLimiter := utils.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 := utils.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 := utils.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 := utils.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 := utils.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)
}
}
}