Upgrade dependencies

Added tests, revised logging

Removed dependency on env file

Try reformatting unit file

Comments clarification

Added readme

Added readme
This commit is contained in:
HugeFrog24
2024-10-22 20:07:06 +02:00
parent c8af457af1
commit d8d0da4704
22 changed files with 420 additions and 156 deletions

0
.gitattributes vendored Normal file → Executable file
View File

0
.github/workflows/go-ci.yaml vendored Normal file → Executable file
View File

0
.gitignore vendored Normal file → Executable file
View File

110
README.md Executable file
View File

@@ -0,0 +1,110 @@
# Go Telegram Multibot
A scalable, multi-bot solution for Telegram using Go, GORM, and the Anthropic API.
## Design Considerations
- AI-powered
- Supports multiple bot profiles
- Uses SQLite for persistence
- Implements rate limiting and user management
- Modular architecture
- Comprehensive unit tests
## Usage
1. Clone the repository or install using `go get`:
- Option 1: Clone the repository
```bash
git clone https://github.com/HugeFrog24/go-telegram-bot.git
```
- Option 2: Install using go get
```bash
go get -u github.com/HugeFrog24/go-telegram-bot
```
- Navigate to the project directory:
```bash
cd go-telegram-bot
```
2. Copy the default config template and edit it:
```bash
cp config/default.json config/config-mybot.json
```
Replace `config-mybot.json` with the name of your bot.
```bash
nano config/config-mybot.json
```
You can set up as many bots as you want. Just copy the template and edit the parameters.
> [!IMPORTANT]
> Keep your config files secret and do not commit them to version control.
3. Build the application:
```bash
go build -o telegram-multibot
```
## Systemd Unit Setup
To enable the bot to start automatically on system boot and run in the background, set up a systemd unit.
1. Copy the systemd unit template and edit it:
```bash
sudo cp examples/systemd/telegram-bot.service /etc/systemd/system/telegram-bot.service
```
Edit the service file:
```bash
nano /etc/systemd/system/telegram-bot.service
```
Adjust the following parameters:
- `WorkingDirectory`
- `ExecStart`
- `User`
3. Enable and start the service:
```bash
sudo systemctl daemon-reload
```
```bash
sudo systemctl enable telegram-bot.service
```
```bash
sudo systemctl start telegram-bot.service
```
4. Check the status:
```bash
sudo systemctl status telegram-bot
```
For more details on the systemd setup, refer to the [demo service file](examples/systemd/telegram-bot.service).
## Logs
View logs using journalctl:
```bash
journalctl -u telegram-bot
```
Follow logs:
```bash
journalctl -u telegram-bot -f
```
View errors:
```bash
journalctl -u telegram-bot -p err
```

0
anthropic.go Normal file → Executable file
View File

45
bot.go Normal file → Executable file
View File

