go-telegram-bot/rate_limiter_test.go

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