add a lot

net/middleware/throttler/throttler_test.go

@@ -0,0 +1,63 @@
+package throttler
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestController_Add(t *testing.T) {
+	resetCycle := 2 * time.Second
+	limit := uint64(3)
+	controller := New(resetCycle, limit)
+
+	// Test case: Add within limit
+	assert.True(t, controller.Add("user1"), "Expected to add user1 successfully within limit")
+	assert.True(t, controller.Add("user1"), "Expected to add user1 successfully within limit")
+	assert.True(t, controller.Add("user1"), "Expected to add user1 successfully within limit")
+
+	// Test case: Surpass limit
+	assert.False(t, controller.Add("user1"), "Expected to not add user1 as it surpasses the limit")
+
+	// Test case: Reset after cycle
+	time.Sleep(resetCycle)
+	assert.True(t, controller.Add("user1"), "Expected to add user1 successfully after reset cycle")
+}
+
+func TestController_AutoClean(t *testing.T) {
+	resetCycle := 100 * time.Millisecond
+	cleanCycle := 50 * time.Millisecond
+	limit := uint64(1)
+	controller := New(resetCycle, limit)
+
+	controller.Add("user1")
+	controller.Add("user2")
+
+	controller.AutoClean(cleanCycle)
+	time.Sleep(resetCycle + 50*time.Millisecond) // Wait for reset cycle and a bit more
+
+	assert.Equal(t, 0, len(controller.bucket), "Expected bucket to be cleaned after reset cycle")
+}
+
+func BenchmarkController_Add(b *testing.B) {
+	controller := New(2*time.Second, 10)
+
+	for i := 0; i < b.N; i++ {
+		controller.Add("user1")
+	}
+}
+
+func BenchmarkController_AutoClean(b *testing.B) {
+	resetCycle := 1 * time.Second
+	cleanCycle := 500 * time.Millisecond
+	controller := New(resetCycle, 10)
+
+	controller.AutoClean(cleanCycle)
+	defer controller.StopCleaner()
+
+	for i := 0; i < b.N; i++ {
+		controller.Add("user" + string(rune(i)))
+		time.Sleep(10 * time.Millisecond)
+	}
+}