fix linting, formatting, and add tests
This commit is contained in:
@@ -0,0 +1,530 @@
|
||||
"""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" # noqa: E501
|
||||
|
||||
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))
|
||||
|
||||
mock_tts_engine.generate_audio.assert_called_once_with(
|
||||
"hello world",
|
||||
voice=main_module.TTS_VOICE, # type: ignore[attr-defined]
|
||||
speed=main_module.TTS_SPEED, # type: ignore[attr-defined]
|
||||
)
|
||||
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
|
||||
Reference in New Issue
Block a user