go-telegram-bot/user_management_test.go

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)
	}
}

// TestGetOrCreateUser tests the getOrCreateUser method of the Bot.
// It verifies that a new user is created when one does not exist,
// and an existing user is returned when one does exist.
func TestGetOrCreateUser(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 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
	config := BotConfig{
		ID:              "bot1",
		MemorySize:      10,
		MessagePerHour:  5,
		MessagePerDay:   10,
		TempBanDuration: "1m",
		SystemPrompts:   make(map[string]string),
		TelegramToken:   "YOUR_TELEGRAM_BOT_TOKEN",
		OwnerTelegramID: 123456789,
	}

	// 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")
	}

	// Create a new user
	newUser, err := bot.getOrCreateUser(987654321, "TestUser", false)
	if err != nil {
		t.Fatalf("Failed to create a new user: %v", err)
	}

	// Verify that the new user was created
	var userInDB User
	err = db.Where("telegram_id = ?", newUser.TelegramID).First(&userInDB).Error
	if err != nil {
		t.Fatalf("New user was not created in the database: %v", err)
	}

	// Get the existing user
	existingUser, err := bot.getOrCreateUser(987654321, "TestUser", false)
	if err != nil {
		t.Fatalf("Failed to get existing user: %v", err)
	}

	// Verify that the existing user is the same as the new user
	if existingUser.ID != userInDB.ID {
		t.Fatalf("Expected to get the existing user, but got a different user")
	}
}

// To ensure thread safety and avoid race conditions during testing,
// you can run the tests with the `-race` flag:
// go test -race -v