mirror of
https://github.com/HugeFrog24/go-telegram-bot.git
synced 2026-03-02 00:14:34 +00:00
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:
0
.gitattributes
vendored
Normal file → Executable file
0
.gitattributes
vendored
Normal file → Executable file
0
.github/workflows/go-ci.yaml
vendored
Normal file → Executable file
0
.github/workflows/go-ci.yaml
vendored
Normal file → Executable file
0
.gitignore
vendored
Normal file → Executable file
0
.gitignore
vendored
Normal file → Executable file
110
README.md
Executable file
110
README.md
Executable 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
0
anthropic.go
Normal file → Executable file
45
bot.go
Normal file → Executable file
45
bot.go
Normal file → Executable 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
|
||||
}
|
||||
|
||||
2
config.go
Normal file → Executable file
2
config.go
Normal file → Executable 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
0
config/default.json
Normal file → Executable file
2
database.go
Normal file → Executable file
2
database.go
Normal file → Executable 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
|
||||
}
|
||||
|
||||
35
examples/systemd/telegram-bot.service
Executable file
35
examples/systemd/telegram-bot.service
Executable 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
2
go.mod
Normal file → Executable 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
4
go.sum
Normal file → Executable 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
19
handlers.go
Normal file → Executable 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
27
logger.go
Executable 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
38
main.go
Normal file → Executable 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
rate_limiter.go
Normal file → Executable file
0
rate_limiter.go
Normal file → Executable file
103
rate_limiter_test.go
Normal file → Executable file
103
rate_limiter_test.go
Normal file → Executable 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
0
telegram_client.go
Normal file → Executable file
0
telegram_client_mock.go
Normal file → Executable file
0
telegram_client_mock.go
Normal file → Executable file
189
user_management_test.go
Executable file
189
user_management_test.go
Executable 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
|
||||
Reference in New Issue
Block a user