package mining import ( "net/http" "net/http/httptest" "testing" "time" "github.com/gin-gonic/gin" ) // rl := NewRateLimiter(10, 20) // rl.Stop() func TestRatelimiter_NewRateLimiter_Good(t *testing.T) { rl := NewRateLimiter(10, 20) if rl == nil { t.Fatal("NewRateLimiter returned nil") } defer rl.Stop() if rl.requestsPerSecond != 10 { t.Errorf("Expected requestsPerSecond 10, got %d", rl.requestsPerSecond) } if rl.burst != 20 { t.Errorf("Expected burst 20, got %d", rl.burst) } } // rl.Stop() // idempotent — calling twice must not panic func TestRatelimiter_Stop_Good(t *testing.T) { rl := NewRateLimiter(10, 20) defer func() { if r := recover(); r != nil { t.Errorf("Stop panicked: %v", r) } }() rl.Stop() rl.Stop() } // router.Use(rl.Middleware()) — allows requests within burst, rejects beyond func TestRatelimiter_Middleware_Good(t *testing.T) { gin.SetMode(gin.TestMode) rl := NewRateLimiter(10, 5) defer rl.Stop() router := gin.New() router.Use(rl.Middleware()) router.GET("/test", func(c *gin.Context) { c.String(http.StatusOK, "ok") }) for i := 0; i < 5; i++ { req := httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w := httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Errorf("Request %d: expected 200, got %d", i+1, w.Code) } } } // router.Use(rl.Middleware()) — rejects the 6th request after burst is exhausted func TestRatelimiter_Middleware_Bad(t *testing.T) { gin.SetMode(gin.TestMode) rl := NewRateLimiter(10, 5) defer rl.Stop() router := gin.New() router.Use(rl.Middleware()) router.GET("/test", func(c *gin.Context) { c.String(http.StatusOK, "ok") }) for i := 0; i < 5; i++ { req := httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w := httptest.NewRecorder() router.ServeHTTP(w, req) } req := httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w := httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != http.StatusTooManyRequests { t.Errorf("Expected 429 Too Many Requests, got %d", w.Code) } } // router.Use(rl.Middleware()) — rate limit per IP; exhausted IP1 does not affect IP2 func TestRatelimiter_Middleware_Ugly(t *testing.T) { gin.SetMode(gin.TestMode) rl := NewRateLimiter(10, 2) defer rl.Stop() router := gin.New() router.Use(rl.Middleware()) router.GET("/test", func(c *gin.Context) { c.String(http.StatusOK, "ok") }) for i := 0; i < 2; i++ { req := httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w := httptest.NewRecorder() router.ServeHTTP(w, req) } req := httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w := httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != http.StatusTooManyRequests { t.Errorf("IP1 should be rate limited, got %d", w.Code) } req = httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.2:12345" w = httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Errorf("IP2 should not be rate limited, got %d", w.Code) } } // count := rl.ClientCount() // returns number of tracked IPs func TestRatelimiter_ClientCount_Good(t *testing.T) { rl := NewRateLimiter(10, 5) defer rl.Stop() gin.SetMode(gin.TestMode) router := gin.New() router.Use(rl.Middleware()) router.GET("/test", func(c *gin.Context) { c.String(http.StatusOK, "ok") }) if count := rl.ClientCount(); count != 0 { t.Errorf("Expected 0 clients, got %d", count) } req := httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w := httptest.NewRecorder() router.ServeHTTP(w, req) if count := rl.ClientCount(); count != 1 { t.Errorf("Expected 1 client, got %d", count) } req = httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.2:12345" w = httptest.NewRecorder() router.ServeHTTP(w, req) if count := rl.ClientCount(); count != 2 { t.Errorf("Expected 2 clients, got %d", count) } } // rl.Middleware() — token refills at requestsPerSecond rate; request succeeds after wait func TestRatelimiter_TokenRefill_Good(t *testing.T) { gin.SetMode(gin.TestMode) rl := NewRateLimiter(100, 1) defer rl.Stop() router := gin.New() router.Use(rl.Middleware()) router.GET("/test", func(c *gin.Context) { c.String(http.StatusOK, "ok") }) req := httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w := httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Errorf("First request should succeed, got %d", w.Code) } req = httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w = httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != http.StatusTooManyRequests { t.Errorf("Second request should be rate limited, got %d", w.Code) } time.Sleep(20 * time.Millisecond) req = httptest.NewRequest("GET", "/test", nil) req.RemoteAddr = "192.168.1.1:12345" w = httptest.NewRecorder() router.ServeHTTP(w, req) if w.Code != http.StatusOK { t.Errorf("Third request should succeed after refill, got %d", w.Code) } }