1261 lines
37 KiB
Python
1261 lines
37 KiB
Python
"""Tests for the main module (Discord bot commands)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC
|
|
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, tzinfo=UTC)
|
|
ctx.author.created_at = datetime(2023, 6, 1, tzinfo=UTC)
|
|
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 required intents are 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
|
|
|
|
|
|
def test_bot_intents_members_and_presence(mock_discord: dict[str, MagicMock]) -> None:
|
|
"""Test that members and presence intents are enabled."""
|
|
intents = mock_discord["Intents"].default.return_value
|
|
intents.message_content = True
|
|
intents.members = True
|
|
intents.presences = True
|
|
assert intents.message_content is True
|
|
assert intents.members is True
|
|
assert intents.presences is True
|
|
|
|
|
|
@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
|
|
|
|
|
|
def test_debug_no_subcommand(mock_ctx: MagicMock) -> None:
|
|
"""Test debug command without subcommand shows menu."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
asyncio.run(main_module.debug(mock_ctx, subcommand=None))
|
|
|
|
mock_ctx.send.assert_called_once()
|
|
call_args = mock_ctx.send.call_args[0][0]
|
|
assert "Debug Menu" in call_args
|
|
assert "members" in call_args
|
|
|
|
|
|
def test_debug_members_no_guild(mock_ctx: MagicMock) -> None:
|
|
"""Test debug members command when channel has no guild."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_ctx.channel.guild = None
|
|
|
|
asyncio.run(main_module.debug(mock_ctx, subcommand="members"))
|
|
|
|
mock_ctx.send.assert_called_once()
|
|
call_args = mock_ctx.send.call_args[0][0]
|
|
assert "No members found in this channel." in call_args
|
|
|
|
|
|
def test_debug_members_with_members(mock_ctx: MagicMock) -> None:
|
|
"""Test debug members command with channel members."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_member = MagicMock()
|
|
mock_member.display_name = "Alice"
|
|
mock_member.name = "alice"
|
|
mock_member.nick = None
|
|
mock_member.global_name = None
|
|
mock_member.status = MagicMock(value="online")
|
|
|
|
mock_ctx.channel.guild.members = [mock_member]
|
|
|
|
asyncio.run(main_module.debug(mock_ctx, subcommand="members"))
|
|
|
|
assert mock_ctx.send.called
|
|
call_args = mock_ctx.send.call_args[0][0]
|
|
assert "Alice" in call_args
|
|
assert "1 total" in call_args
|
|
|
|
|
|
def test_debug_unknown_subcommand(mock_ctx: MagicMock) -> None:
|
|
"""Test debug command with unknown subcommand."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
asyncio.run(main_module.debug(mock_ctx, subcommand="unknown"))
|
|
|
|
mock_ctx.send.assert_called_once()
|
|
call_args = mock_ctx.send.call_args[0][0]
|
|
assert "Unknown debug sub-command" in call_args
|
|
|
|
|
|
def test_debug_members_many_chunks(mock_ctx: MagicMock) -> None:
|
|
"""Test debug members command with many members that exceed 1900 chars."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_members = []
|
|
for i in range(50):
|
|
mock_member = MagicMock()
|
|
mock_member.display_name = f"User{i}_with_a_very_long_display_name"
|
|
mock_member.name = f"user{i}"
|
|
mock_member.nick = None
|
|
mock_member.global_name = None
|
|
mock_member.status = MagicMock(value="online")
|
|
mock_members.append(mock_member)
|
|
|
|
mock_ctx.channel.guild.members = mock_members
|
|
|
|
asyncio.run(main_module.debug(mock_ctx, subcommand="members"))
|
|
|
|
assert mock_ctx.send.call_count >= 2
|
|
first_chunk = mock_ctx.send.call_args_list[0][0][0]
|
|
assert "Members in this channel (50 total)" in first_chunk
|
|
|
|
|
|
def test_debug_whoami(mock_ctx: MagicMock) -> None:
|
|
"""Test debug whoami command shows user info."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
asyncio.run(main_module.debug(mock_ctx, subcommand="whoami"))
|
|
|
|
mock_ctx.send.assert_called_once()
|
|
call_args = mock_ctx.send.call_args[0][0]
|
|
assert "Username: testuser" in call_args
|
|
assert "User ID: 12345" in call_args
|
|
assert "Global Name: Test User" in call_args
|
|
assert "Nickname: tester" in call_args
|
|
|
|
|
|
def test_debug_whoami_minimal(mock_ctx: MagicMock) -> None:
|
|
"""Test debug whoami command with minimal user data."""
|
|
import asyncio
|
|
|
|
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 = []
|
|
mock_ctx.author.joined_at = None
|
|
mock_ctx.author.created_at = None
|
|
|
|
asyncio.run(main_module.debug(mock_ctx, subcommand="whoami"))
|
|
|
|
mock_ctx.send.assert_called_once()
|
|
call_args = mock_ctx.send.call_args[0][0]
|
|
assert "Username: testuser" in call_args
|
|
assert "User ID: 12345" in call_args
|
|
assert "Global Name" not in call_args
|
|
assert "Nickname" not in call_args
|
|
assert "Activities" not in call_args
|
|
assert "Joined" not in call_args
|
|
|
|
|
|
def test_debug_tools(mock_ctx: MagicMock) -> None:
|
|
"""Test debug tools command shows LLM tools."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
asyncio.run(main_module.debug(mock_ctx, subcommand="tools"))
|
|
|
|
mock_ctx.send.assert_called_once()
|
|
call_args = mock_ctx.send.call_args[0][0]
|
|
assert "LLM Tools" in call_args
|
|
assert "get_channel_members" in call_args
|
|
assert "members" in call_args.lower()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Image layout selection (doodlebob)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
LAYOUT_TEST_SIZES = {
|
|
"portrait": "1024x1536",
|
|
"landscape": "1536x1024",
|
|
"square": "1024x1024",
|
|
}
|
|
|
|
|
|
def _sent_texts(mock_ctx: MagicMock) -> list[str]:
|
|
"""Collect positional (text) arguments sent via ctx.send."""
|
|
return [c.args[0] for c in mock_ctx.send.call_args_list if c.args]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("response", "expected"),
|
|
[
|
|
("portrait", "portrait"),
|
|
("landscape", "landscape"),
|
|
("square", "square"),
|
|
(" Portrait ", "portrait"),
|
|
("LANDSCAPE", "landscape"),
|
|
("square.", "square"),
|
|
("I would use portrait.", "portrait"),
|
|
("This scene is best as landscape.", "landscape"),
|
|
("A square composition works here.", "square"),
|
|
],
|
|
)
|
|
def test_parse_image_layout_valid(response: str, expected: str) -> None:
|
|
"""Test that valid LLM layout responses parse to the right layout."""
|
|
import vibe_bot.main as main_module
|
|
|
|
assert main_module.parse_image_layout(response) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"response",
|
|
[
|
|
"",
|
|
" ",
|
|
"banana",
|
|
"1024x1024",
|
|
"tall and wide",
|
|
"squareness", # word boundary should prevent a match on "square"
|
|
],
|
|
)
|
|
def test_parse_image_layout_defaults_to_square(response: str) -> None:
|
|
"""Test that empty or malformed responses fall back to square."""
|
|
import vibe_bot.main as main_module
|
|
|
|
assert main_module.parse_image_layout(response) == "square"
|
|
|
|
|
|
def test_image_prompt_system_prompt_covers_key_details() -> None:
|
|
"""The prompt-rewrite system prompt forces explicit detail on all aspects."""
|
|
import vibe_bot.main as main_module
|
|
|
|
prompt = main_module.IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE.format(layout="square")
|
|
lowered = prompt.lower()
|
|
assert "square" in prompt
|
|
assert "exact text" in lowered
|
|
assert "composition" in lowered
|
|
assert "style" in lowered
|
|
assert "canada goose" in lowered
|
|
assert "only the image generation prompt" in lowered
|
|
|
|
|
|
def test_doodlebob_prompt_rewrite_uses_detailed_system_prompt(
|
|
mock_ctx: MagicMock,
|
|
mock_llama_wrapper: MagicMock,
|
|
mock_base64: MagicMock,
|
|
mock_database: MagicMock,
|
|
) -> None:
|
|
"""The prompt rewrite call uses the detailed system prompt with the layout."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_llama_wrapper.chat_completion_instruct.side_effect = [
|
|
"landscape", # layout selection
|
|
"a very detailed prompt", # prompt rewrite
|
|
]
|
|
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
|
|
mock_database.get_image_generation_time_estimate.return_value = None
|
|
|
|
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
|
|
asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
|
|
|
|
rewrite_call = mock_llama_wrapper.chat_completion_instruct.call_args_list[1]
|
|
expected = main_module.IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE.format(
|
|
layout="landscape"
|
|
)
|
|
assert rewrite_call.kwargs["system_prompt"] == expected
|
|
assert rewrite_call.kwargs["user_prompt"] == "a scene"
|
|
assert mock_llama_wrapper.image_generation.call_args.kwargs["prompt"] == (
|
|
"a very detailed prompt"
|
|
)
|
|
|
|
|
|
def test_doodlebob_selects_portrait(
|
|
mock_ctx: MagicMock,
|
|
mock_llama_wrapper: MagicMock,
|
|
mock_base64: MagicMock,
|
|
mock_database: MagicMock,
|
|
) -> None:
|
|
"""Test doodlebob picks portrait and passes the portrait size."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_llama_wrapper.chat_completion_instruct.side_effect = [
|
|
"portrait", # layout selection
|
|
"a tall portrait of a lighthouse", # prompt rewrite
|
|
]
|
|
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
|
|
mock_database.get_image_generation_time_estimate.return_value = None
|
|
|
|
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
|
|
asyncio.run(main_module.doodlebob(mock_ctx, message="a tall lighthouse"))
|
|
|
|
# Layout selection is the first instruct call, using the layout prompt.
|
|
layout_call = mock_llama_wrapper.chat_completion_instruct.call_args_list[0]
|
|
assert layout_call.kwargs["system_prompt"] == main_module.IMAGE_LAYOUT_SYSTEM_PROMPT
|
|
assert layout_call.kwargs["user_prompt"] == "a tall lighthouse"
|
|
|
|
mock_llama_wrapper.image_generation.assert_called_once()
|
|
assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1536"
|
|
|
|
sent = _sent_texts(mock_ctx)
|
|
assert any("shopping for a canvas" in m for m in sent)
|
|
assert any("selected portrait" in m for m in sent)
|
|
assert any("drone strike" in m for m in sent)
|
|
|
|
|
|
def test_doodlebob_selects_landscape(
|
|
mock_ctx: MagicMock,
|
|
mock_llama_wrapper: MagicMock,
|
|
mock_base64: MagicMock,
|
|
mock_database: MagicMock,
|
|
) -> None:
|
|
"""Test doodlebob picks landscape and passes the landscape size."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_llama_wrapper.chat_completion_instruct.side_effect = [
|
|
"landscape",
|
|
"a wide panoramic coastline",
|
|
]
|
|
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
|
|
mock_database.get_image_generation_time_estimate.return_value = None
|
|
|
|
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
|
|
asyncio.run(main_module.doodlebob(mock_ctx, message="wide coastline"))
|
|
|
|
assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1536x1024"
|
|
|
|
|
|
def test_doodlebob_malformed_layout_defaults_square(
|
|
mock_ctx: MagicMock,
|
|
mock_llama_wrapper: MagicMock,
|
|
mock_base64: MagicMock,
|
|
mock_database: MagicMock,
|
|
) -> None:
|
|
"""Test a malformed layout response falls back to the square size."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_llama_wrapper.chat_completion_instruct.side_effect = [
|
|
"banana", # malformed layout
|
|
"a balanced composition",
|
|
]
|
|
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
|
|
mock_database.get_image_generation_time_estimate.return_value = None
|
|
|
|
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
|
|
asyncio.run(main_module.doodlebob(mock_ctx, message="a logo"))
|
|
|
|
assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1024"
|
|
sent = _sent_texts(mock_ctx)
|
|
assert any("selected square" in m for m in sent)
|
|
|
|
|
|
def test_doodlebob_empty_layout_defaults_square(
|
|
mock_ctx: MagicMock,
|
|
mock_llama_wrapper: MagicMock,
|
|
mock_base64: MagicMock,
|
|
mock_database: MagicMock,
|
|
) -> None:
|
|
"""Test an empty layout response (LLM failure) falls back to square."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_llama_wrapper.chat_completion_instruct.side_effect = [
|
|
"", # empty layout response
|
|
"a balanced composition",
|
|
]
|
|
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
|
|
mock_database.get_image_generation_time_estimate.return_value = None
|
|
|
|
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
|
|
asyncio.run(main_module.doodlebob(mock_ctx, message="a logo"))
|
|
|
|
assert mock_llama_wrapper.image_generation.call_args.kwargs["size"] == "1024x1024"
|
|
|
|
|
|
def test_doodlebob_reports_estimate_and_elapsed(
|
|
mock_ctx: MagicMock,
|
|
mock_llama_wrapper: MagicMock,
|
|
mock_base64: MagicMock,
|
|
mock_database: MagicMock,
|
|
) -> None:
|
|
"""Doodlebob posts an ETA from history and the final elapsed time."""
|
|
import asyncio
|
|
import re
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_llama_wrapper.chat_completion_instruct.side_effect = [
|
|
"square",
|
|
"a test scene",
|
|
]
|
|
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
|
|
mock_database.get_image_generation_time_estimate.return_value = 12.34
|
|
|
|
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
|
|
asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
|
|
|
|
mock_database.record_image_generation_time.assert_called_once()
|
|
recorded = mock_database.record_image_generation_time.call_args[0][0]
|
|
assert isinstance(recorded, float)
|
|
assert recorded >= 0.0
|
|
|
|
sent = _sent_texts(mock_ctx)
|
|
assert any("drone strike" in m for m in sent)
|
|
assert any("~12 seconds" in m for m in sent)
|
|
assert any(re.search(r"generated in \d+\.\d+ seconds", m) for m in sent)
|
|
|
|
|
|
def test_doodlebob_no_estimate_without_history(
|
|
mock_ctx: MagicMock,
|
|
mock_llama_wrapper: MagicMock,
|
|
mock_base64: MagicMock,
|
|
mock_database: MagicMock,
|
|
) -> None:
|
|
"""No ETA is posted when there is no generation history yet."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_llama_wrapper.chat_completion_instruct.side_effect = [
|
|
"square",
|
|
"a test scene",
|
|
]
|
|
mock_llama_wrapper.image_generation.return_value = "aW1hZ2U="
|
|
mock_database.get_image_generation_time_estimate.return_value = None
|
|
|
|
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
|
|
asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
|
|
|
|
sent = _sent_texts(mock_ctx)
|
|
assert any("drone strike" in m for m in sent)
|
|
assert not any("ETA" in m for m in sent)
|
|
mock_database.record_image_generation_time.assert_called_once()
|
|
|
|
|
|
def test_doodlebob_failed_generation_not_recorded(
|
|
mock_ctx: MagicMock,
|
|
mock_llama_wrapper: MagicMock,
|
|
mock_base64: MagicMock,
|
|
mock_database: MagicMock,
|
|
) -> None:
|
|
"""A failed generation records no time and reports no elapsed seconds."""
|
|
import asyncio
|
|
|
|
import vibe_bot.main as main_module
|
|
|
|
mock_llama_wrapper.chat_completion_instruct.side_effect = [
|
|
"square",
|
|
"a test scene",
|
|
]
|
|
mock_llama_wrapper.image_generation.return_value = ""
|
|
mock_database.get_image_generation_time_estimate.return_value = 8.0
|
|
|
|
with patch.object(main_module, "LAYOUT_SIZES", LAYOUT_TEST_SIZES):
|
|
asyncio.run(main_module.doodlebob(mock_ctx, message="a scene"))
|
|
|
|
mock_database.record_image_generation_time.assert_not_called()
|
|
sent = _sent_texts(mock_ctx)
|
|
assert any("Failed to generate image" in m for m in sent)
|
|
assert not any("generated in" in m for m in sent)
|