1 files changed, 128 insertions, 0 deletions
diff --git a/internal/bot/ratelimit_test.go b/internal/bot/ratelimit_test.go
new file mode 100644
index 0000000..7f09787
--- /dev/null
+++ b/internal/bot/ratelimit_test.go
@@ -0,0 +1,128 @@
+package bot
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+type fakeClock struct {
+	mu  sync.Mutex
+	now time.Time
+}
+
+func (c *fakeClock) Now() time.Time {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.now
+}
+
+func (c *fakeClock) Advance(d time.Duration) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.now = c.now.Add(d)
+}
+
+func TestLimiterAllowsTheBurstThenStops(t *testing.T) {
+	clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
+	limiter := NewLimiterWith(3, time.Second, 30*time.Second, clock.Now)
+
+	for i := 0; i < 3; i++ {
+		if decision := limiter.Allow(1); !decision.Allowed {
+			t.Fatalf("request %d was refused inside the burst", i+1)
+		}
+	}
+	decision := limiter.Allow(1)
+	if decision.Allowed {
+		t.Fatal("the fourth request should be refused")
+	}
+	if !decision.Warn {
+		t.Fatal("the first refusal should warn the user")
+	}
+	if decision.RetryAfter <= 0 {
+		t.Fatalf("retry after = %v, want a positive hint", decision.RetryAfter)
+	}
+}
+
+func TestLimiterWarnsOncePerWindow(t *testing.T) {
+	clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
+	limiter := NewLimiterWith(1, time.Minute, 30*time.Second, clock.Now)
+
+	limiter.Allow(1)
+	if !limiter.Allow(1).Warn {
+		t.Fatal("the first refusal should warn")
+	}
+	if limiter.Allow(1).Warn {
+		t.Fatal("a second refusal inside the window should stay silent")
+	}
+
+	clock.Advance(31 * time.Second)
+	if !limiter.Allow(1).Warn {
+		t.Fatal("a refusal after the window should warn again")
+	}
+}
+
+func TestLimiterRefills(t *testing.T) {
+	clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
+	limiter := NewLimiterWith(2, time.Second, time.Minute, clock.Now)
+
+	limiter.Allow(1)
+	limiter.Allow(1)
+	if limiter.Allow(1).Allowed {
+		t.Fatal("the bucket should be empty")
+	}
+
+	clock.Advance(time.Second)
+	if !limiter.Allow(1).Allowed {
+		t.Fatal("one second should have bought one token")
+	}
+	if limiter.Allow(1).Allowed {
+		t.Fatal("only one token should have been refilled")
+	}
+}
+
+func TestLimiterNeverExceedsTheBurst(t *testing.T) {
+	clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
+	limiter := NewLimiterWith(2, time.Second, time.Minute, clock.Now)
+
+	limiter.Allow(1)
+	clock.Advance(time.Hour)
+
+	if !limiter.Allow(1).Allowed || !limiter.Allow(1).Allowed {
+		t.Fatal("the bucket should be full again")
+	}
+	if limiter.Allow(1).Allowed {
+		t.Fatal("an idle hour must not bank more than the burst")
+	}
+}
+
+func TestLimiterKeysAreIndependent(t *testing.T) {
+	clock := &fakeClock{now: time.Unix(1_700_000_000, 0)}
+	limiter := NewLimiterWith(1, time.Minute, time.Minute, clock.Now)
+
+	if !limiter.Allow(1).Allowed {
+		t.Fatal("first user refused")
+	}
+	if limiter.Allow(1).Allowed {
+		t.Fatal("first user should now be limited")
+	}
+	if !limiter.Allow(2).Allowed {
+		t.Fatal("one user's flood must not silence another")
+	}
+}
+
+func TestLimiterIsConcurrencySafe(t *testing.T) {
+	limiter := NewLimiterWith(100, time.Millisecond, time.Second, nil)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 50; i++ {
+		wg.Add(1)
+		go func(key int64) {
+			defer wg.Done()
+			for j := 0; j < 20; j++ {
+				limiter.Allow(key % 5)
+			}
+		}(int64(i))
+	}
+	wg.Wait()
+}