@@ -4,7 +4,6 @@ import (
"context"
"errors"
"fmt"
"log"
"strings"
"sync"
"time"
@@ -282,7 +281,7 @@ func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, busin
// Pass the outgoing message through the centralized screen for storage
_, err := b.screenOutgoingMessage(chatID, text, businessConnectionID)
if err != nil {
log.Printf("Error storing assistant message: %v", err)
ErrorLogger.Printf("Error storing assistant message: %v", err)
return err
}
@@ -299,7 +298,7 @@ func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, busin
// Send the message via Telegram client
_, err = b.tgBot.SendMessage(ctx, params)
if err != nil {
log.Printf("[%s] [ERROR] Error sending message to chat %d with BusinessConnectionID %s: %v",
ErrorLogger.Printf("[%s] Error sending message to chat %d with BusinessConnectionID %s: %v",
b.config.ID, chatID, businessConnectionID, err)
return err
}
@@ -310,9 +309,9 @@ func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, busin
func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, username string, businessConnectionID string) {
totalUsers, totalMessages, err := b.getStats()
if err != nil {
fmt.Printf("Error fetching stats: %v\n", err)
ErrorLogger.Printf("Error fetching stats: %v\n", err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't retrieve the stats at this time.", businessConnectionID); err != nil {
log.Printf("Error sending response: %v", err)
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
@@ -328,7 +327,7 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, usernam
// Send the response through the centralized screen
if err := b.sendResponse(ctx, chatID, statsMessage, businessConnectionID); err != nil {
log.Printf("Error sending stats message: %v", err)
ErrorLogger.Printf("Error sending stats message: %v", err)
}
}
@@ -370,18 +369,18 @@ func isEmoji(r rune) bool {
func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, username string, businessConnectionID string) {
user, err := b.getOrCreateUser(userID, username, false)
if err != nil {
log.Printf("Error getting or creating user: %v", err)
ErrorLogger.Printf("Error getting or creating user: %v", err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't retrieve your information.", businessConnectionID); err != nil {
log.Printf("Error sending response: %v", err)
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
role, err := b.getRoleByName(user.Role.Name)
if err != nil {
log.Printf("Error getting role by name: %v", err)
ErrorLogger.Printf("Error getting role by name: %v", err)
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't retrieve your role information.", businessConnectionID); err != nil {
log.Printf("Error sending response: %v", err)
ErrorLogger.Printf("Error sending response: %v", err)
}
return
}
@@ -396,7 +395,7 @@ func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, userna
// Send the response through the centralized screen
if err := b.sendResponse(ctx, chatID, whoAmIMessage, businessConnectionID); err != nil {
log.Printf("Error sending /whoami message: %v", err)
ErrorLogger.Printf("Error sending /whoami message: %v", err)
}
}
@@ -440,3 +439,27 @@ func (b *Bot) screenOutgoingMessage(chatID int64, response string, businessConne
return assistantMessage, nil
}
func (b *Bot) promoteUserToAdmin(promoterID, userToPromoteID int64) error {
// Check if the promoter is an owner or admin
if !b.isAdminOrOwner(promoterID) {
return errors.New("only admins or owners can promote users to admin")
}
// Get the user to promote
userToPromote, err := b.getOrCreateUser(userToPromoteID, "", false)
if err != nil {
return err
}
// Get the admin role
var adminRole Role
if err := b.db.Where("name = ?", "admin").First(&adminRole).Error; err != nil {
return err
}
// Update the user's role
userToPromote.RoleID = adminRole.ID
userToPromote.Role = adminRole
return b.db.Save(&userToPromote).Error
}

0
clock.go Normal file → Executable file
View File

2
config.go Normal file → Executable file
View File

@@ -61,7 +61,7 @@ func loadAllConfigs(dir string) ([]BotConfig, error) {
// Skip inactive bots
if !config.Active {
fmt.Printf("Skipping inactive bot: %s\n", config.ID)
InfoLogger.Printf("Skipping inactive bot: %s", config.ID)
continue
}

0
config/default.json Normal file → Executable file
View File

2
database.go Normal file → Executable file
View File

@@ -56,8 +56,10 @@ func createDefaultRoles(db *gorm.DB) error {
for _, roleName := range roles {
var role Role
if err := db.FirstOrCreate(&role, Role{Name: roleName}).Error; err != nil {
ErrorLogger.Printf("Failed to create default role %s: %v", roleName, err)
return fmt.Errorf("failed to create default role %s: %w", roleName, err)
}
InfoLogger.Printf("Created or confirmed default role: %s", roleName)
}
return nil
}

View File

@@ -0,0 +1,35 @@
[Unit]
# A concise description of the service
Description=Telegram Bot Service
# Postpone starting until network is available
After=network.target
[Service]
# The user that runs the bot
User=tibik
# The directory where the bot is located
WorkingDirectory=/home/tibik/go-telegram-bot
# The command to start the bot
ExecStart=/home/tibik/go-telegram-bot/telegram-bot
# Restart if crashed
Restart=always
# Delay between restarts to avoid resource exhaustion
RestartSec=5
# Capture stdout (INFO logs)
StandardOutput=journal
# Capture stderr (ERROR logs)
StandardError=journal
# Identifier for journalctl filtering
SyslogIdentifier=telegram-bot
[Install]
# The bot will start automatically at system boot
WantedBy=multi-user.target
# NOTE:
# New line comments: good
# Inline comments: no good, they mess up the service file
# View logs: journalctl -u telegram-bot
# Follow logs: journalctl -u telegram-bot -f
# View errors: journalctl -u telegram-bot -p err

2
go.mod Normal file → Executable file
View File

@@ -3,7 +3,7 @@ module github.com/HugeFrog24/go-telegram-bot
go 1.23
require (
github.com/go-telegram/bot v1.9.0
github.com/go-telegram/bot v1.9.1
github.com/joho/godotenv v1.5.1
github.com/liushuangls/go-anthropic/v2 v2.8.1
golang.org/x/time v0.7.0

4
go.sum Normal file → Executable file
View File

@@ -1,5 +1,5 @@
github.com/go-telegram/bot v1.9.0 h1:z9g0Fgk9B7G/xoVMqji30hpJPlr3Dz3aVW2nzSGfPuI=
github.com/go-telegram/bot v1.9.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/go-telegram/bot v1.9.1 h1:4vkNV6vDmEPZaYP7sZYaagOaJyV4GerfOPkjg/Ki5ic=
github.com/go-telegram/bot v1.9.1/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc=
github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ=

19
handlers.go Normal file → Executable file
View File

@@ -2,7 +2,6 @@ package main
import (
"context"
"log"
"strings"
"github.com/go-telegram/bot"
@@ -38,7 +37,7 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
// Pass the incoming message through the centralized screen for storage
_, err := b.screenIncomingMessage(message)
if err != nil {
log.Printf("Error storing user message: %v", err)
ErrorLogger.Printf("Error storing user message: %v", err)
return
}
@@ -73,7 +72,7 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
// Proceed only if the message contains text
if text == "" {
log.Printf("Received a non-text message from user %d in chat %d", userID, chatID)
InfoLogger.Printf("Received a non-text message from user %d in chat %d", userID, chatID)
return
}
@@ -86,7 +85,7 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
user, err := b.getOrCreateUser(userID, username, isOwner)
if err != nil {
log.Printf("Error getting or creating user: %v", err)
ErrorLogger.Printf("Error getting or creating user: %v", err)
return
}
@@ -94,7 +93,7 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
if user.Username != username {
user.Username = username
if err := b.db.Save(&user).Error; err != nil {
log.Printf("Error updating user username: %v", err)
ErrorLogger.Printf("Error updating user username: %v", err)
}
}
@@ -109,20 +108,20 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
// Get response from Anthropic
response, err := b.getAnthropicResponse(ctx, contextMessages, b.isNewChat(chatID), isOwner, isEmojiOnly)
if err != nil {
log.Printf("Error getting Anthropic response: %v", err)
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
response = "I'm sorry, I'm having trouble processing your request right now."
}
// Send the response through the centralized screen
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
log.Printf("Error sending response: %v", err)
ErrorLogger.Printf("Error sending response: %v", err)
return
}
}
func (b *Bot) sendRateLimitExceededMessage(ctx context.Context, chatID int64, businessConnectionID string) {
if err := b.sendResponse(ctx, chatID, "Rate limit exceeded. Please try again later.", businessConnectionID); err != nil {
log.Printf("Error sending rate limit exceeded message: %v", err)
ErrorLogger.Printf("Error sending rate limit exceeded message: %v", err)
}
}
@@ -145,7 +144,7 @@ func (b *Bot) handleStickerMessage(ctx context.Context, chatID, userID int64, me
// Generate AI response about the sticker
response, err := b.generateStickerResponse(ctx, userMessage)
if err != nil {
log.Printf("Error generating sticker response: %v", err)
ErrorLogger.Printf("Error generating sticker response: %v", err)
// Provide a fallback dynamic response based on sticker type
if message.Sticker.IsAnimated {
response = "Wow, that's a cool animated sticker!"
@@ -158,7 +157,7 @@ func (b *Bot) handleStickerMessage(ctx context.Context, chatID, userID int64, me
// Send the response through the centralized screen
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
log.Printf("Error sending response: %v", err)
ErrorLogger.Printf("Error sending response: %v", err)
return
}
}

27
logger.go Executable file
View File

@@ -0,0 +1,27 @@
package main
import (
"log"
"os"
)
// For log management, use journalctl commands:
// - View logs: journalctl -u telegram-bot
// - Follow logs: journalctl -u telegram-bot -f
// - View errors: journalctl -u telegram-bot -p err
// Refer to the documentation for details on systemd unit setup.
// Initialize loggers for informational and error messages.
var (
InfoLogger *log.Logger
ErrorLogger *log.Logger
)
// initLoggers sets up separate loggers for stdout and stderr.
func initLoggers() {
// InfoLogger writes to stdout with specific flags.
InfoLogger = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
// ErrorLogger writes to stderr with specific flags.
ErrorLogger = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
}

38
main.go Normal file → Executable file
View File

@@ -2,38 +2,28 @@ package main
import (
"context"
"io"
"log"
"os"
"os/signal"
"sync"
"github.com/joho/godotenv"
)
func main() {
// Initialize logger
logFile, err := initLogger()
if err != nil {
log.Fatalf("Error initializing logger: %v", err)
}
defer logFile.Close()
// Initialize custom loggers
initLoggers()
// Load environment variables
if err := godotenv.Load(); err != nil {
log.Printf("Error loading .env file: %v", err)
}
// Log the start of the application
InfoLogger.Println("Starting Telegram Bot Application")
// Initialize database
db, err := initDB()
if err != nil {
log.Fatalf("Error initializing database: %v", err)
ErrorLogger.Fatalf("Error initializing database: %v", err)
}
// Load all bot configurations
configs, err := loadAllConfigs("config")
if err != nil {
log.Fatalf("Error loading configurations: %v", err)
ErrorLogger.Fatalf("Error loading configurations: %v", err)
}
// Create a WaitGroup to manage goroutines
@@ -53,14 +43,14 @@ func main() {
realClock := RealClock{}
bot, err := NewBot(db, cfg, realClock, nil)
if err != nil {
log.Printf("Error creating bot %s: %v", cfg.ID, err)
ErrorLogger.Printf("Error creating bot %s: %v", cfg.ID, err)
return
}
// Initialize TelegramClient with the bot's handleUpdate method
tgClient, err := initTelegramBot(cfg.TelegramToken, bot.handleUpdate)
if err != nil {
log.Printf("Error initializing Telegram client for bot %s: %v", cfg.ID, err)
ErrorLogger.Printf("Error initializing Telegram client for bot %s: %v", cfg.ID, err)
return
}
@@ -68,21 +58,13 @@ func main() {
bot.tgBot = tgClient
// Start the bot
log.Printf("Starting bot %s...", cfg.ID)
InfoLogger.Printf("Starting bot %s...", cfg.ID)
bot.Start(ctx)
}(config)
}
// Wait for all bots to finish
wg.Wait()
}
func initLogger() (*os.File, error) {
logFile, err := os.OpenFile("bot.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
if err != nil {
return nil, err
}
mw := io.MultiWriter(os.Stdout, logFile)
log.SetOutput(mw)
return logFile, nil
InfoLogger.Println("All bots have stopped. Exiting application.")
}

0
models.go Normal file → Executable file
View File

0
rate_limiter.go Normal file → Executable file
View File

103
rate_limiter_test.go Normal file → Executable file
View File

@@ -1,15 +1,8 @@
package main
import (
"context"
"fmt"
"testing"
"time"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
)
// TestCheckRateLimits tests the checkRateLimits method of the Bot.
@@ -87,102 +80,6 @@ func TestCheckRateLimits(t *testing.T) {
}
}
func TestOwnerAssignment(t *testing.T) {
// Initialize in-memory database for testing
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("Failed to open in-memory database: %v", err)
}
// Migrate the schema
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
if err != nil {
t.Fatalf("Failed to migrate database schema: %v", err)
}
// Create default roles
err = createDefaultRoles(db)
if err != nil {
t.Fatalf("Failed to create default roles: %v", err)
}
// Create a bot configuration
config := BotConfig{
ID: "test_bot",
TelegramToken: "TEST_TELEGRAM_TOKEN",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1m",
SystemPrompts: make(map[string]string),
Active: true,
OwnerTelegramID: 111111111,
}
// Initialize MockClock
mockClock := &MockClock{
currentTime: time.Now(),
}
// Initialize MockTelegramClient
mockTGClient := &MockTelegramClient{
SendMessageFunc: func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
chatID, ok := params.ChatID.(int64)
if !ok {
return nil, fmt.Errorf("ChatID is not of type int64")
}
// Simulate successful message sending
return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil
},
}
// Create the bot with the mock Telegram client
bot, err := NewBot(db, config, mockClock, mockTGClient)
if err != nil {
t.Fatalf("Failed to create bot: %v", err)
}
// Verify that the owner exists
var owner User
err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error
if err != nil {
t.Fatalf("Owner was not created: %v", err)
}
// Attempt to create another owner for the same bot
_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
if err == nil {
t.Fatalf("Expected error when creating a second owner, but got none")
}
// Verify that the error message is appropriate
expectedErrorMsg := "an owner already exists for this bot"
if err.Error() != expectedErrorMsg {
t.Fatalf("Unexpected error message: %v", err)
}
// Assign admin role to a new user
adminUser, err := bot.getOrCreateUser(333333333, "AdminUser", false)
if err != nil {
t.Fatalf("Failed to create admin user: %v", err)
}
if adminUser.Role.Name != "admin" {
t.Fatalf("Expected role 'admin', got '%s'", adminUser.Role.Name)
}
// Attempt to change an existing user to owner
_, err = bot.getOrCreateUser(333333333, "AdminUser", true)
if err == nil {
t.Fatalf("Expected error when changing existing user to owner, but got none")
}
expectedErrorMsg = "cannot change existing user to owner"
if err.Error() != expectedErrorMsg {
t.Fatalf("Unexpected error message: %v", err)
}
}
// To ensure thread safety and avoid race conditions during testing,
// you can run the tests with the `-race` flag:
// go test -race -v

0
telegram_client.go Normal file → Executable file
View File

0
telegram_client_mock.go Normal file → Executable file
View File

189
user_management_test.go Executable file
View File

@@ -0,0 +1,189 @@
package main
import (
"context"
"fmt"
"testing"
"time"
"github.com/go-telegram/bot"
"github.com/go-telegram/bot/models"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
)
func TestOwnerAssignment(t *testing.T) {
// Initialize loggers
initLoggers()
// Initialize in-memory database for testing
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("Failed to open in-memory database: %v", err)
}
// Migrate the schema
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
if err != nil {
t.Fatalf("Failed to migrate database schema: %v", err)
}
// Create default roles
err = createDefaultRoles(db)
if err != nil {
t.Fatalf("Failed to create default roles: %v", err)
}
// Create a bot configuration
config := BotConfig{
ID: "test_bot",
TelegramToken: "TEST_TELEGRAM_TOKEN",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1m",
SystemPrompts: make(map[string]string),
Active: true,
OwnerTelegramID: 111111111,
}
// Initialize MockClock
mockClock := &MockClock{
currentTime: time.Now(),
}
// Initialize MockTelegramClient
mockTGClient := &MockTelegramClient{
SendMessageFunc: func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) {
chatID, ok := params.ChatID.(int64)
if !ok {
return nil, fmt.Errorf("ChatID is not of type int64")
}
// Simulate successful message sending
return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil
},
}
// Create the bot with the mock Telegram client
bot, err := NewBot(db, config, mockClock, mockTGClient)
if err != nil {
t.Fatalf("Failed to create bot: %v", err)
}
// Verify that the owner exists
var owner User
err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error
if err != nil {
t.Fatalf("Owner was not created: %v", err)
}
// Attempt to create another owner for the same bot
_, err = bot.getOrCreateUser(222222222, "AnotherOwner", true)
if err == nil {
t.Fatalf("Expected error when creating a second owner, but got none")
}
// Verify that the error message is appropriate
expectedErrorMsg := "an owner already exists for this bot"
if err.Error() != expectedErrorMsg {
t.Fatalf("Unexpected error message: %v", err)
}
// Assign admin role to a new user
regularUser, err := bot.getOrCreateUser(333333333, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to create regular user: %v", err)
}
if regularUser.Role.Name != "user" {
t.Fatalf("Expected role 'user', got '%s'", regularUser.Role.Name)
}
// Attempt to change an existing user to owner
_, err = bot.getOrCreateUser(333333333, "AdminUser", true)
if err == nil {
t.Fatalf("Expected error when changing existing user to owner, but got none")
}
expectedErrorMsg = "cannot change existing user to owner"
if err.Error() != expectedErrorMsg {
t.Fatalf("Unexpected error message: %v", err)
}
// If you need to test admin creation, you should do it through a separate admin creation function
// or by updating an existing user's role with proper authorization checks
}
func TestPromoteUserToAdmin(t *testing.T) {
// Initialize loggers
initLoggers()
// Initialize in-memory database for testing
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("Failed to open in-memory database: %v", err)
}
// Migrate the schema
err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{})
if err != nil {
t.Fatalf("Failed to migrate database schema: %v", err)
}
// Create default roles
err = createDefaultRoles(db)
if err != nil {
t.Fatalf("Failed to create default roles: %v", err)
}
config := BotConfig{
ID: "test_bot",
TelegramToken: "TEST_TELEGRAM_TOKEN",
MemorySize: 10,
MessagePerHour: 5,
MessagePerDay: 10,
TempBanDuration: "1m",
SystemPrompts: make(map[string]string),
Active: true,
OwnerTelegramID: 111111111,
}
mockClock := &MockClock{currentTime: time.Now()}
mockTGClient := &MockTelegramClient{}
bot, err := NewBot(db, config, mockClock, mockTGClient)
if err != nil {
t.Fatalf("Failed to create bot: %v", err)
}
// Create an owner
owner, err := bot.getOrCreateUser(config.OwnerTelegramID, "OwnerUser", true)
if err != nil {
t.Fatalf("Failed to create owner: %v", err)
}
// Test promoting a user to admin
regularUser, err := bot.getOrCreateUser(444444444, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to create regular user: %v", err)
}
err = bot.promoteUserToAdmin(owner.TelegramID, regularUser.TelegramID)
if err != nil {
t.Fatalf("Failed to promote user to admin: %v", err)
}
// Refresh user data
promotedUser, err := bot.getOrCreateUser(444444444, "RegularUser", false)
if err != nil {
t.Fatalf("Failed to get promoted user: %v", err)
}
if promotedUser.Role.Name != "admin" {
t.Fatalf("Expected role 'admin', got '%s'", promotedUser.Role.Name)
}
}
// To ensure thread safety and avoid race conditions during testing,
// you can run the tests with the `-race` flag:
// go test -race -v