Multibot finished

rate_limiter_test.go

@@ -0,0 +1,84 @@
+package main
+
+import (
+	"testing"
+	"time"
+)
+
+// TestCheckRateLimits tests the checkRateLimits method of the Bot.
+// It verifies that users are allowed or denied based on their message rates.
+func TestCheckRateLimits(t *testing.T) {
+	// Create a mock clock starting at a fixed time
+	mockClock := &MockClock{
+		currentTime: time.Date(2023, 10, 1, 0, 0, 0, 0, time.UTC),
+	}
+
+	// Create a mock configuration with reduced timeframes for testing
+	config := BotConfig{
+		ID:              "bot1",
+		MemorySize:      10,
+		MessagePerHour:  5,    // Allow 5 messages per hour
+		MessagePerDay:   10,   // Allow 10 messages per day
+		TempBanDuration: "1m", // Temporary ban duration of 1 minute for testing
+		SystemPrompts:   make(map[string]string),
+		TelegramToken:   "YOUR_TELEGRAM_BOT_TOKEN",
+	}
+
+	// Initialize the Bot with mock data and MockClock
+	bot := &Bot{
+		config:       config,
+		userLimiters: make(map[int64]*userLimiter),
+		clock:        mockClock,
+	}
+
+	userID := int64(12345)
+
+	// Helper function to simulate message sending
+	sendMessage := func() bool {
+		return bot.checkRateLimits(userID)
+	}
+
+	// Send 5 messages within the hourly limit
+	for i := 0; i < config.MessagePerHour; i++ {
+		if !sendMessage() {
+			t.Errorf("Expected message %d to be allowed", i+1)
+		}
+	}
+
+	// 6th message should exceed the hourly limit and trigger a ban
+	if sendMessage() {
+		t.Errorf("Expected message to be denied due to hourly limit exceeded")
+	}
+
+	// Attempt to send another message immediately, should still be banned
+	if sendMessage() {
+		t.Errorf("Expected message to be denied while user is banned")
+	}
+
+	// Fast-forward time by TempBanDuration to lift the ban
+	mockClock.Advance(time.Minute) // Banned for 1 minute
+
+	// Advance time to allow hourly limiter to replenish
+	mockClock.Advance(time.Hour) // Advance by 1 hour
+
+	// Send another message, should be allowed now
+	if !sendMessage() {
+		t.Errorf("Expected message to be allowed after ban duration")
+	}
+
+	// Send additional messages to reach the daily limit
+	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)
+		}
+	}
+
+	// Attempt to exceed the daily limit
+	if sendMessage() {
+		t.Errorf("Expected message to be denied due to daily limit exceeded")
+	}
+}
+
+// To ensure thread safety and avoid race conditions during testing,
+// you can run the tests with the `-race` flag:
+// go test -race -v