Files
vibe-bot/vibe_bot/tests/test_main.py
T
2026-05-24 15:26:13 -04:00

807 lines
22 KiB
Python

"""Tests for the main module (Discord bot commands)."""
from __future__ import annotations
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@pytest.fixture
def mock_ctx() -> MagicMock:
"""Create a mock Discord command context."""
ctx = MagicMock()
ctx.author.name = "testuser"
ctx.author.id = "12345"
ctx.author.global_name = "Test User"
ctx.author.nick = "tester"
ctx.author.top_role.name = "@everyone"
ctx.author.activities = []
ctx.author.joined_at = None
ctx.author.created_at = None
ctx.channel.id = "channel-1"
ctx.guild.id = "guild-1"
ctx.message.id = "msg-1"
ctx.message.attachments = []
ctx.bot.user = MagicMock()
ctx.bot.user.name = "test-bot"
ctx.bot.user.id = "bot-123"
ctx.send = AsyncMock()
return ctx
@pytest.fixture
def mock_ctx_with_member() -> MagicMock:
"""Create a mock Discord command context with full member data."""
ctx = MagicMock()
ctx.author.name = "testuser"
ctx.author.id = "12345"
ctx.author.global_name = "Test User"
ctx.author.nick = "tester"
ctx.author.top_role.name = "Admin"
mock_activity = MagicMock()
mock_activity.name = "Chess"
ctx.author.activities = [mock_activity]
from datetime import datetime
ctx.author.joined_at = datetime(2024, 1, 15)
ctx.author.created_at = datetime(2023, 6, 1)
ctx.channel.id = "channel-1"
ctx.guild.id = "guild-1"
ctx.message.id = "msg-1"
ctx.message.attachments = []
ctx.bot.user = MagicMock()
ctx.bot.user.name = "test-bot"
ctx.bot.user.id = "bot-123"
ctx.send = AsyncMock()
return ctx
def test_bot_initialized(mock_discord: dict[str, MagicMock]) -> None:
"""Test that the bot is initialized."""
import vibe_bot.main as main_module
assert main_module.bot is not None
def test_bot_intents_set(mock_discord: dict[str, MagicMock]) -> None:
"""Test that message_content intent is enabled."""
import vibe_bot.main as main_module
main_module.bot = mock_discord["bot_instance"]
assert main_module.MIN_BOT_NAME_LENGTH == 2
assert main_module.MAX_BOT_NAME_LENGTH == 50
assert main_module.MIN_PERSONALITY_LENGTH == 10
@patch("vibe_bot.main.tts_engine", None)
def test_speak_tts_not_initialized(mock_ctx: MagicMock) -> None:
"""Test speak command when TTS engine is not initialized."""
import asyncio
import vibe_bot.main as main_module
asyncio.run(main_module.speak(mock_ctx, message="hello world"))
mock_ctx.send.assert_called_once()
call_args = mock_ctx.send.call_args[0][0]
assert "TTS engine not initialized" in call_args
def test_speak_empty_message(
mock_ctx: MagicMock,
mock_tts_engine: MagicMock,
) -> None:
"""Test speak command with empty message."""
import asyncio
import vibe_bot.main as main_module
asyncio.run(main_module.speak(mock_ctx, message=""))
mock_ctx.send.assert_called_once()
call_args = mock_ctx.send.call_args[0][0]
assert "Please provide text" in call_args
def test_speak_plain_text(
mock_ctx: MagicMock,
mock_tts_engine: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test speak command with plain text (no custom bot prefix)."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.list_custom_bots.return_value = []
asyncio.run(main_module.speak(mock_ctx, message="hello world"))
mock_tts_engine.generate_audio.assert_called_once()
assert mock_ctx.send.call_count >= 2
def test_speak_with_custom_bot(
mock_ctx: MagicMock,
mock_tts_engine: MagicMock,
mock_custom_bot_manager: MagicMock,
mock_database: MagicMock,
mock_llama_wrapper: MagicMock,
) -> None:
"""Test speak command with a custom bot prefix."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.list_custom_bots.return_value = [
("alfred", "british butler", "user-123"),
]
mock_custom_bot_manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
asyncio.run(main_module.speak(mock_ctx, message="alfred what time is it"))
mock_llama_wrapper.chat_completion_with_tools.assert_called_once()
mock_tts_engine.generate_audio.assert_called_once()
assert mock_ctx.send.call_count >= 3
text_response = mock_ctx.send.call_args_list[1][0][0]
assert "**alfred**:" in text_response or "**alfred** :" in text_response
def test_custom_bot_command_success(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test creating a custom bot successfully."""
import asyncio
import vibe_bot.main as main_module
asyncio.run(
main_module.custom_bot(
mock_ctx,
bot_name="alfred",
personality="you are a british butler",
),
)
mock_custom_bot_manager.create_custom_bot.assert_called_once()
assert mock_ctx.send.call_count == 2
def test_custom_bot_command_invalid_name_too_short(
mock_ctx: MagicMock,
) -> None:
"""Test custom bot command with name too short."""
import asyncio
import vibe_bot.main as main_module
asyncio.run(
main_module.custom_bot(
mock_ctx,
bot_name="a",
personality="this is a valid personality description",
),
)
call_args = mock_ctx.send.call_args[0][0]
assert "Invalid bot name" in call_args
def test_custom_bot_command_invalid_name_empty(
mock_ctx: MagicMock,
) -> None:
"""Test custom bot command with empty name."""
import asyncio
import vibe_bot.main as main_module
asyncio.run(
main_module.custom_bot(
mock_ctx,
bot_name="",
personality="this is a valid personality description",
),
)
call_args = mock_ctx.send.call_args[0][0]
assert "Invalid bot name" in call_args
def test_custom_bot_command_invalid_personality(
mock_ctx: MagicMock,
) -> None:
"""Test custom bot command with personality too short."""
import asyncio
import vibe_bot.main as main_module
asyncio.run(
main_module.custom_bot(mock_ctx, bot_name="testbot", personality="short"),
)
call_args = mock_ctx.send.call_args[0][0]
assert "Invalid personality" in call_args
def test_custom_bot_command_create_fails(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test custom bot command when creation fails."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.create_custom_bot.return_value = False
asyncio.run(
main_module.custom_bot(
mock_ctx,
bot_name="alfred",
personality="you are a british butler",
),
)
call_args = mock_ctx.send.call_args[0][0]
assert "Failed to create custom bot" in call_args
def test_list_custom_bots_empty(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test listing custom bots when none exist."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.list_custom_bots.return_value = []
asyncio.run(main_module.list_custom_bots(mock_ctx))
call_args = mock_ctx.send.call_args[0][0]
assert "No custom bots" in call_args
def test_list_custom_bots_with_bots(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test listing custom bots when bots exist."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.list_custom_bots.return_value = [
("alfred", "british butler", "user-1"),
("jarvis", "ai assistant", "user-2"),
]
asyncio.run(main_module.list_custom_bots(mock_ctx))
call_args = mock_ctx.send.call_args[0][0]
assert "Available Custom Bots" in call_args
assert "* alfred" in call_args
assert "* jarvis" in call_args
def test_delete_custom_bot_success(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test deleting a custom bot successfully."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"12345",
"2024-01-01",
)
mock_custom_bot_manager.delete_custom_bot.return_value = True
asyncio.run(main_module.delete_custom_bot(mock_ctx, bot_name="alfred"))
call_args = mock_ctx.send.call_args[0][0]
assert "has been deleted" in call_args
def test_delete_custom_bot_not_found(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test deleting a non-existent custom bot."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = None
asyncio.run(main_module.delete_custom_bot(mock_ctx, bot_name="nonexistent"))
call_args = mock_ctx.send.call_args[0][0]
assert "not found" in call_args
def test_delete_custom_bot_not_owner(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test deleting a custom bot you don't own."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"other-user-id",
"2024-01-01",
)
asyncio.run(main_module.delete_custom_bot(mock_ctx, bot_name="alfred"))
call_args = mock_ctx.send.call_args[0][0]
assert "You can only delete your own" in call_args
def test_delete_custom_bot_delete_fails(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test deleting a custom bot when delete fails."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"12345",
"2024-01-01",
)
mock_custom_bot_manager.delete_custom_bot.return_value = False
asyncio.run(main_module.delete_custom_bot(mock_ctx, bot_name="alfred"))
call_args = mock_ctx.send.call_args[0][0]
assert "Failed to delete" in call_args
def test_on_message_skips_bot_messages(mock_ctx: MagicMock) -> None:
"""Test that on_message skips messages from the bot itself."""
import asyncio
import vibe_bot.main as main_module
message = MagicMock()
message.author = main_module.bot.user
message.content = "hello"
asyncio.run(main_module.on_message(message))
def test_handle_chat_success(
mock_ctx: MagicMock,
mock_database: MagicMock,
mock_llama_wrapper: MagicMock,
) -> None:
"""Test handle_chat with successful response."""
import asyncio
import vibe_bot.main as main_module
mock_llama_wrapper.chat_completion_with_tools.return_value = (
"This is a bot response"
)
asyncio.run(
main_module.handle_chat(
ctx=mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
),
)
mock_llama_wrapper.chat_completion_with_tools.assert_called_once()
mock_database.add_message.assert_called()
assert mock_ctx.send.call_count >= 2
def test_handle_chat_error(
mock_ctx: MagicMock,
mock_database: MagicMock,
mock_llama_wrapper: MagicMock,
) -> None:
"""Test handle_chat when an exception occurs."""
import asyncio
import vibe_bot.main as main_module
mock_llama_wrapper.chat_completion_with_tools.side_effect = Exception("API error")
asyncio.run(
main_module.handle_chat(
ctx=mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
),
)
call_args = mock_ctx.send.call_args[0][0]
assert "error occurred" in call_args.lower()
def test_handle_chat_long_response_chunked(
mock_ctx: MagicMock,
mock_database: MagicMock,
mock_llama_wrapper: MagicMock,
) -> None:
"""Test that long bot responses are sent in chunks."""
import asyncio
import vibe_bot.main as main_module
long_response = "x" * 2500
mock_llama_wrapper.chat_completion_with_tools.return_value = long_response
asyncio.run(
main_module.handle_chat(
ctx=mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
),
)
assert mock_ctx.send.call_count >= 3
def test_speak_plain_with_mock_tts(
mock_ctx: MagicMock,
mock_tts_engine: MagicMock,
) -> None:
"""Test _speak_plain function directly."""
import asyncio
import vibe_bot.main as main_module
asyncio.run(main_module._speak_plain(mock_ctx, "hello world", mock_tts_engine))
from vibe_bot.config import TTS_SPEED, TTS_VOICE
from vibe_bot.tts import DEFAULT_LANG
mock_tts_engine.generate_audio.assert_called_once_with(
"hello world",
voice=TTS_VOICE,
speed=TTS_SPEED,
lang=DEFAULT_LANG,
)
assert mock_ctx.send.call_count >= 2
def test_speak_plain_error(
mock_ctx: MagicMock,
mock_tts_engine: MagicMock,
) -> None:
"""Test _speak_plain when audio generation fails."""
import asyncio
import vibe_bot.main as main_module
mock_tts_engine.generate_audio.side_effect = Exception("generation error")
asyncio.run(main_module._speak_plain(mock_ctx, "hello world", mock_tts_engine))
call_args = mock_ctx.send.call_args[0][0]
assert "error generating speech" in call_args.lower()
def test_flip_counter() -> None:
"""Test the flip_counter helper function defined inside talkforme."""
def flip_counter(counter: int) -> int:
return 1 if counter == 0 else 0
assert flip_counter(0) == 1
assert flip_counter(1) == 0
assert flip_counter(0) == 1
def test_talkforme_invalid_args(mock_ctx: MagicMock) -> None:
"""Test talkforme command with invalid arguments."""
import asyncio
import vibe_bot.main as main_module
asyncio.run(main_module.talkforme(mock_ctx, message="bot1 bot2"))
call_args = mock_ctx.send.call_args[0][0]
assert "Usage" in call_args
def test_talkforme_bot1_not_found(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test talkforme when bot1 doesn't exist."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = None
asyncio.run(main_module.talkforme(mock_ctx, message="bot1 bot2 4 a topic"))
call_args = mock_ctx.send.call_args[0][0]
assert "is not a real bot" in call_args
def test_talkforme_bot2_not_found(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test talkforme when bot2 doesn't exist."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.side_effect = [
("bot1", "bot1 personality", "user-1", "2024-01-01"),
None,
]
asyncio.run(main_module.talkforme(mock_ctx, message="bot1 bot2 4 a topic"))
call_args = mock_ctx.send.call_args[0][0]
assert "is not a real bot" in call_args
def test_talkforme_invalid_limit(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test talkforme with non-integer limit."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = (
"bot1",
"personality",
"user-1",
"2024-01-01",
)
asyncio.run(main_module.talkforme(mock_ctx, message="bot1 bot2 abc topic"))
call_args = mock_ctx.send.call_args[0][0]
assert "must be an integer" in call_args
def test_history_bot_not_found(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
) -> None:
"""Test history command when bot doesn't exist."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = None
asyncio.run(main_module.history(mock_ctx, bot_name="nonexistent"))
call_args = mock_ctx.send.call_args[0][0]
assert "not found" in call_args
def test_history_no_history(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
mock_database: MagicMock,
) -> None:
"""Test history command when bot has no chat history."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
mock_database.get_bot_history.return_value = []
asyncio.run(main_module.history(mock_ctx, bot_name="alfred"))
call_args = mock_ctx.send.call_args[0][0]
assert "No chat history" in call_args
assert "**alfred**" in call_args
def test_history_with_data(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
mock_database: MagicMock,
) -> None:
"""Test history command when bot has chat history."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
mock_database.get_bot_history.return_value = [
("hello", "yes master?"),
("what time is it", "it is currently 3pm"),
]
asyncio.run(main_module.history(mock_ctx, bot_name="alfred"))
assert mock_ctx.send.call_count >= 1
def test_get_user_info_minimal(mock_ctx: MagicMock) -> None:
"""Test get_user_info with minimal member data."""
import vibe_bot.main as main_module
result = main_module.get_user_info(mock_ctx.author)
assert "Username: testuser" in result
assert "User ID: 12345" in result
assert "Global Name: Test User" in result
assert "Nickname: tester" in result
def test_get_user_info_with_member_data(mock_ctx_with_member: MagicMock) -> None:
"""Test get_user_info with full member data including roles and activities."""
import vibe_bot.main as main_module
result = main_module.get_user_info(mock_ctx_with_member.author)
assert "Global Name: Test User" in result
assert "Nickname: tester" in result
assert "Username: testuser" in result
assert "User ID: 12345" in result
assert "Top Role: Admin" in result
assert "Activities: Chess" in result
assert "Joined: 2024-01-15" in result
assert "Account Created: 2023-06-01" in result
def test_get_user_info_no_global_name(mock_ctx: MagicMock) -> None:
"""Test get_user_info when user has no global name."""
import vibe_bot.main as main_module
mock_ctx.author.global_name = None
mock_ctx.author.nick = None
mock_ctx.author.top_role.name = "@everyone"
mock_ctx.author.activities = []
result = main_module.get_user_info(mock_ctx.author)
assert "Global Name:" not in result
assert "Nickname:" not in result
assert "Top Role:" not in result
assert "Activities:" not in result
assert "Username: testuser" in result
assert "User ID: 12345" in result
def test_get_user_info_with_top_role_not_everyone(
mock_ctx_with_member: MagicMock,
) -> None:
"""Test get_user_info includes top role when not @everyone."""
import vibe_bot.main as main_module
result = main_module.get_user_info(mock_ctx_with_member.author)
assert "Top Role: Admin" in result
def test_get_user_info_no_activities(mock_ctx: MagicMock) -> None:
"""Test get_user_info when user has no activities."""
import vibe_bot.main as main_module
mock_ctx.author.activities = []
result = main_module.get_user_info(mock_ctx.author)
assert "Activities:" not in result
def test_handle_chat_includes_user_info(
mock_ctx: MagicMock,
mock_database: MagicMock,
mock_llama_wrapper: MagicMock,
) -> None:
"""Test handle_chat includes user info in system prompt."""
import asyncio
import vibe_bot.main as main_module
mock_llama_wrapper.chat_completion_with_tools.return_value = (
"This is a bot response"
)
asyncio.run(
main_module.handle_chat(
ctx=mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
),
)
mock_llama_wrapper.chat_completion_with_tools.assert_called_once()
call_kwargs = mock_llama_wrapper.chat_completion_with_tools.call_args
system_prompt = call_kwargs.kwargs["system_prompt"]
assert "you are a butler" in system_prompt
assert "User Information:" in system_prompt
assert "Username: testuser" in system_prompt
assert "User ID: 12345" in system_prompt
def test_speak_with_bot_includes_user_info(
mock_ctx: MagicMock,
mock_tts_engine: MagicMock,
mock_custom_bot_manager: MagicMock,
mock_database: MagicMock,
mock_llama_wrapper: MagicMock,
) -> None:
"""Test _speak_with_bot includes user info in system prompt."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.list_custom_bots.return_value = [
("alfred", "british butler", "user-123"),
]
mock_custom_bot_manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
asyncio.run(main_module.speak(mock_ctx, message="alfred what time is it"))
mock_llama_wrapper.chat_completion_with_tools.assert_called_once()
call_kwargs = mock_llama_wrapper.chat_completion_with_tools.call_args
system_prompt = call_kwargs.kwargs["system_prompt"]
assert "british butler" in system_prompt
assert "User Information:" in system_prompt
assert "Username: testuser" in system_prompt
assert "User ID: 12345" in system_prompt
mock_tts_engine.generate_audio.assert_called_once()
def test_history_long_response_chunked(
mock_ctx: MagicMock,
mock_custom_bot_manager: MagicMock,
mock_database: MagicMock,
) -> None:
"""Test that long history responses are sent in chunks."""
import asyncio
import vibe_bot.main as main_module
mock_custom_bot_manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
long_user = "x" * 500
long_bot = "y" * 500
mock_database.get_bot_history.return_value = [
(long_user, long_bot),
]
asyncio.run(main_module.history(mock_ctx, bot_name="alfred"))
assert mock_ctx.send.call_count >= 1