7638fd3a90
Implement global and per ip rate limiting * all tests passing * test fix * Enhances test environment for rate limiting Updates the test environment to better support rate limiting tests. - Increases rate limits in test container to avoid interference with legitimate test traffic. - Increases the polling interval for client status checks to reduce load on the rate limiter. - Adds logic to retry status checks if rate limiting is encountered. * Merge branch 'main' of bitbucket.org:aarete/query-orchestration into feature/rate-limiting * build fix * test fix
191 lines
4.8 KiB
Go
191 lines
4.8 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
func TestDefaultConfig(t *testing.T) {
|
|
cfg := DefaultConfig()
|
|
|
|
// Test global defaults
|
|
if cfg.GlobalRate != 500 {
|
|
t.Errorf("Expected GlobalRate=500, got %v", cfg.GlobalRate)
|
|
}
|
|
if cfg.GlobalBurst != 1000 {
|
|
t.Errorf("Expected GlobalBurst=1000, got %d", cfg.GlobalBurst)
|
|
}
|
|
|
|
// Test per-IP defaults
|
|
if cfg.DefaultRate != 5 {
|
|
t.Errorf("Expected DefaultRate=5, got %v", cfg.DefaultRate)
|
|
}
|
|
if cfg.DefaultBurst != 10 {
|
|
t.Errorf("Expected DefaultBurst=10, got %d", cfg.DefaultBurst)
|
|
}
|
|
if cfg.ExpiresIn != 3*time.Minute {
|
|
t.Errorf("Expected ExpiresIn=3m, got %v", cfg.ExpiresIn)
|
|
}
|
|
if cfg.EndpointOverrides == nil {
|
|
t.Error("Expected EndpointOverrides to be initialized")
|
|
}
|
|
}
|
|
|
|
func TestLoadConfigDefaults(t *testing.T) {
|
|
// Clear all env vars
|
|
os.Unsetenv("RATE_LIMIT_GLOBAL_RATE")
|
|
os.Unsetenv("RATE_LIMIT_GLOBAL_BURST")
|
|
os.Unsetenv("RATE_LIMIT_DEFAULT_RATE")
|
|
os.Unsetenv("RATE_LIMIT_DEFAULT_BURST")
|
|
os.Unsetenv("RATE_LIMIT_EXPIRES_IN")
|
|
os.Unsetenv("RATE_LIMIT_ENDPOINT_OVERRIDES")
|
|
|
|
cfg, err := LoadConfig()
|
|
if err != nil {
|
|
t.Fatalf("LoadConfig failed: %v", err)
|
|
}
|
|
|
|
// Should use defaults
|
|
if cfg.GlobalRate != 500 {
|
|
t.Errorf("Expected GlobalRate=500, got %v", cfg.GlobalRate)
|
|
}
|
|
if cfg.DefaultRate != 5 {
|
|
t.Errorf("Expected DefaultRate=5, got %v", cfg.DefaultRate)
|
|
}
|
|
}
|
|
|
|
func TestLoadConfigFromEnvironment(t *testing.T) {
|
|
// Set env vars
|
|
t.Setenv("RATE_LIMIT_GLOBAL_RATE", "500")
|
|
t.Setenv("RATE_LIMIT_GLOBAL_BURST", "1000")
|
|
t.Setenv("RATE_LIMIT_DEFAULT_RATE", "5")
|
|
t.Setenv("RATE_LIMIT_DEFAULT_BURST", "10")
|
|
t.Setenv("RATE_LIMIT_EXPIRES_IN", "300")
|
|
|
|
cfg, err := LoadConfig()
|
|
if err != nil {
|
|
t.Fatalf("LoadConfig failed: %v", err)
|
|
}
|
|
|
|
// Test global limits
|
|
if cfg.GlobalRate != rate.Limit(500) {
|
|
t.Errorf("Expected GlobalRate=500, got %v", cfg.GlobalRate)
|
|
}
|
|
if cfg.GlobalBurst != 1000 {
|
|
t.Errorf("Expected GlobalBurst=1000, got %d", cfg.GlobalBurst)
|
|
}
|
|
|
|
// Test per-IP limits
|
|
if cfg.DefaultRate != rate.Limit(5) {
|
|
t.Errorf("Expected DefaultRate=5, got %v", cfg.DefaultRate)
|
|
}
|
|
if cfg.DefaultBurst != 10 {
|
|
t.Errorf("Expected DefaultBurst=10, got %d", cfg.DefaultBurst)
|
|
}
|
|
if cfg.ExpiresIn != 5*time.Minute {
|
|
t.Errorf("Expected ExpiresIn=5m, got %v", cfg.ExpiresIn)
|
|
}
|
|
}
|
|
|
|
func TestLoadConfigWithEndpointOverrides(t *testing.T) {
|
|
// Set endpoint overrides
|
|
overridesJSON := `{
|
|
"POST /documents": {
|
|
"global_rate": 100,
|
|
"global_burst": 200,
|
|
"rate": 5,
|
|
"burst": 10
|
|
},
|
|
"GET /health": {
|
|
"global_rate": 500,
|
|
"global_burst": 1000,
|
|
"rate": 100,
|
|
"burst": 200
|
|
}
|
|
}`
|
|
t.Setenv("RATE_LIMIT_ENDPOINT_OVERRIDES", overridesJSON)
|
|
|
|
cfg, err := LoadConfig()
|
|
if err != nil {
|
|
t.Fatalf("LoadConfig failed: %v", err)
|
|
}
|
|
|
|
// Test POST /documents override
|
|
override, exists := cfg.EndpointOverrides["POST /documents"]
|
|
if !exists {
|
|
t.Fatal("Expected 'POST /documents' override to exist")
|
|
}
|
|
if override.GlobalRate != rate.Limit(100) {
|
|
t.Errorf("Expected GlobalRate=100, got %v", override.GlobalRate)
|
|
}
|
|
if override.GlobalBurst != 200 {
|
|
t.Errorf("Expected GlobalBurst=200, got %d", override.GlobalBurst)
|
|
}
|
|
if override.Rate != rate.Limit(5) {
|
|
t.Errorf("Expected Rate=5, got %v", override.Rate)
|
|
}
|
|
if override.Burst != 10 {
|
|
t.Errorf("Expected Burst=10, got %d", override.Burst)
|
|
}
|
|
|
|
// Test GET /health override
|
|
override, exists = cfg.EndpointOverrides["GET /health"]
|
|
if !exists {
|
|
t.Fatal("Expected 'GET /health' override to exist")
|
|
}
|
|
if override.GlobalRate != rate.Limit(500) {
|
|
t.Errorf("Expected GlobalRate=500, got %v", override.GlobalRate)
|
|
}
|
|
if override.Rate != rate.Limit(100) {
|
|
t.Errorf("Expected Rate=100, got %v", override.Rate)
|
|
}
|
|
}
|
|
|
|
func TestLoadConfigWithEmptyOverrides(t *testing.T) {
|
|
t.Setenv("RATE_LIMIT_ENDPOINT_OVERRIDES", "{}")
|
|
|
|
cfg, err := LoadConfig()
|
|
if err != nil {
|
|
t.Fatalf("LoadConfig failed: %v", err)
|
|
}
|
|
|
|
if len(cfg.EndpointOverrides) != 0 {
|
|
t.Errorf("Expected no overrides, got %d", len(cfg.EndpointOverrides))
|
|
}
|
|
}
|
|
|
|
func TestLoadConfigWithInvalidJSON(t *testing.T) {
|
|
t.Setenv("RATE_LIMIT_ENDPOINT_OVERRIDES", "invalid json")
|
|
|
|
cfg, err := LoadConfig()
|
|
if err != nil {
|
|
t.Fatalf("LoadConfig failed: %v", err)
|
|
}
|
|
|
|
// Should ignore invalid JSON and use defaults
|
|
if len(cfg.EndpointOverrides) != 0 {
|
|
t.Errorf("Expected no overrides for invalid JSON, got %d", len(cfg.EndpointOverrides))
|
|
}
|
|
}
|
|
|
|
func TestLoadConfigWithInvalidNumbers(t *testing.T) {
|
|
t.Setenv("RATE_LIMIT_GLOBAL_RATE", "invalid")
|
|
t.Setenv("RATE_LIMIT_DEFAULT_BURST", "not_a_number")
|
|
|
|
cfg, err := LoadConfig()
|
|
if err != nil {
|
|
t.Fatalf("LoadConfig failed: %v", err)
|
|
}
|
|
|
|
// Should use defaults for invalid values
|
|
if cfg.GlobalRate != 500 {
|
|
t.Errorf("Expected GlobalRate=500 (default), got %v", cfg.GlobalRate)
|
|
}
|
|
if cfg.DefaultBurst != 10 {
|
|
t.Errorf("Expected DefaultBurst=10 (default), got %d", cfg.DefaultBurst)
|
|
}
|
|
}
|