Mining/pkg/mining/ratelimiter_test.go

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) {
	rateLimiter := NewRateLimiter(10, 20)
	if rateLimiter == nil {
		t.Fatal("NewRateLimiter returned nil")
	}
	defer rateLimiter.Stop()

	if rateLimiter.requestsPerSecond != 10 {
		t.Errorf("Expected requestsPerSecond 10, got %d", rateLimiter.requestsPerSecond)
	}
	if rateLimiter.burst != 20 {
		t.Errorf("Expected burst 20, got %d", rateLimiter.burst)
	}
}

// rl.Stop() // idempotent — calling twice must not panic
func TestRatelimiter_Stop_Good(t *testing.T) {
	rateLimiter := NewRateLimiter(10, 20)

	defer func() {
		if r := recover(); r != nil {
			t.Errorf("Stop panicked: %v", r)
		}
	}()

	rateLimiter.Stop()
	rateLimiter.Stop()
}

// router.Use(rl.Middleware()) — allows requests within burst, rejects beyond
func TestRatelimiter_Middleware_Good(t *testing.T) {
	gin.SetMode(gin.TestMode)
	rateLimiter := NewRateLimiter(10, 5)
	defer rateLimiter.Stop()

	router := gin.New()
	router.Use(rateLimiter.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)
	rateLimiter := NewRateLimiter(10, 5)
	defer rateLimiter.Stop()

	router := gin.New()
	router.Use(rateLimiter.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)
	rateLimiter := NewRateLimiter(10, 2)
	defer rateLimiter.Stop()

	router := gin.New()
	router.Use(rateLimiter.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) {
	rateLimiter := NewRateLimiter(10, 5)
	defer rateLimiter.Stop()

	gin.SetMode(gin.TestMode)
	router := gin.New()
	router.Use(rateLimiter.Middleware())
	router.GET("/test", func(c *gin.Context) {
		c.String(http.StatusOK, "ok")
	})

	if count := rateLimiter.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 := rateLimiter.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 := rateLimiter.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)
	rateLimiter := NewRateLimiter(100, 1)
	defer rateLimiter.Stop()

	router := gin.New()
	router.Use(rateLimiter.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)
	}
}