Adaptive thinking

rate_limiter_test.go

@@ -0,0 +1,67 @@
+package main
+
+import (
+	"testing"
+	"time"
+)
+
+func TestCheckRateLimits(t *testing.T) {
+	mockClock := &MockClock{
+		currentTime: time.Date(2023, 10, 1, 0, 0, 0, 0, time.UTC),
+	}
+
+	config := BotConfig{
+		ID:              "bot1",
+		MemorySize:      10,
+		MessagePerHour:  5,
+		MessagePerDay:   10,
+		TempBanDuration: "1m",
+		SystemPrompts:   make(map[string]string),
+		TelegramToken:   "YOUR_TELEGRAM_BOT_TOKEN",
+		OwnerTelegramID: 123456789,
+	}
+
+	bot := &Bot{
+		config:       config,
+		userLimiters: make(map[int64]*userLimiter),
+		clock:        mockClock,
+	}
+
+	userID := int64(12345)
+
+	sendMessage := func() bool {
+		return bot.checkRateLimits(userID)
+	}
+
+	for i := 0; i < config.MessagePerHour; i++ {
+		if !sendMessage() {
+			t.Errorf("Expected message %d to be allowed", i+1)
+		}
+	}
+
+	if sendMessage() {
+		t.Errorf("Expected message to be denied due to hourly limit exceeded")
+	}
+
+	if sendMessage() {
+		t.Errorf("Expected message to be denied while user is banned")
+	}
+
+	mockClock.Advance(time.Minute)
+
+	mockClock.Advance(time.Hour)
+
+	if !sendMessage() {
+		t.Errorf("Expected message to be allowed after ban duration")
+	}
+
+	for i := 0; i < config.MessagePerDay-config.MessagePerHour-1; i++ {
+		if !sendMessage() {
+			t.Errorf("Expected message %d to be allowed towards daily limit", i+1)
+		}
+	}
+
+	if sendMessage() {
+		t.Errorf("Expected message to be denied due to daily limit exceeded")
+	}
+}