diff --git a/net/middleware/throttler/throttler_test.go b/net/middleware/throttler/throttler_test.go new file mode 100644 index 0000000..4d2cb70 --- /dev/null +++ b/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) + } +}