638 lines
17 KiB
Python
638 lines
17 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.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_history.assert_called_once()
|
|
mock_tts_engine.generate_audio.assert_called_once()
|
|
|
|
|
|
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_history.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_history.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_history.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_history.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
|
|
|
|
mock_tts_engine.generate_audio.assert_called_once_with(
|
|
"hello world",
|
|
voice=TTS_VOICE,
|
|
speed=TTS_SPEED,
|
|
)
|
|
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
|
|
first_call = mock_ctx.send.call_args_list[0][0][0]
|
|
assert "Chat History for **alfred**" in first_call
|
|
assert "hello" in first_call
|
|
assert "alfred: yes master?" in first_call
|
|
assert "what time is it" in first_call
|
|
assert "alfred: it is currently 3pm" in first_call
|
|
|
|
|
|
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
|