services/vibe-bot
2
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity

complete restructure

Browse Source
This commit is contained in:
ducoterra
2026-08-19 13:16:43 -04:00
parent d7b6f28cbd
commit f87e1d51ef
60 changed files with 8176 additions and 5143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vibe_bot/tests/_helpers.py
+23
View File
@@ -0,0 +1,23 @@
"""Shared helpers for wiring tests (command invocation, sent-text asserts)."""
from __future__ import annotations
import asyncio
from collections.abc import Callable
from typing import Any, cast
from unittest.mock import MagicMock
from discord.ext import commands
def invoke(bot: commands.Bot, name: str, *args: Any, **kwargs: Any) -> None:
"""Invoke a registered command's callback directly with the given args."""
cmd = bot.get_command(name)
assert cmd is not None
callback = cast("Callable[..., Any]", cmd.callback)
asyncio.run(callback(*args, **kwargs))
def sent_texts(ctx: MagicMock) -> list[str]:
"""All positional text messages sent through ctx.send."""
return [c.args[0] for c in ctx.send.call_args_list if c.args]
vibe_bot/tests/conftest.py
+77 -148
View File
@@ -5,12 +5,16 @@ from __future__ import annotations
import tempfile
import warnings
from collections.abc import Generator
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
from unittest.mock import AsyncMock, MagicMock, patch
import numpy as np
import pytest
from discord.ext import commands
from vibe_bot.app import App, build_bot
warnings.filterwarnings(
"ignore",
@@ -21,41 +25,79 @@ if TYPE_CHECKING:
from vibe_bot.database import ChatDatabase, CustomBotManager
@dataclass(frozen=True)
class AppMocks:
"""The App under test plus its mock dependencies, for setup and asserts."""
app: App
db: MagicMock
manager: MagicMock
registry: MagicMock
chat: MagicMock
image: MagicMock
speech: MagicMock
conversation: MagicMock
@pytest.fixture
def mock_env_vars() -> Generator[None]:
"""Provide minimal env vars for config loading."""
with patch.dict(
"os.environ",
{
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
"CHAT_ENDPOINT_KEY": "test-key",
"COMPLETION_ENDPOINT_KEY": "test-completion-key",
"IMAGE_GEN_ENDPOINT_KEY": "test-image-key",
"IMAGE_EDIT_ENDPOINT_KEY": "test-image-edit-key",
"EMBEDDING_ENDPOINT_KEY": "test-embedding-key",
"MAX_COMPLETION_TOKENS": "1000",
"MAX_HISTORY_MESSAGES": "1000",
"SIMILARITY_THRESHOLD": "0.7",
"TOP_K_RESULTS": "5",
"TTS_MODEL_PATH": "/tmp/test-model.onnx",
"TTS_VOICES_PATH": "/tmp/test-voices.bin",
"TTS_VOICE": "af_sarah",
"TTS_SPEED": "1.0",
"DB_PATH": ":memory:",
},
clear=False,
):
yield
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 app_mocks() -> AppMocks:
"""An App built entirely from mocks, with alfred in the bot cache."""
db = MagicMock()
manager = MagicMock()
manager.list_custom_bots.return_value = [
("alfred", "british butler", "user123"),
]
registry = MagicMock()
chat = MagicMock()
chat.handle = AsyncMock()
image = MagicMock()
image.generate = AsyncMock()
image.edit = AsyncMock()
speech = MagicMock()
speech.speak = AsyncMock()
conversation = MagicMock()
conversation.run = AsyncMock()
app = App(
db=db,
manager=manager,
registry=registry,
tts=MagicMock(),
chat=chat,
image=image,
speech=speech,
conversation=conversation,
bot_cache={"alfred": ("british butler", "user123")},
)
return AppMocks(app, db, manager, registry, chat, image, speech, conversation)
@pytest.fixture
def bot(app_mocks: AppMocks) -> commands.Bot:
"""A real Bot built from the mock App (never connected)."""
return build_bot(app_mocks.app)
@pytest.fixture
@@ -71,22 +113,13 @@ def temp_db_path() -> Generator[str]:
def mock_embedding() -> Generator[MagicMock]:
"""Provide a mock embedding function returning a fixed vector."""
vector: list[float] = [0.1] * 2048
with patch("vibe_bot.llama_wrapper.embedding", return_value=vector) as mock:
yield mock
@pytest.fixture
def mock_openai_client() -> Generator[MagicMock]:
"""Provide a mock OpenAI client."""
mock_client = MagicMock()
with patch("vibe_bot.database.OpenAI", return_value=mock_client) as mock:
with patch("vibe_bot.llm_client.embedding", return_value=vector) as mock:
yield mock
@pytest.fixture
def chat_db(
temp_db_path: str,
mock_openai_client: MagicMock,
mock_embedding: MagicMock,
) -> Generator[ChatDatabase]:
"""Provide a ChatDatabase instance with a temp database."""
@@ -94,7 +127,6 @@ def chat_db(
db = ChatDatabase(db_path=temp_db_path)
yield db
db.client.close()
@pytest.fixture
@@ -133,106 +165,3 @@ def mock_kokoro_tts() -> Generator[dict[str, Any]]:
"mock_samples": mock_samples,
"mock_sr": 24000,
}
@pytest.fixture
def mock_discord() -> Generator[dict[str, MagicMock]]:
"""Mock discord module components."""
mock_intents = MagicMock()
mock_intents.default.return_value = MagicMock()
mock_intents.default.return_value.message_content = True
mock_bot_class = MagicMock()
mock_bot_instance = MagicMock()
mock_bot_instance.user = MagicMock()
mock_bot_instance.user.name = "test-bot"
mock_bot_instance.user.id = "123456789"
with (
patch("vibe_bot.main.discord") as mock_discord_module,
patch("vibe_bot.main.commands", MagicMock()),
patch("vibe_bot.main.commands.Bot", mock_bot_class),
):
mock_bot_class.return_value = mock_bot_instance
mock_discord_module.Intents = mock_intents
mock_discord_module.Message = MagicMock
mock_discord_module.File = MagicMock
yield {
"Intents": mock_intents,
"Bot": mock_bot_class,
"bot_instance": mock_bot_instance,
}
@pytest.fixture
def mock_tts_engine() -> Generator[MagicMock]:
"""Provide a mock TTSEngine."""
mock_engine = MagicMock()
mock_engine.generate_audio.return_value = MagicMock()
with (
patch("vibe_bot.main.tts_engine", mock_engine),
patch("vibe_bot.main.tts.TTSEngine", return_value=mock_engine),
):
yield mock_engine
@pytest.fixture
def mock_requests() -> Generator[MagicMock]:
"""Provide mock requests module."""
with patch("vibe_bot.main.requests") as mock_requests_module:
mock_response = MagicMock()
mock_response.content = b"fake image data"
mock_requests_module.get.return_value = mock_response
yield mock_requests_module
@pytest.fixture
def mock_base64() -> Generator[MagicMock]:
"""Provide mock base64 module."""
with patch("vibe_bot.main.base64") as mock_base64_module:
mock_base64_module.b64decode.return_value = b"fake image data"
yield mock_base64_module
@pytest.fixture
def mock_llama_wrapper() -> Generator[MagicMock]:
"""Provide mock llama_wrapper module."""
with patch("vibe_bot.main.llama_wrapper") as mock_wrapper:
mock_wrapper.chat_completion_with_history.return_value = "Bot response"
mock_wrapper.chat_completion_with_tools = AsyncMock(return_value="Bot response")
mock_wrapper.chat_completion_instruct.return_value = "image prompt"
mock_wrapper.image_generation.return_value = ""
mock_wrapper.image_edit.return_value = ""
mock_wrapper.embedding.return_value = [0.1] * 2048
yield mock_wrapper
@pytest.fixture
def mock_database() -> Generator[MagicMock]:
"""Provide mock database module."""
with patch("vibe_bot.main.get_database") as mock_get_db:
mock_db = MagicMock()
mock_db.get_conversation_context.return_value = []
mock_db.add_message.return_value = True
mock_get_db.return_value = mock_db
yield mock_db
@pytest.fixture
def mock_custom_bot_manager() -> Generator[MagicMock]:
"""Provide mock CustomBotManager."""
with patch("vibe_bot.main.CustomBotManager") as mock_manager_class:
mock_manager = MagicMock()
mock_manager.create_custom_bot.return_value = True
mock_manager.get_custom_bot.return_value = (
"alfred",
"british butler personality",
"user123",
"2024-01-01",
)
mock_manager.list_custom_bots.return_value = [
("alfred", "british butler personality", "user123"),
]
mock_manager.delete_custom_bot.return_value = True
mock_manager_class.return_value = mock_manager
yield mock_manager
vibe_bot/tests/test_app.py
+322
View File
@@ -0,0 +1,322 @@
"""Tests for the app composition root (singletons, services, bot handlers)."""
from __future__ import annotations
import asyncio
import logging
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from discord.ext import commands
from vibe_bot.app import (
configure_logging,
create_app,
invalidate_bot_cache,
)
from vibe_bot.commands import speech as speech_commands
from vibe_bot.config import TTS_MODEL_PATH, TTS_VOICES_PATH
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.conversation_service import ConversationService
from vibe_bot.services.image_service import ImageService
from vibe_bot.services.speech_service import SpeechService
from vibe_bot.tests._helpers import sent_texts
from vibe_bot.tests.conftest import AppMocks
def test_create_app_wires_singletons(
chat_db: Any,
custom_bot_manager: Any,
) -> None:
"""create_app shares one db/manager with the services and seeds the cache."""
engine = MagicMock()
with (
patch("vibe_bot.app.ChatDatabase", return_value=chat_db),
patch("vibe_bot.app.CustomBotManager", return_value=custom_bot_manager),
patch("vibe_bot.app.TTSEngine", return_value=engine) as mock_tts,
patch(
"vibe_bot.llm_client.get_tool_registry",
return_value=MagicMock(),
) as mock_registry,
):
app = create_app()
mock_tts.assert_called_once_with(TTS_MODEL_PATH, TTS_VOICES_PATH)
assert app.db is chat_db
assert app.manager is custom_bot_manager
assert app.tts is engine
assert app.registry is mock_registry.return_value
assert isinstance(app.chat, ChatService)
assert isinstance(app.image, ImageService)
assert isinstance(app.speech, SpeechService)
assert isinstance(app.conversation, ConversationService)
assert app.chat._db is chat_db
assert app.chat._registry is mock_registry.return_value
assert app.image._db is chat_db
assert app.speech._db is chat_db
assert app.speech._manager is custom_bot_manager
assert app.speech._tts is engine
assert app.conversation._manager is custom_bot_manager
assert app.bot_cache == {}
def test_create_app_tts_failure_tolerant(
chat_db: Any,
custom_bot_manager: Any,
) -> None:
"""A failing TTS engine degrades to None instead of crashing startup."""
with (
patch("vibe_bot.app.ChatDatabase", return_value=chat_db),
patch("vibe_bot.app.CustomBotManager", return_value=custom_bot_manager),
patch("vibe_bot.app.TTSEngine", side_effect=OSError("no model file")),
):
app = create_app()
assert app.tts is None
def test_invalidate_bot_cache_rebuilds(app_mocks: AppMocks) -> None:
"""invalidate_bot_cache rebuilds the cache from the manager."""
app_mocks.manager.list_custom_bots.return_value = [
("newbot", "a personality", "user999"),
("alfred", "british butler", "user123"),
]
invalidate_bot_cache(app_mocks.app)
assert app_mocks.app.bot_cache == {
"newbot": ("a personality", "user999"),
"alfred": ("british butler", "user123"),
}
def test_configure_logging_configures_root() -> None:
"""configure_logging is the sole basicConfig: it adds a root handler."""
root = logging.getLogger()
original = root.handlers
root.handlers.clear()
try:
configure_logging()
assert len(root.handlers) == 1
handler = root.handlers[0]
assert handler.formatter is not None
assert handler.formatter._fmt == (
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
finally:
root.handlers = original
# ---------------------------------------------------------------------------
# build_bot
# ---------------------------------------------------------------------------
def test_build_bot_registers_commands(bot: commands.Bot) -> None:
"""build_bot returns a commands.Bot with every command registered."""
assert isinstance(bot, commands.Bot)
for name in (
"custom-bot",
"list-custom-bots",
"delete-custom-bot",
"lobotomize",
"debug",
"voices",
"speak",
"doodlebob",
"retcon",
"history",
"talkforme",
):
assert bot.get_command(name) is not None
for event in ("on_ready", "on_message", "on_command_error"):
assert event in bot.__dict__
def test_build_bot_intents(bot: commands.Bot) -> None:
"""message_content, members, and presences intents are enabled."""
assert bot.intents.message_content is True
assert bot.intents.members is True
assert bot.intents.presences is True
def test_handlers_refuse_to_run_before_build_bot() -> None:
"""Event and command handlers raise if invoked before the wiring exists."""
import vibe_bot.app as app_module
from vibe_bot.commands import _state as commands_state
saved_app, saved_bot = app_module._app, app_module._bot
saved_commands_app = commands_state._app
app_module._app = None
app_module._bot = None
commands_state._app = None
try:
with pytest.raises(RuntimeError, match="App is not initialized"):
asyncio.run(app_module.on_message(MagicMock()))
with pytest.raises(RuntimeError, match="Bot is not initialized"):
asyncio.run(app_module.on_ready())
with pytest.raises(RuntimeError, match="App is not initialized"):
asyncio.run(speech_commands.speak(MagicMock(), message="hello"))
finally:
app_module._app = saved_app
app_module._bot = saved_bot
commands_state._app = saved_commands_app
# ---------------------------------------------------------------------------
# on_message guards
# ---------------------------------------------------------------------------
def test_on_message_ignores_non_bang_messages(
app_mocks: AppMocks,
bot: commands.Bot,
) -> None:
"""A non-! message returns before any service or DB touch."""
message = MagicMock()
message.author = MagicMock()
message.content = "hello there"
with patch.object(bot, "process_commands", new=AsyncMock()) as mock_process:
asyncio.run(bot.on_message(message))
mock_process.assert_not_called()
app_mocks.chat.handle.assert_not_awaited()
app_mocks.speech.speak.assert_not_awaited()
app_mocks.image.generate.assert_not_awaited()
app_mocks.image.edit.assert_not_awaited()
app_mocks.conversation.run.assert_not_awaited()
assert app_mocks.manager.mock_calls == []
assert app_mocks.db.mock_calls == []
def test_on_message_skips_bot_authors(app_mocks: AppMocks, bot: commands.Bot) -> None:
"""A bot-authored message returns without calling any service or DB."""
message = MagicMock()
message.author = bot.user
message.content = "!alfred hi"
with patch.object(bot, "process_commands", new=AsyncMock()) as mock_process:
asyncio.run(bot.on_message(message))
mock_process.assert_not_called()
app_mocks.chat.handle.assert_not_awaited()
assert app_mocks.manager.mock_calls == []
assert app_mocks.db.mock_calls == []
def test_on_message_routes_custom_bot(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!alfred hi dispatches to the chat service with the parsed args."""
message = MagicMock()
message.author = MagicMock()
message.author.name = "testuser"
message.content = "!alfred hi"
with (
patch.object(bot, "get_context", new=AsyncMock(return_value=mock_ctx)),
patch.object(bot, "process_commands", new=AsyncMock()) as mock_process,
):
asyncio.run(bot.on_message(message))
app_mocks.chat.handle.assert_awaited_once_with(
ctx=mock_ctx,
bot_name="alfred",
message="hi",
system_prompt="british butler",
response_prefix="alfred response",
)
mock_process.assert_not_called()
def test_on_message_falls_through_to_commands(
app_mocks: AppMocks,
bot: commands.Bot,
) -> None:
"""An unmatched ! message falls through to process_commands."""
message = MagicMock()
message.author = MagicMock()
message.author.name = "testuser"
message.content = "!unknown hi"
with patch.object(bot, "process_commands", new=AsyncMock()) as mock_process:
asyncio.run(bot.on_message(message))
mock_process.assert_awaited_once_with(message)
app_mocks.chat.handle.assert_not_awaited()
def test_on_message_uses_cache_not_manager(
app_mocks: AppMocks,
bot: commands.Bot,
) -> None:
"""Bot-name matching reads bot_cache, never list_custom_bots()."""
message = MagicMock()
message.author = MagicMock()
message.author.name = "testuser"
message.content = "!alfred hi"
with patch.object(bot, "get_context", new=AsyncMock(return_value=MagicMock())):
asyncio.run(bot.on_message(message))
app_mocks.manager.list_custom_bots.assert_not_called()
# ---------------------------------------------------------------------------
# on_ready / on_command_error
# ---------------------------------------------------------------------------
def test_on_ready(bot: commands.Bot) -> None:
"""on_ready logs startup without raising."""
handler = bot.__dict__["on_ready"]
asyncio.run(handler())
def test_on_command_error_cooldown(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A cooldown error becomes a friendly message."""
from discord.ext.commands.cooldowns import Cooldown
error = commands.CommandOnCooldown(
Cooldown(1.0, 60.0), 0.5, commands.BucketType.user
)
asyncio.run(bot.on_command_error(mock_ctx, error))
texts = sent_texts(mock_ctx)
assert any("too quickly" in t for t in texts)
def test_on_command_error_generic(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Any other command error is logged and not re-raised."""
error = commands.CommandError("boom")
asyncio.run(bot.on_command_error(mock_ctx, error))
# ---------------------------------------------------------------------------
# main.py entrypoint
# ---------------------------------------------------------------------------
def test_main_builds_and_runs_bot() -> None:
"""main() validates config, configures logging, and runs the built bot."""
from vibe_bot import main as main_module
with (
patch.object(main_module, "validate_config") as mock_validate,
patch.object(main_module, "configure_logging") as mock_logging,
patch.object(
main_module, "create_app", return_value=MagicMock()
) as mock_create,
patch.object(main_module, "build_bot", return_value=MagicMock()) as mock_build,
patch.object(main_module, "DISCORD_TOKEN", "test-token"),
):
main_module.main()
mock_validate.assert_called_once_with()
mock_logging.assert_called_once_with()
mock_create.assert_called_once_with()
mock_build.assert_called_once()
mock_build.return_value.run.assert_called_once_with("test-token")
vibe_bot/tests/test_commands.py
+643
View File
@@ -0,0 +1,643 @@
"""Wiring tests for the commands package: every handler reaches its service."""
from __future__ import annotations
import asyncio
from datetime import UTC, datetime
from unittest.mock import MagicMock
import pytest
from discord.ext import commands
from vibe_bot.commands.custom_bots import MAX_PERSONALITY_LENGTH
from vibe_bot.tests._helpers import invoke, sent_texts
from vibe_bot.tests.conftest import AppMocks
# ---------------------------------------------------------------------------
# custom_bots: create / list / delete
# ---------------------------------------------------------------------------
def test_custom_bot_success_updates_cache(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A successful create invalidates the cache with the new bot."""
app_mocks.manager.create_custom_bot.return_value = "created"
app_mocks.manager.list_custom_bots.return_value = [
("newbot", "a personality", "user123"),
]
invoke(
bot, "custom-bot", mock_ctx, "newbot", personality="you are a british butler"
)
app_mocks.manager.create_custom_bot.assert_called_once_with(
bot_name="newbot",
system_prompt="you are a british butler",
created_by="12345",
)
assert app_mocks.app.bot_cache == {"newbot": ("a personality", "user123")}
assert "alfred" not in app_mocks.app.bot_cache
texts = sent_texts(mock_ctx)
assert any("has been created" in t for t in texts)
assert any("You can now use this bot" in t for t in texts)
def test_custom_bot_replaced_message(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""When the name already exists, the command reports a replace, not a create."""
app_mocks.manager.create_custom_bot.return_value = "replaced"
invoke(
bot, "custom-bot", mock_ctx, "alfred", personality="you are a british butler"
)
texts = sent_texts(mock_ctx)
assert any("already existed" in t and "replaced" in t for t in texts)
def test_custom_bot_invalid_name_too_short(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A one-character name is rejected before any manager call."""
invoke(bot, "custom-bot", mock_ctx, "a", personality="this is a valid personality")
texts = sent_texts(mock_ctx)
assert any("Invalid bot name" in t for t in texts)
app_mocks.manager.create_custom_bot.assert_not_called()
def test_custom_bot_invalid_name_empty(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""An empty name is rejected before any manager call."""
invoke(bot, "custom-bot", mock_ctx, "", personality="this is a valid personality")
texts = sent_texts(mock_ctx)
assert any("Invalid bot name" in t for t in texts)
app_mocks.manager.create_custom_bot.assert_not_called()
def test_custom_bot_invalid_personality(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A personality under 10 characters is rejected."""
invoke(bot, "custom-bot", mock_ctx, "testbot", personality="short")
texts = sent_texts(mock_ctx)
assert any("Invalid personality" in t for t in texts)
app_mocks.manager.create_custom_bot.assert_not_called()
def test_custom_bot_personality_too_long(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A personality over MAX_PERSONALITY_LENGTH is rejected before any DB call."""
invoke(
bot,
"custom-bot",
mock_ctx,
"testbot",
personality="a" * (MAX_PERSONALITY_LENGTH + 1),
)
texts = sent_texts(mock_ctx)
assert any("Personality too long" in t for t in texts)
app_mocks.manager.create_custom_bot.assert_not_called()
def test_custom_bot_create_fails(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A failed create reports an error and does not invalidate the cache."""
app_mocks.manager.create_custom_bot.return_value = False
invoke(
bot, "custom-bot", mock_ctx, "alfred", personality="you are a british butler"
)
texts = sent_texts(mock_ctx)
assert any("Failed to create custom bot" in t for t in texts)
app_mocks.manager.list_custom_bots.assert_not_called()
assert app_mocks.app.bot_cache == {"alfred": ("british butler", "user123")}
def test_list_custom_bots_empty(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Listing with no bots suggests creating one."""
app_mocks.manager.list_custom_bots.return_value = []
invoke(bot, "list-custom-bots", mock_ctx)
texts = sent_texts(mock_ctx)
assert any("No custom bots" in t for t in texts)
def test_list_custom_bots_with_bots(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Listing shows every bot name."""
app_mocks.manager.list_custom_bots.return_value = [
("alfred", "british butler", "user-1"),
("jarvis", "ai assistant", "user-2"),
]
invoke(bot, "list-custom-bots", mock_ctx)
texts = sent_texts(mock_ctx)
assert any("Available Custom Bots" in t for t in texts)
assert any("* alfred" in t and "* jarvis" in t for t in texts)
def test_delete_custom_bot_success(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""The creator can delete their bot; the cache is invalidated."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"12345",
"2024-01-01",
)
app_mocks.manager.delete_custom_bot.return_value = True
app_mocks.manager.list_custom_bots.return_value = []
invoke(bot, "delete-custom-bot", mock_ctx, "alfred")
app_mocks.manager.delete_custom_bot.assert_called_once_with("alfred")
texts = sent_texts(mock_ctx)
assert any("has been deleted" in t for t in texts)
assert app_mocks.app.bot_cache == {}
def test_delete_custom_bot_not_found(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Deleting a non-existent bot reports not found."""
app_mocks.manager.get_custom_bot.return_value = None
invoke(bot, "delete-custom-bot", mock_ctx, "nonexistent")
texts = sent_texts(mock_ctx)
assert any("not found" in t for t in texts)
app_mocks.manager.delete_custom_bot.assert_not_called()
def test_delete_custom_bot_not_owner(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A non-owner cannot delete the bot."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"other-user-id",
"2024-01-01",
)
invoke(bot, "delete-custom-bot", mock_ctx, "alfred")
texts = sent_texts(mock_ctx)
assert any("You can only delete your own" in t for t in texts)
app_mocks.manager.delete_custom_bot.assert_not_called()
def test_delete_custom_bot_delete_fails(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A failed delete reports an error and keeps the cache."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"prompt",
"12345",
"2024-01-01",
)
app_mocks.manager.delete_custom_bot.return_value = False
invoke(bot, "delete-custom-bot", mock_ctx, "alfred")
texts = sent_texts(mock_ctx)
assert any("Failed to delete" in t for t in texts)
app_mocks.manager.list_custom_bots.assert_not_called()
assert app_mocks.app.bot_cache == {"alfred": ("british butler", "user123")}
# ---------------------------------------------------------------------------
# speech: speak / voices
# ---------------------------------------------------------------------------
def test_speak_delegates_to_service(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!speak forwards to SpeechService.speak."""
invoke(bot, "speak", mock_ctx, message="hello world")
app_mocks.speech.speak.assert_awaited_once_with(mock_ctx, message="hello world")
def test_speak_cooldown_blocks_third_invocation(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""speak allows 3 per 30s; the 4th immediate call is rate-limited."""
cmd = bot.get_command("speak")
assert cmd is not None
assert cmd.cooldown is not None
mock_ctx.message.edited_at = None
mock_ctx.message.created_at = datetime(2026, 1, 1, tzinfo=UTC)
for _ in range(3):
cmd._prepare_cooldowns(mock_ctx)
with pytest.raises(commands.CommandOnCooldown) as exc_info:
cmd._prepare_cooldowns(mock_ctx)
asyncio.run(bot.on_command_error(mock_ctx, exc_info.value))
app_mocks.speech.speak.assert_not_awaited()
texts = sent_texts(mock_ctx)
assert any("too quickly" in t for t in texts)
def test_voices(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!voices lists the voice catalog."""
invoke(bot, "voices", mock_ctx)
full = "\n".join(sent_texts(mock_ctx))
assert "Available Voices" in full
assert "af_sarah" in full
assert "Use `!speak" in full
# ---------------------------------------------------------------------------
# images: doodlebob / retcon
# ---------------------------------------------------------------------------
def test_doodlebob_delegates_to_service(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!doodlebob forwards to ImageService.generate."""
invoke(bot, "doodlebob", mock_ctx, message="a centaur")
app_mocks.image.generate.assert_awaited_once_with(mock_ctx, message="a centaur")
def test_retcon_delegates_to_service(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!retcon forwards to ImageService.edit."""
invoke(bot, "retcon", mock_ctx, message="make it blue")
app_mocks.image.edit.assert_awaited_once_with(mock_ctx, message="make it blue")
def test_doodlebob_cooldown_blocks_second_invocation(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A second immediate doodlebob invocation is rate-limited before the LLM."""
cmd = bot.get_command("doodlebob")
assert cmd is not None
assert cmd.cooldown is not None
mock_ctx.message.edited_at = None
mock_ctx.message.created_at = datetime(2026, 1, 1, tzinfo=UTC)
# First invocation consumes the single token (no error).
cmd._prepare_cooldowns(mock_ctx)
# Second immediate invocation is rejected before the service is called.
with pytest.raises(commands.CommandOnCooldown) as exc_info:
cmd._prepare_cooldowns(mock_ctx)
# The error handler turns it into a friendly message.
asyncio.run(bot.on_command_error(mock_ctx, exc_info.value))
app_mocks.image.generate.assert_not_awaited()
texts = sent_texts(mock_ctx)
assert any("too quickly" in t for t in texts)
# ---------------------------------------------------------------------------
# conversation: talkforme
# ---------------------------------------------------------------------------
def test_talkforme_delegates_to_service(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!talkforme parses its args, then forwards to ConversationService.run."""
invoke(bot, "talkforme", mock_ctx, message="a b 3 talking cats")
app_mocks.conversation.run.assert_awaited_once_with(
mock_ctx, "a", "b", "3", "talking cats"
)
def test_talkforme_invalid_args(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!talkforme with too few parts shows usage and calls no service."""
invoke(bot, "talkforme", mock_ctx, message="bot1 bot2")
texts = sent_texts(mock_ctx)
assert any("Usage" in t for t in texts)
app_mocks.conversation.run.assert_not_awaited()
def test_talkforme_cooldown_blocks_second_invocation(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""A second immediate talkforme invocation is rate-limited before the LLM."""
cmd = bot.get_command("talkforme")
assert cmd is not None
assert cmd.cooldown is not None
mock_ctx.message.edited_at = None
mock_ctx.message.created_at = datetime(2026, 1, 1, tzinfo=UTC)
cmd._prepare_cooldowns(mock_ctx)
with pytest.raises(commands.CommandOnCooldown) as exc_info:
cmd._prepare_cooldowns(mock_ctx)
asyncio.run(bot.on_command_error(mock_ctx, exc_info.value))
app_mocks.conversation.run.assert_not_awaited()
texts = sent_texts(mock_ctx)
assert any("too quickly" in t for t in texts)
# ---------------------------------------------------------------------------
# admin: lobotomize / debug / history
# ---------------------------------------------------------------------------
def test_lobotomize(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!lobotomize clears all messages on the shared db."""
invoke(bot, "lobotomize", mock_ctx)
app_mocks.db.clear_all_messages.assert_called_once_with()
texts = sent_texts(mock_ctx)
assert any("cleared" in t for t in texts)
def test_history_bot_not_found(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!history on an unknown bot reports not found."""
app_mocks.manager.get_custom_bot.return_value = None
invoke(bot, "history", mock_ctx, "nonexistent")
texts = sent_texts(mock_ctx)
assert any("not found" in t for t in texts)
app_mocks.db.get_bot_history.assert_not_called()
def test_history_no_history(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!history on a bot with no messages says so."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
app_mocks.db.get_bot_history.return_value = []
invoke(bot, "history", mock_ctx, "alfred")
texts = sent_texts(mock_ctx)
assert any("No chat history" in t and "**alfred**" in t for t in texts)
def test_history_with_data(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!history formats the stored exchange, newest last."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
app_mocks.db.get_bot_history.return_value = [
("hello", "yes master?"),
("what time is it", "it is currently 3pm"),
]
invoke(bot, "history", mock_ctx, "alfred")
full = "\n".join(sent_texts(mock_ctx))
assert "Chat History for **alfred**" in full
assert "what time is it" in full
assert "alfred: it is currently 3pm" in full
def test_history_long_response_chunked(
app_mocks: AppMocks,
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""Long history payloads are split into multiple sends."""
app_mocks.manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
app_mocks.db.get_bot_history.return_value = [
("x" * 2000, "y" * 2000),
]
invoke(bot, "history", mock_ctx, "alfred")
assert mock_ctx.send.call_count >= 2
def test_debug_no_subcommand(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug without a subcommand shows the menu."""
invoke(bot, "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(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug members on a channel without guild members says so."""
mock_ctx.channel.guild = None
invoke(bot, "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(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug members lists the guild members."""
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]
invoke(bot, "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(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug with an unknown subcommand explains the usage."""
invoke(bot, "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(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug members with many members exceeds the chunk limit."""
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
invoke(bot, "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(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug whoami shows user info."""
invoke(bot, "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(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug whoami omits fields the user does not have."""
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
invoke(bot, "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(
bot: commands.Bot,
mock_ctx: MagicMock,
) -> None:
"""!debug tools shows the LLM tool schema."""
invoke(bot, "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()
vibe_bot/tests/test_config.py
+64 -71
View File
@@ -4,6 +4,9 @@ from __future__ import annotations
import subprocess
import sys
from unittest.mock import patch
import pytest
def test_config_defaults() -> None:
@@ -12,17 +15,14 @@ def test_config_defaults() -> None:
for k, v in {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
"CHAT_ENDPOINT_KEY": "test-key",
"COMPLETION_ENDPOINT_KEY": "test-completion-key",
"IMAGE_GEN_ENDPOINT_KEY": "test-image-key",
"IMAGE_EDIT_ENDPOINT_KEY": "test-image-edit-key",
"EMBEDDING_ENDPOINT_KEY": "test-embedding-key",
@@ -39,8 +39,6 @@ def test_config_defaults() -> None:
env_str += f'os.environ["{k}"] = "{v}"\n'
code = f"""
import sys
sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots")
import os
os.environ.clear()
os.environ["PATH"] = "/usr/bin:/bin"
@@ -48,12 +46,10 @@ os.environ["PATH"] = "/usr/bin:/bin"
import vibe_bot.config
assert vibe_bot.config.DISCORD_TOKEN == "test-token"
assert vibe_bot.config.CHAT_ENDPOINT == "https://chat.example.com/v1"
assert vibe_bot.config.COMPLETION_ENDPOINT == "https://completion.example.com/v1"
assert vibe_bot.config.IMAGE_GEN_ENDPOINT == "https://image.example.com/v1"
assert vibe_bot.config.IMAGE_EDIT_ENDPOINT == "https://image-edit.example.com/v1"
assert vibe_bot.config.EMBEDDING_ENDPOINT == "https://embedding.example.com/v1"
assert vibe_bot.config.CHAT_MODEL == "test-chat-model"
assert vibe_bot.config.COMPLETION_MODEL == "test-completion-model"
assert vibe_bot.config.IMAGE_GEN_MODEL == "test-image-model"
assert vibe_bot.config.IMAGE_EDIT_MODEL == "test-image-edit-model"
assert vibe_bot.config.EMBEDDING_MODEL == "test-embedding-model"
@@ -78,20 +74,23 @@ print("OK")
def _run_config_check(env_vars: dict[str, str], expected_error: str) -> None:
"""Run a subprocess that imports config and checks for expected RuntimeError."""
"""Run a subprocess that imports config and calls validate_config().
The import itself must never raise; only validate_config() may raise the
expected RuntimeError for the missing required setting.
"""
env_str = ""
for k, v in env_vars.items():
env_str += f'os.environ["{k}"] = "{v}"\n'
code = f"""
import sys
sys.path.insert(0, "/var/home/ducoterra/Projects/vibe_discord_bots")
import os
os.environ.clear()
os.environ["PATH"] = "/usr/bin:/bin"
{env_str}
try:
import vibe_bot.config
vibe_bot.config.validate_config()
print("NO_ERROR")
except RuntimeError as e:
print(f"ERROR: {{e}}")
@@ -116,12 +115,10 @@ def test_config_missing_discord_token() -> None:
env: dict[str, str] = {
"DISCORD_TOKEN": "",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
@@ -134,12 +131,10 @@ def test_config_missing_chat_endpoint() -> None:
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
@@ -147,35 +142,15 @@ def test_config_missing_chat_endpoint() -> None:
_run_config_check(env, "CHAT_ENDPOINT required")
def test_config_missing_completion_endpoint() -> None:
"""Test that RuntimeError is raised when COMPLETION_ENDPOINT is missing."""
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
}
_run_config_check(env, "COMPLETION_ENDPOINT required")
def test_config_missing_image_gen_endpoint() -> None:
"""Test that RuntimeError is raised when IMAGE_GEN_ENDPOINT is missing."""
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
@@ -188,12 +163,10 @@ def test_config_missing_image_edit_endpoint() -> None:
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
@@ -206,12 +179,10 @@ def test_config_missing_embedding_endpoint() -> None:
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
@@ -224,12 +195,10 @@ def test_config_missing_chat_model() -> None:
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
@@ -237,35 +206,15 @@ def test_config_missing_chat_model() -> None:
_run_config_check(env, "CHAT_MODEL required")
def test_config_missing_completion_model() -> None:
"""Test that RuntimeError is raised when COMPLETION_MODEL is missing."""
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
}
_run_config_check(env, "COMPLETION_MODEL required")
def test_config_missing_image_gen_model() -> None:
"""Test that RuntimeError is raised when IMAGE_GEN_MODEL is missing."""
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "test-embedding-model",
@@ -278,12 +227,10 @@ def test_config_missing_image_edit_model() -> None:
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "",
"EMBEDDING_MODEL": "test-embedding-model",
@@ -296,12 +243,10 @@ def test_config_missing_embedding_model() -> None:
env: dict[str, str] = {
"DISCORD_TOKEN": "test-token",
"CHAT_ENDPOINT": "https://chat.example.com/v1",
"COMPLETION_ENDPOINT": "https://completion.example.com/v1",
"IMAGE_GEN_ENDPOINT": "https://image.example.com/v1",
"IMAGE_EDIT_ENDPOINT": "https://image-edit.example.com/v1",
"EMBEDDING_ENDPOINT": "https://embedding.example.com/v1",
"CHAT_MODEL": "test-chat-model",
"COMPLETION_MODEL": "test-completion-model",
"IMAGE_GEN_MODEL": "test-image-model",
"IMAGE_EDIT_MODEL": "test-image-edit-model",
"EMBEDDING_MODEL": "",
@@ -309,16 +254,64 @@ def test_config_missing_embedding_model() -> None:
_run_config_check(env, "EMBEDDING_MODEL required")
REQUIRED_VARS = (
"DISCORD_TOKEN",
"CHAT_ENDPOINT",
"IMAGE_GEN_ENDPOINT",
"IMAGE_EDIT_ENDPOINT",
"EMBEDDING_ENDPOINT",
"CHAT_MODEL",
"IMAGE_GEN_MODEL",
"IMAGE_EDIT_MODEL",
"EMBEDDING_MODEL",
)
def test_validate_config_passes_with_full_env() -> None:
"""With all required settings present, validate_config() is a no-op."""
from vibe_bot import config
config.validate_config()
@pytest.mark.parametrize("var_name", REQUIRED_VARS)
def test_validate_config_missing_var(var_name: str) -> None:
"""Blanking any single required setting makes validate_config() raise."""
from vibe_bot import config
with (
patch.object(config, var_name, ""),
pytest.raises(RuntimeError, match=f"{var_name} required"),
):
config.validate_config()
def test_import_config_empty_env_no_raise_no_logging() -> None:
"""In an empty env the import succeeds and leaves the root logger unconfigured."""
code = """
import logging
import os
os.environ.clear()
os.environ["PATH"] = "/usr/bin:/bin"
import vibe_bot.config
handlers = logging.getLogger().handlers
assert handlers == [], f"config configured logging: {handlers}"
print("OK")
"""
result = subprocess.run( # noqa: PLW1510
[sys.executable, "-c", code],
capture_output=True,
text=True,
timeout=30,
)
assert result.returncode == 0, f"Subprocess failed: {result.stderr}"
assert "OK" in result.stdout
def test_config_logging_exists() -> None:
"""Test that logging is configured in config module."""
from vibe_bot.config import logger
assert logger is not None
assert logger.name == "vibe_bot.config"
def test_config_embedding_dimension() -> None:
"""Test that EMBEDDING_DIMENSION has expected default value."""
from vibe_bot.config import EMBEDDING_DIMENSION
assert EMBEDDING_DIMENSION == 2048
vibe_bot/tests/test_database.py
+774 -44
View File
@@ -2,17 +2,40 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from datetime import datetime
from typing import TYPE_CHECKING, Any, cast
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
if TYPE_CHECKING:
import sqlite3
from vibe_bot.database import ChatDatabase
def _recent_messages(
db_path: str,
limit: int,
) -> list[tuple[str, str, str, datetime]]:
"""Read the newest rows straight from the database (test-side helper)."""
from vibe_bot.db.connection import connect
conn = connect(db_path)
try:
rows = conn.execute(
"SELECT message_id, username, content, timestamp "
"FROM chat_messages ORDER BY timestamp DESC LIMIT ?",
(limit,),
).fetchall()
finally:
conn.close()
return [
(str(row[0]), str(row[1]), str(row[2]), cast(datetime, row[3])) for row in rows
]
def test_vector_to_bytes(chat_db: ChatDatabase) -> None:
"""Test converting a vector to bytes and back."""
vector: list[float] = [0.1, 0.2, 0.3, 0.4]
@@ -55,6 +78,13 @@ def test_calculate_similarity_negative(chat_db: ChatDatabase) -> None:
assert similarity == pytest.approx(-1.0, abs=1e-6)
def test_calculate_similarity_zero_norm(chat_db: ChatDatabase) -> None:
"""A zero vector has no direction, so its similarity is 0."""
zero = np.zeros(3, dtype=np.float32)
other = np.array([1.0, 0.0, 0.0], dtype=np.float32)
assert chat_db._calculate_similarity(zero, other) == 0.0
def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
"""Test adding a message to the database."""
result = chat_db.add_message(
@@ -67,7 +97,7 @@ def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
)
assert result is True
messages = chat_db.get_recent_messages(limit=10)
messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 1
assert messages[0][0] == "msg-1"
assert messages[0][1] == "testuser"
@@ -76,7 +106,7 @@ def test_add_message(chat_db: ChatDatabase, mock_embedding: MagicMock) -> None:
def test_add_message_no_embedding(chat_db: ChatDatabase) -> None:
"""Test adding a message when embedding generation fails."""
with patch("vibe_bot.llama_wrapper.embedding", return_value=None):
with patch("vibe_bot.llm_client.embedding", return_value=None):
result = chat_db.add_message(
message_id="msg-no-embed",
user_id="user-1",
@@ -106,14 +136,14 @@ def test_add_message_duplicate(
content="Second content",
)
messages = chat_db.get_recent_messages(limit=10)
messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 1
assert messages[0][2] == "Second content"
def test_add_message_failure(chat_db: ChatDatabase) -> None:
"""Test that add_message returns False on database error."""
with patch.object(chat_db, "_vector_to_bytes", side_effect=Exception("fail")):
with patch("vibe_bot.db.messages.connect", return_value=_broken_connection()):
result = chat_db.add_message(
message_id="msg-fail",
user_id="user-1",
@@ -123,11 +153,304 @@ def test_add_message_failure(chat_db: ChatDatabase) -> None:
assert result is False
def test_get_recent_messages(
def _embedding_row_count(db_path: str) -> int:
"""Count the rows stored in message_embeddings."""
import sqlite3
conn = sqlite3.connect(db_path)
count = conn.execute("SELECT COUNT(*) FROM message_embeddings").fetchone()[0]
conn.close()
return int(count)
def test_add_message_embed_false_skips_embedding(chat_db: ChatDatabase) -> None:
"""embed=False neither calls the embedding API nor stores a row."""
with patch("vibe_bot.llm_client.embedding") as mock_embedding:
result = chat_db.add_message(
message_id="msg-assist",
user_id="bot-1",
username="some-bot",
content="assistant reply",
role="assistant",
embed=False,
)
assert result is True
mock_embedding.assert_not_called()
assert _embedding_row_count(chat_db.db_path) == 0
def test_add_message_stores_embedding_by_default(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""Test retrieving recent messages."""
"""The default embed=True still stores exactly one embedding row."""
assert chat_db.add_message(
message_id="msg-embed",
user_id="user-1",
username="testuser",
content="default embed",
)
assert _embedding_row_count(chat_db.db_path) == 1
def test_add_message_stores_embedding_norm(chat_db: ChatDatabase) -> None:
"""add_message stores the L2 norm of the float32-encoded embedding."""
import sqlite3
vector: list[float] = [0.6, 0.8]
with patch("vibe_bot.llm_client.embedding", return_value=vector):
assert chat_db.add_message(
message_id="norm-1",
user_id="u1",
username="alice",
content="normed message",
)
conn = sqlite3.connect(chat_db.db_path)
blob, norm = conn.execute(
"SELECT embedding, norm FROM message_embeddings WHERE message_id = 'norm-1'"
).fetchone()
conn.close()
stored = np.frombuffer(blob, dtype=np.float32)
assert float(norm) == pytest.approx(float(np.linalg.norm(stored)), abs=1e-6)
def test_cleanup_old_messages_no_orphaned_embeddings(chat_db: ChatDatabase) -> None:
"""Deleting the oldest rows must delete their embeddings, not the next ones.
Regression test for the bug where the embedding cleanup re-queried
``chat_messages`` *after* the message delete, so it stripped embeddings
from the next-oldest live rows while leaving the deleted rows' embeddings
behind as orphans.
"""
import sqlite3
with patch("vibe_bot.db.messages.MAX_HISTORY_MESSAGES", 5):
# Seed 7 rows with distinct ascending timestamps and an embedding for
# each, so the "oldest" ordering is deterministic.
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
for i in range(1, 8):
ts = f"2024-01-0{i} 00:00:00"
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, bot_name, timestamp) "
"VALUES (?, ?, ?, ?, ?, ?)",
(f"m{i}", "u1", "alice", f"content {i}", "bot", ts),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding) "
"VALUES (?, ?)",
(f"m{i}", chat_db._vector_to_bytes([0.1, 0.2, 0.3])),
)
conn.commit()
conn.close()
# The 8th insert pushes the count to 8 (> 5), so add_message's
# cleanup must delete exactly the 3 oldest rows and their embeddings.
assert chat_db.add_message(
message_id="m8",
user_id="u1",
username="alice",
content="content 8",
)
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
cursor.execute("SELECT message_id FROM chat_messages")
live = {row[0] for row in cursor.fetchall()}
cursor.execute("SELECT message_id FROM message_embeddings")
embedded = {row[0] for row in cursor.fetchall()}
conn.close()
# The three oldest rows are gone; the rest (incl. the new one) remain.
assert live == {"m4", "m5", "m6", "m7", "m8"}
# No orphaned embeddings and no stripped survivors: the embedding set
# exactly matches the live message rows.
assert embedded == live
def test_role_scopes_history_and_search(chat_db: ChatDatabase) -> None:
"""get_user_history excludes responses; search matches only user rows."""
chat_db.add_message(
message_id="r-1",
user_id="u1",
username="alice",
content="User asks about the weather",
role="user",
)
chat_db.add_message(
message_id="r-1_response",
user_id="bot",
username="some-bot",
content="Bot answers the weather",
role="assistant",
)
# get_user_history returns only the user row (paired with its response).
conversations = chat_db.get_user_history("u1")
assert len(conversations) == 1
assert conversations[0][0] == "User asks about the weather"
assert conversations[0][1] == "Bot answers the weather"
# search_similar_messages only considers user rows, never responses.
results = chat_db.search_similar_messages(
"User asks about the weather", top_k=5, min_similarity=0.0
)
assert len(results) == 1
assert results[0][0] == "User asks about the weather"
assert results[0][1] == "Bot answers the weather"
def test_role_migration_backfills_legacy_rows(
temp_db_path: str,
) -> None:
"""Legacy rows (no role column) are backfilled when ChatDatabase inits."""
import sqlite3
from vibe_bot.database import ChatDatabase
# Create the legacy schema (no role column) with a user row and its
# response row inserted directly.
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"bot_name TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT)"
)
cursor.execute(
"INSERT INTO chat_messages (message_id, user_id, username, content) "
"VALUES ('legacy-1', 'u1', 'alice', 'old question')"
)
cursor.execute(
"INSERT INTO chat_messages (message_id, user_id, username, content) "
"VALUES ('legacy-1_response', 'bot', 'old-bot', 'old answer')"
)
conn.commit()
conn.close()
# Initializing ChatDatabase should add and backfill the role column.
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute("SELECT message_id, role FROM chat_messages ORDER BY message_id")
roles = {row[0]: row[1] for row in cursor.fetchall()}
conn.close()
assert roles["legacy-1"] == "user"
assert roles["legacy-1_response"] == "assistant"
def test_bot_name_migration_adds_column(
temp_db_path: str,
) -> None:
"""A pre-bot_name schema gets the column added when ChatDatabase inits."""
import sqlite3
from vibe_bot.database import ChatDatabase
conn = sqlite3.connect(temp_db_path)
conn.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT, role TEXT)"
)
conn.commit()
conn.close()
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
cursor.execute("PRAGMA table_info(chat_messages)")
columns = {row[1] for row in cursor.fetchall()}
conn.close()
assert "bot_name" in columns
def _seed_pre_norm_db(db_path: str) -> list[tuple[str, str, list[float]]]:
"""Create a pre-norm-schema database (no norm column) with seeded rows."""
import sqlite3
rows: list[tuple[str, str, list[float]]] = [
("n-1", "ask about the sky", [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-2", "ask about the sea", [0.7, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-3", "ask about the sun", [0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-4", "ask about the sand", [0.2, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("n-5", "ask about nothing", [0.0] * 8),
]
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute(
"CREATE TABLE chat_messages ("
"id INTEGER PRIMARY KEY AUTOINCREMENT,"
"message_id TEXT UNIQUE, user_id TEXT, username TEXT, content TEXT,"
"bot_name TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
"channel_id TEXT, guild_id TEXT, role TEXT)"
)
cursor.execute(
"CREATE TABLE message_embeddings ("
"message_id TEXT PRIMARY KEY, embedding BLOB,"
"FOREIGN KEY (message_id) REFERENCES chat_messages(message_id))"
)
for i, (message_id, content, vector) in enumerate(rows):
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'u1', 'alice', ?, 'user', ?)",
(message_id, content, f"2024-01-0{i + 1} 00:00:00"),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'bot-1', 'some-bot', ?, 'assistant', ?)",
(
f"{message_id}_response",
f"response {i + 1}",
f"2024-01-0{i + 1} 00:00:01",
),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding) VALUES (?, ?)",
(message_id, np.array(vector, dtype=np.float32).tobytes()),
)
conn.commit()
conn.close()
return rows
def test_norm_migration_backfills_stored_norms(temp_db_path: str) -> None:
"""A pre-norm database gets the norm column added and backfilled on init."""
import sqlite3
from vibe_bot.database import ChatDatabase
rows = _seed_pre_norm_db(temp_db_path)
ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
stored = dict(conn.execute("SELECT message_id, norm FROM message_embeddings"))
conn.close()
for message_id, _content, vector in rows:
expected = float(np.linalg.norm(np.array(vector, dtype=np.float32)))
assert stored[message_id] == pytest.approx(expected, abs=1e-5)
def test_recent_messages_desc_order(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""Newest-first ordering of stored messages."""
chat_db.add_message(
message_id="msg-1",
user_id="u1",
@@ -147,17 +470,17 @@ def test_get_recent_messages(
content="Third",
)
messages = chat_db.get_recent_messages(limit=2)
messages = _recent_messages(chat_db.db_path, 2)
assert len(messages) == 2
assert messages[0][2] == "Third"
assert messages[1][2] == "Second"
def test_get_recent_messages_limit(
def test_recent_messages_limit(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""Test that get_recent_messages respects the limit."""
"""The newest-rows query respects the limit."""
for i in range(5):
chat_db.add_message(
message_id=f"msg-{i}",
@@ -166,10 +489,25 @@ def test_get_recent_messages_limit(
content=f"Message {i}",
)
messages = chat_db.get_recent_messages(limit=3)
messages = _recent_messages(chat_db.db_path, 3)
assert len(messages) == 3
def test_recent_messages_returns_datetime(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""The timestamp column comes back as a real datetime, not a string."""
chat_db.add_message(
message_id="dt-1",
user_id="u1",
username="alice",
content="fresh message",
)
messages = _recent_messages(chat_db.db_path, 1)
assert isinstance(messages[0][3], datetime)
def test_clear_all_messages(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
@@ -190,7 +528,7 @@ def test_clear_all_messages(
chat_db.clear_all_messages()
messages = chat_db.get_recent_messages(limit=10)
messages = _recent_messages(chat_db.db_path, 10)
assert len(messages) == 0
@@ -264,6 +602,25 @@ def test_image_generation_estimate_after_capping(
assert estimate == pytest.approx(99.5)
def _broken_connection() -> MagicMock:
"""A mock connection whose first cursor.execute raises."""
fake_conn = MagicMock()
fake_conn.cursor.return_value.execute.side_effect = Exception("db error")
return fake_conn
def test_record_image_generation_time_failure(chat_db: ChatDatabase) -> None:
"""A database error while recording yields False, not an exception."""
with patch("vibe_bot.db.timing.connect", return_value=_broken_connection()):
assert chat_db.record_image_generation_time(1.0) is False
def test_image_generation_estimate_failure(chat_db: ChatDatabase) -> None:
"""A database error while reading yields None, not an exception."""
with patch("vibe_bot.db.timing.connect", return_value=_broken_connection()):
assert chat_db.get_image_generation_time_estimate() is None
def test_get_user_history(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
@@ -320,6 +677,39 @@ def test_get_user_history_excludes_bot(
assert len(conversations) == 0
def test_get_bot_history(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
) -> None:
"""get_bot_history pairs user messages with responses for one bot."""
chat_db.add_message(
message_id="bh-1",
user_id="u1",
username="alice",
content="bot question",
bot_name="alfred",
)
chat_db.add_message(
message_id="bh-1_response",
user_id="bot-1",
username="some-bot",
content="bot answer",
bot_name="alfred",
role="assistant",
embed=False,
)
chat_db.add_message(
message_id="bh-2",
user_id="u1",
username="alice",
content="unanswered question",
bot_name="alfred",
)
history = chat_db.get_bot_history("alfred")
assert history == [("bot question", "bot answer")]
def test_get_conversation_context(
chat_db: ChatDatabase,
mock_embedding: MagicMock,
@@ -349,21 +739,333 @@ def test_get_conversation_context_empty(chat_db: ChatDatabase) -> None:
assert context == []
def _query_vector() -> list[float]:
"""A 8-dim query vector along the first axis."""
return [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
def _seed_search_rows(chat_db: ChatDatabase) -> list[tuple[str, str, list[float]]]:
"""Seed user/response rows with known embeddings (plus exclusion traps)."""
import sqlite3
rows: list[tuple[str, str, list[float]]] = [
("s-1", "ask about the sky", [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-2", "ask about the sea", [0.7, 0.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-3", "ask about the sun", [0.5, 0.5, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
("s-4", "ask about the sand", [0.2, 0.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
]
def stored_norm(vector: list[float]) -> float:
return float(np.linalg.norm(np.array(vector, dtype=np.float32)))
conn = sqlite3.connect(chat_db.db_path)
cursor = conn.cursor()
for i, (message_id, content, vector) in enumerate(rows):
timestamp = f"2024-01-0{i + 1} 00:00:00"
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, ?, ?, ?, 'user', ?)",
(message_id, "u1", "alice", content, timestamp),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, ?, ?, ?, 'assistant', ?)",
(
f"{message_id}_response",
"bot-1",
"some-bot",
f"response {i + 1}",
f"2024-01-0{i + 1} 00:00:01",
),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(message_id, chat_db._vector_to_bytes(vector), stored_norm(vector)),
)
# A user row with the top possible similarity but no response row: the
# JOIN must exclude it instead of returning a NULL response.
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES ('s-5', 'u1', 'alice', 'orphan question', 'user', "
"'2024-01-05 00:00:00')",
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(
"s-5",
chat_db._vector_to_bytes(_query_vector()),
stored_norm(_query_vector()),
),
)
# An assistant row carrying an embedding: the role filter must skip it.
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES ('s-6', 'bot-1', 'some-bot', 'assistant noise', 'assistant', "
"'2024-01-06 00:00:00')",
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(
"s-6",
chat_db._vector_to_bytes(_query_vector()),
stored_norm(_query_vector()),
),
)
conn.commit()
conn.close()
return rows
def test_search_matches_reference_topk_and_ordering(
chat_db: ChatDatabase,
) -> None:
"""The JOINed search matches the per-row reference: same top-k and order."""
rows = _seed_search_rows(chat_db)
query = _query_vector()
expected: list[tuple[str, str, float]] = []
query_arr = np.array(query, dtype=np.float32)
for i, (_message_id, content, vector) in enumerate(rows):
stored = np.frombuffer(chat_db._vector_to_bytes(vector), dtype=np.float32)
similarity = float(
np.dot(query_arr, stored)
/ (np.linalg.norm(query_arr) * np.linalg.norm(stored))
)
expected.append((content, f"response {i + 1}", similarity))
expected.sort(key=lambda item: item[2], reverse=True)
expected_top = expected[:3]
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = chat_db.search_similar_messages(
"query text", top_k=3, min_similarity=0.0
)
assert len(results) == 3
for (content, response, actual), (_ec, er, reference) in zip(
results, expected_top, strict=True
):
assert content == _ec
assert response == er
assert actual == pytest.approx(reference, abs=1e-5)
def test_search_with_stored_norms_matches_pre_norm_reference(
temp_db_path: str,
) -> None:
"""Stored-norm search is identical to per-vector renormalization.
Seeds a pre-norm database, migrates it, then asserts the stored-norm
search (top-k + ordering + similarities) matches a reference that
renormalizes every candidate vector inline — the pre-optimization
algorithm.
"""
from vibe_bot.database import ChatDatabase
rows = _seed_pre_norm_db(temp_db_path)
db = ChatDatabase(db_path=temp_db_path)
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=3, min_similarity=0.0)
query_arr = np.array(query, dtype=np.float32)
query_norm = float(np.linalg.norm(query_arr))
expected: list[tuple[str, str, float]] = []
for i, (_message_id, content, vector) in enumerate(rows):
stored = np.array(vector, dtype=np.float32)
stored_norm = float(np.linalg.norm(stored))
similarity = (
0.0
if stored_norm == 0
else float(np.dot(query_arr, stored) / (query_norm * stored_norm))
)
expected.append((content, f"response {i + 1}", similarity))
expected.sort(key=lambda item: item[2], reverse=True)
assert len(results) == 3
for (content, response, actual), (ec, er, reference) in zip(
results, expected[:3], strict=True
):
assert content == ec
assert response == er
assert actual == pytest.approx(reference, abs=1e-6)
def test_search_null_norm_row_scores_zero(temp_db_path: str) -> None:
"""A row whose norm was never written scores 0 instead of crashing."""
import sqlite3
from vibe_bot.database import ChatDatabase
_rows = _seed_pre_norm_db(temp_db_path)
db = ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
conn.execute("UPDATE message_embeddings SET norm = NULL WHERE message_id = 'n-1'")
conn.commit()
conn.close()
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=10, min_similarity=0.0)
scores = {content: similarity for content, _response, similarity in results}
assert scores["ask about the sky"] == 0.0
# All four scored rows plus the zero-vector row (0.0 passes min_similarity=0.0).
assert len(results) == 5
def test_search_mixed_embedding_dims_falls_back_to_per_row(temp_db_path: str) -> None:
"""Rows with different blob lengths (mid-life model change) don't crash."""
import sqlite3
from vibe_bot.database import ChatDatabase
wide = [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
narrow = [0.5, 0.5, 0.0, 0.0]
db = ChatDatabase(db_path=temp_db_path)
conn = sqlite3.connect(temp_db_path)
cursor = conn.cursor()
for message_id, content, vector in (
("m-1", "wide question", wide),
("m-2", "narrow question", narrow),
):
blob = np.array(vector, dtype=np.float32).tobytes()
stored_norm = float(np.linalg.norm(np.frombuffer(blob, dtype=np.float32)))
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'u1', 'alice', ?, 'user', '2024-01-01 00:00:00')",
(message_id, content),
)
cursor.execute(
"INSERT INTO chat_messages "
"(message_id, user_id, username, content, role, timestamp) "
"VALUES (?, 'bot-1', 'some-bot', ?, 'assistant', '2024-01-01 00:00:01')",
(f"{message_id}_response", f"{message_id} answer"),
)
cursor.execute(
"INSERT INTO message_embeddings (message_id, embedding, norm) "
"VALUES (?, ?, ?)",
(message_id, blob, stored_norm),
)
conn.commit()
conn.close()
query = _query_vector()
with patch("vibe_bot.llm_client.embedding", return_value=query):
results = db.search_similar_messages("query text", top_k=10, min_similarity=0.0)
scores = {content: similarity for content, _response, similarity in results}
# Both rows are returned: the 8-dim row scores exactly, the 4-dim row is
# zero-padded to the query dim and still scores sanely.
assert set(scores) == {"wide question", "narrow question"}
wide_norm = float(np.linalg.norm(np.array(wide, dtype=np.float32)))
narrow_norm = float(np.linalg.norm(np.array(narrow, dtype=np.float32)))
assert scores["wide question"] == pytest.approx(0.9 / wide_norm, abs=1e-5)
assert scores["narrow question"] == pytest.approx(0.5 / narrow_norm, abs=1e-5)
assert scores["wide question"] > scores["narrow question"]
def test_search_issues_single_select_per_call(chat_db: ChatDatabase) -> None:
"""search_similar_messages issues exactly one SELECT per call (no N+1)."""
import vibe_bot.db.search as db_search
from vibe_bot.db.connection import connect as realconnect
chat_db.add_message(
message_id="n-1",
user_id="u1",
username="alice",
content="one question",
)
chat_db.add_message(
message_id="n-1_response",
user_id="bot-1",
username="some-bot",
content="one answer",
role="assistant",
embed=False,
)
select_statements: list[str] = []
class _TracingCursor:
def __init__(self, cursor: sqlite3.Cursor) -> None:
self._cursor = cursor
def execute(self, sql: str, *args: Any) -> Any:
if sql.strip().upper().startswith("SELECT"):
select_statements.append(sql)
return self._cursor.execute(sql, *args)
def fetchall(self) -> Any:
return self._cursor.fetchall()
def fetchone(self) -> Any:
return self._cursor.fetchone()
class _TracingConnection:
def __init__(self, conn: sqlite3.Connection) -> None:
self._conn = conn
def cursor(self) -> _TracingCursor:
return _TracingCursor(self._conn.cursor())
def close(self) -> None:
self._conn.close()
def tracingconnect(db_path: str) -> _TracingConnection:
return _TracingConnection(realconnect(db_path))
with patch.object(db_search, "connect", side_effect=tracingconnect):
results = chat_db.search_similar_messages(
"one question", top_k=5, min_similarity=0.0
)
assert len(results) == 1
assert results[0][0] == "one question"
assert results[0][1] == "one answer"
assert len(select_statements) == 1
def test_search_empty_query_embedding_returns_empty(chat_db: ChatDatabase) -> None:
"""A failed query embedding yields no results."""
with patch("vibe_bot.llm_client.embedding", return_value=[]):
assert chat_db.search_similar_messages("anything") == []
def test_search_zero_query_vector_returns_empty(chat_db: ChatDatabase) -> None:
"""A zero-norm query vector yields no results (no division by zero)."""
with patch("vibe_bot.llm_client.embedding", return_value=[0.0] * 8):
assert chat_db.search_similar_messages("anything") == []
def test_custom_bot_create(custom_bot_manager: Any) -> None:
"""Test creating a custom bot."""
"""Test creating a custom bot returns "created" for a new name."""
result = custom_bot_manager.create_custom_bot(
bot_name="alfred",
system_prompt="You are a british butler",
created_by="user-123",
)
assert result is True
assert result == "created"
def test_custom_bot_create_duplicate(
custom_bot_manager: Any,
) -> None:
"""Test creating a duplicate custom bot replaces the old one."""
custom_bot_manager.create_custom_bot(
first = custom_bot_manager.create_custom_bot(
bot_name="alfred",
system_prompt="First personality",
created_by="user-1",
@@ -373,13 +1075,25 @@ def test_custom_bot_create_duplicate(
system_prompt="Second personality",
created_by="user-1",
)
assert result is True
assert first == "created"
assert result == "replaced"
bot = custom_bot_manager.get_custom_bot("alfred")
assert bot is not None
assert bot[1] == "Second personality"
def test_custom_bot_create_failure(custom_bot_manager: Any) -> None:
"""A database error while creating yields False, not an exception."""
with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()):
result = custom_bot_manager.create_custom_bot(
bot_name="failbot",
system_prompt="a long enough personality",
created_by="user-1",
)
assert result is False
def test_custom_bot_create_case_insensitive(
custom_bot_manager: Any,
) -> None:
@@ -399,6 +1113,18 @@ def test_custom_bot_get_not_found(custom_bot_manager: Any) -> None:
assert result is None
def test_custom_bot_get_returns_datetime(custom_bot_manager: Any) -> None:
"""created_at comes back as a real datetime, not a string."""
custom_bot_manager.create_custom_bot(
bot_name="dtbot",
system_prompt="a long enough personality",
created_by="user-1",
)
result = custom_bot_manager.get_custom_bot("dtbot")
assert result is not None
assert isinstance(result[3], datetime)
def test_custom_bot_get_returns_correct_data(
custom_bot_manager: Any,
) -> None:
@@ -413,8 +1139,7 @@ def test_custom_bot_get_returns_correct_data(
assert result[0] == "testbot"
assert result[1] == "test prompt"
assert result[2] == "creator-1"
assert result[3] is not None
assert "20" in result[3]
assert isinstance(result[3], datetime)
def test_custom_bot_list_empty(custom_bot_manager: Any) -> None:
@@ -440,6 +1165,23 @@ def test_custom_bot_list(custom_bot_manager: Any) -> None:
assert len(bots) == 2
def test_custom_bot_list_by_creator(custom_bot_manager: Any) -> None:
"""list_custom_bots filters by creator when user_id is given."""
custom_bot_manager.create_custom_bot(
bot_name="bot-x",
system_prompt="prompt x",
created_by="user-1",
)
custom_bot_manager.create_custom_bot(
bot_name="bot-y",
system_prompt="prompt y",
created_by="user-2",
)
bots = custom_bot_manager.list_custom_bots(user_id="user-1")
assert [bot[0] for bot in bots] == ["bot-x"]
def test_custom_bot_delete(custom_bot_manager: Any) -> None:
"""Test deleting a custom bot."""
custom_bot_manager.create_custom_bot(
@@ -462,32 +1204,18 @@ def test_custom_bot_delete_nonexistent(
assert result is False
def test_custom_bot_deactivate(custom_bot_manager: Any) -> None:
"""Test deactivating a custom bot."""
custom_bot_manager.create_custom_bot(
bot_name="inactive-bot",
system_prompt="will be deactivated",
created_by="user-1",
)
result = custom_bot_manager.deactivate_custom_bot("inactive-bot")
assert result is True
bot = custom_bot_manager.get_custom_bot("inactive-bot")
assert bot is None
def test_custom_bot_deactivate_nonexistent(
custom_bot_manager: Any,
) -> None:
"""Test deactivating a non-existent bot returns False."""
result = custom_bot_manager.deactivate_custom_bot("nonexistent")
assert result is False
def test_custom_bot_delete_failure(custom_bot_manager: Any) -> None:
"""A database error while deleting yields False, not an exception."""
with patch("vibe_bot.db.bots.connect", return_value=_broken_connection()):
assert custom_bot_manager.delete_custom_bot("whatever") is False
def test_custom_bot_list_excludes_inactive(
custom_bot_manager: Any,
) -> None:
"""Test that list_custom_bots excludes deactivated bots."""
"""Test that list_custom_bots excludes bots with is_active = 0."""
import sqlite3
custom_bot_manager.create_custom_bot(
bot_name="active-bot",
system_prompt="stays active",
@@ -498,7 +1226,12 @@ def test_custom_bot_list_excludes_inactive(
system_prompt="should not appear",
created_by="user-1",
)
custom_bot_manager.deactivate_custom_bot("deactivated-bot")
conn = sqlite3.connect(custom_bot_manager.db_path)
conn.execute(
"UPDATE custom_bots SET is_active = 0 WHERE bot_name = 'deactivated-bot'"
)
conn.commit()
conn.close()
bots = custom_bot_manager.list_custom_bots()
assert len(bots) == 1
@@ -532,16 +1265,13 @@ def test_database_get_database_singleton(temp_db_path: str) -> None:
db2 = get_database()
assert db1 is db2
db1.client.close()
def test_database_init_creates_tables(temp_db_path: str) -> None:
"""Test that database initialization creates the expected tables."""
from vibe_bot.database import ChatDatabase, CustomBotManager
db = ChatDatabase(db_path=temp_db_path)
ChatDatabase(db_path=temp_db_path)
CustomBotManager(db_path=temp_db_path)
db.client.close()
import sqlite3
vibe_bot/tests/test_docs.py
+109
View File
@@ -0,0 +1,109 @@
"""Docs-as-tests: README.md stays in sync with the real commands and file tree."""
from __future__ import annotations
import re
import subprocess
from pathlib import Path
from discord.ext import commands
REPO_ROOT = Path(__file__).resolve().parents[2]
README = REPO_ROOT / "README.md"
BINARY_SUFFIXES = {
".bin",
".db",
".gif",
".ico",
".jpeg",
".jpg",
".mp3",
".onnx",
".png",
".wav",
}
def _git(args: list[str]) -> list[str]:
result = subprocess.run(
["git", *args],
cwd=REPO_ROOT,
capture_output=True,
text=True,
check=True,
)
return result.stdout.splitlines()
def _committed_files() -> set[str]:
"""Files that make up the current tree: index minus deletions, plus untracked."""
deleted: set[str] = set()
for line in _git(["status", "--porcelain"]):
if "D" in line[:2]:
deleted.add(line[3:])
tracked = set(_git(["ls-files"])) - deleted
untracked = set(_git(["ls-files", "--others", "--exclude-standard"]))
files: set[str] = set()
for path in tracked | untracked:
parts = Path(path).parts
if any(part.startswith(".") for part in parts):
continue
if Path(path).suffix.lower() in BINARY_SUFFIXES:
continue
files.add(path)
return files
def _readme_tree_paths() -> set[str]:
"""Reconstruct the relative paths (and directory entries) from the README tree."""
text = README.read_text(encoding="utf-8")
match = re.search(r"## File Structure\s*```text\n(.*?)```", text, re.DOTALL)
assert match is not None, "File Structure tree block not found in README"
lines = match.group(1).splitlines()
assert lines, "File Structure tree block is empty"
root = lines[0].split("#")[0].strip().rstrip("/")
assert root == REPO_ROOT.name, f"Tree root {root!r} != repo dir {REPO_ROOT.name!r}"
stack: list[str] = []
paths: set[str] = set()
entry_re = re.compile(
r"^(?P<prefix>(?:[│ ] )*)(?:├── |└── )(?P<name>\S.*?)(?:\s+#.*)?$"
)
for line in lines[1:]:
m = entry_re.match(line)
assert m is not None, f"Unparseable tree line: {line!r}"
depth = len(m.group("prefix")) // 4
name = m.group("name").strip()
if name.endswith("/"):
stack = stack[:depth] + [name.rstrip("/")]
paths.add("/".join(stack) + "/")
else:
paths.add("/".join(stack[:depth] + [name]))
return paths
def test_readme_documents_every_registered_command(bot: commands.Bot) -> None:
"""Every command registered on the bot is documented in README.md."""
readme = README.read_text(encoding="utf-8")
assert len(bot.commands) >= 11
missing = [cmd.name for cmd in bot.commands if cmd.name not in readme]
assert missing == [], f"Commands missing from README: {missing}"
def test_readme_file_tree_matches_actual_tree() -> None:
"""The README tree covers every committed non-dotfile, non-binary file."""
actual = _committed_files()
readme_paths = _readme_tree_paths()
missing_in_readme = actual - readme_paths
assert (
missing_in_readme == set()
), f"Files missing from the README tree: {sorted(missing_in_readme)}"
for path in readme_paths:
target = REPO_ROOT / path
if path.endswith("/"):
assert target.is_dir(), f"README tree lists missing directory: {path}"
else:
assert target.is_file(), f"README tree lists missing file: {path}"
vibe_bot/tests/test_llama_wrapper.py
-150
View File
@@ -1,150 +0,0 @@
"""Tests for the llama_wrapper module."""
from __future__ import annotations
import base64
import tempfile
from io import BytesIO
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch
import numpy as np
from vibe_bot.config import (
CHAT_ENDPOINT,
CHAT_ENDPOINT_KEY,
CHAT_MODEL,
EMBEDDING_ENDPOINT,
EMBEDDING_ENDPOINT_KEY,
IMAGE_EDIT_ENDPOINT,
IMAGE_EDIT_ENDPOINT_KEY,
IMAGE_GEN_ENDPOINT,
IMAGE_GEN_ENDPOINT_KEY,
)
from vibe_bot.llama_wrapper import (
chat_completion,
chat_completion_instruct,
embedding,
image_edit,
image_generation,
)
TEMPDIR = Path(tempfile.mkdtemp())
def test_chat_completion_think() -> None:
"""Test chat completion with think model."""
chat_completion(
system_prompt="You are a helpful assistant.",
user_prompt="Tell me about Everquest",
openai_url=CHAT_ENDPOINT,
openai_api_key=CHAT_ENDPOINT_KEY,
model=CHAT_MODEL,
max_tokens=100,
)
def test_chat_completion_instruct() -> None:
"""Test chat completion with instruct model."""
chat_completion_instruct(
system_prompt="You are a helpful assistant.",
user_prompt="Tell me about Everquest",
openai_url=CHAT_ENDPOINT,
openai_api_key=CHAT_ENDPOINT_KEY,
model=CHAT_MODEL,
max_tokens=100,
)
def test_image_generation() -> None:
"""Test image generation endpoint."""
with patch("vibe_bot.llama_wrapper.openai.OpenAI") as mock_openai:
mock_response = MagicMock()
mock_data = MagicMock()
mock_data.b64_json = base64.b64encode(b"fake image data").decode()
mock_response.data = [mock_data]
mock_openai.return_value.images.generate.return_value = mock_response
result = image_generation(
prompt="Generate an image of a horse",
openai_url=IMAGE_GEN_ENDPOINT,
openai_api_key=IMAGE_GEN_ENDPOINT_KEY,
)
assert result == base64.b64encode(b"fake image data").decode()
def test_image_edit() -> None:
"""Test image edit endpoint."""
with patch("vibe_bot.llama_wrapper.openai.OpenAI") as mock_openai:
mock_response = MagicMock()
mock_data = MagicMock()
mock_data.b64_json = base64.b64encode(b"fake edited image data").decode()
mock_response.data = [mock_data]
mock_openai.return_value.images.edit.return_value = mock_response
result = image_edit(
image=BytesIO(b"fake image"),
prompt="Paint the words 'horse' on the horse.",
openai_url=IMAGE_EDIT_ENDPOINT,
openai_api_key=IMAGE_EDIT_ENDPOINT_KEY,
)
assert result == base64.b64encode(b"fake edited image data").decode()
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
"""Calculate cosine similarity between two arrays.
Returns a value close to 1 for similar vectors,
close to 0 for orthogonal vectors,
and close to -1 for opposite vectors.
"""
a_arr, b_arr = np.array(a), np.array(b)
return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))
EMBEDDING_SIMILARITY_HIGH = 0.9
EMBEDDING_SIMILARITY_LOW = 0.5
def test_embeddings() -> None:
"""Test embedding similarity for similar and different texts."""
mock_horse_vec = [0.8] * 1024 + [0.6] * 1024
mock_horse_also_vec = [0.79] * 1024 + [0.61] * 1024
mock_donkey_vec = [-0.8] * 1024 + [-0.6] * 1024
def mock_post(*args: Any, **kwargs: Any) -> MagicMock:
json_data = kwargs.get("json", {})
text = json_data["input"][0]
if "horse" in text and "donkey" not in text and "also" not in text:
embedding_data = mock_horse_vec
elif "also" in text:
embedding_data = mock_horse_also_vec
else:
embedding_data = mock_donkey_vec
mock_resp = MagicMock()
mock_resp.json.return_value = {"data": [{"embedding": embedding_data}]}
return mock_resp
with patch("vibe_bot.llama_wrapper.requests.post", side_effect=mock_post):
result1 = embedding(
"this is a horse",
openai_url=EMBEDDING_ENDPOINT,
openai_api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
result2 = embedding(
"this is a horse also",
openai_url=EMBEDDING_ENDPOINT,
openai_api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
result3 = embedding(
"this is a donkey",
openai_url=EMBEDDING_ENDPOINT,
openai_api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
similarity_1 = _cosine_similarity(np.array(result1), np.array(result2))
assert similarity_1 > EMBEDDING_SIMILARITY_HIGH
similarity_2 = _cosine_similarity(np.array(result1), np.array(result3))
assert similarity_2 < EMBEDDING_SIMILARITY_LOW
vibe_bot/tests/test_llm_client.py
+419
View File
@@ -0,0 +1,419 @@
"""Tests for the llm_client module."""
from __future__ import annotations
import base64
from io import BytesIO
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock, patch
import numpy as np
import pytest
from vibe_bot.config import (
CHAT_MODEL,
EMBEDDING_ENDPOINT,
EMBEDDING_ENDPOINT_KEY,
)
from vibe_bot.llm_client import (
chat_complete,
chat_completion_instruct,
embedding,
image_edit,
image_generation,
)
@pytest.mark.live
def test_chat_complete_live() -> None:
"""Live call to the chat endpoint via the core async ``chat_complete``.
Unmocked: requires network access to the configured chat API. Ported from
the former ``test_chat_completion_think`` (its sync ``chat_completion``
wrapper was deleted).
"""
import asyncio
result = asyncio.run(
chat_complete(
[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Tell me about Everquest"},
],
model=CHAT_MODEL,
max_tokens=100,
)
)
assert isinstance(result, str)
@pytest.mark.live
def test_chat_completion_instruct_live() -> None:
"""Live call to the chat endpoint via the async instruct adapter.
Unmocked: requires network access to the configured chat API.
"""
import asyncio
result = asyncio.run(
chat_completion_instruct(
system_prompt="You are a helpful assistant.",
user_prompt="Tell me about Everquest",
model=CHAT_MODEL,
max_tokens=100,
)
)
assert isinstance(result, str)
def test_image_generation() -> None:
"""Image generation returns the first b64 payload from the API."""
import asyncio
mock_client = MagicMock()
mock_data = MagicMock()
mock_data.b64_json = base64.b64encode(b"fake image data").decode()
mock_response = MagicMock()
mock_response.data = [mock_data]
mock_client.images.generate = AsyncMock(return_value=mock_response)
with patch("vibe_bot.llm.images.get_image_gen_client", return_value=mock_client):
result = asyncio.run(
image_generation(
prompt="Generate an image of a horse",
model="test-image-model",
)
)
assert result == base64.b64encode(b"fake image data").decode()
def test_image_generation_api_error_returns_empty() -> None:
"""A 4xx/5xx (APIStatusError) from the image API returns "" without raising."""
import asyncio
import openai
mock_client = MagicMock()
mock_client.images.generate = AsyncMock(
side_effect=openai.APIStatusError(
"boom",
response=MagicMock(),
body=None,
)
)
with patch("vibe_bot.llm.images.get_image_gen_client", return_value=mock_client):
result = asyncio.run(
image_generation(
prompt="Generate an image of a horse",
model="test-image-model",
)
)
assert result == ""
def test_image_edit() -> None:
"""Image edit returns the first b64 payload from the API."""
import asyncio
mock_client = MagicMock()
mock_data = MagicMock()
mock_data.b64_json = base64.b64encode(b"fake edited image data").decode()
mock_response = MagicMock()
mock_response.data = [mock_data]
mock_client.images.edit = AsyncMock(return_value=mock_response)
with patch("vibe_bot.llm.images.get_image_edit_client", return_value=mock_client):
result = asyncio.run(
image_edit(
image=BytesIO(b"fake image"),
prompt="Paint the words 'horse' on the horse.",
model="test-image-edit-model",
)
)
assert result == base64.b64encode(b"fake edited image data").decode()
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
"""Calculate cosine similarity between two arrays.
Returns a value close to 1 for similar vectors,
close to 0 for orthogonal vectors,
and close to -1 for opposite vectors.
"""
a_arr, b_arr = np.array(a), np.array(b)
return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))
EMBEDDING_SIMILARITY_HIGH = 0.9
EMBEDDING_SIMILARITY_LOW = 0.5
def _mock_embedding_session(
post: MagicMock,
) -> MagicMock:
"""Build a mock requests.Session whose .post is ``post``."""
session = MagicMock()
session.post = post
return session
def test_embeddings() -> None:
"""Embedding similarity for similar and different texts."""
mock_horse_vec = [0.8] * 1024 + [0.6] * 1024
mock_horse_also_vec = [0.79] * 1024 + [0.61] * 1024
mock_donkey_vec = [-0.8] * 1024 + [-0.6] * 1024
def mock_post(*args: Any, **kwargs: Any) -> MagicMock:
json_data = kwargs.get("json", {})
text = json_data["input"][0]
if "horse" in text and "donkey" not in text and "also" not in text:
embedding_data = mock_horse_vec
elif "also" in text:
embedding_data = mock_horse_also_vec
else:
embedding_data = mock_donkey_vec
mock_resp = MagicMock()
mock_resp.json.return_value = {"data": [{"embedding": embedding_data}]}
return mock_resp
session = _mock_embedding_session(MagicMock(side_effect=mock_post))
with patch("vibe_bot.llm_client.get_embedding_session", return_value=session):
result1 = embedding(
"this is a horse",
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
result2 = embedding(
"this is a horse also",
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
result3 = embedding(
"this is a donkey",
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
similarity_1 = _cosine_similarity(np.array(result1), np.array(result2))
assert similarity_1 > EMBEDDING_SIMILARITY_HIGH
similarity_2 = _cosine_similarity(np.array(result1), np.array(result3))
assert similarity_2 < EMBEDDING_SIMILARITY_LOW
def test_embedding_non_json_2xx_returns_empty() -> None:
"""A 2xx response with a non-JSON body must return [] without raising.
Regression test for ``resp.json()`` sitting outside the try block, so an
HTML error page (or any non-JSON 2xx body) raised JSONDecodeError out of
``embedding`` and, through it, out of ``get_conversation_context``.
"""
mock_resp = MagicMock()
mock_resp.raise_for_status.return_value = None
mock_resp.json.side_effect = ValueError("<html>rate limited</html>")
session = _mock_embedding_session(MagicMock(return_value=mock_resp))
with patch("vibe_bot.llm_client.get_embedding_session", return_value=session):
result = embedding(
"this is a horse",
url=EMBEDDING_ENDPOINT,
api_key=EMBEDDING_ENDPOINT_KEY,
model="embed",
)
assert result == []
def test_chat_client_singleton_identity() -> None:
"""The shared chat client is built once and reused across calls."""
from vibe_bot import llm_client
client1 = llm_client.get_chat_client()
client2 = llm_client.get_chat_client()
assert client1 is client2
def test_image_gen_client_singleton_identity() -> None:
"""The shared image-generation client is built once, from a cold start."""
import vibe_bot.llm.images as images_mod
saved = images_mod._image_gen_client
images_mod._image_gen_client = None
try:
client1 = images_mod.get_image_gen_client()
client2 = images_mod.get_image_gen_client()
assert client1 is client2
finally:
images_mod._image_gen_client = saved
def test_image_edit_client_singleton_identity() -> None:
"""The shared image-edit client is built once, from a cold start."""
import vibe_bot.llm.images as images_mod
saved = images_mod._image_edit_client
images_mod._image_edit_client = None
try:
client1 = images_mod.get_image_edit_client()
client2 = images_mod.get_image_edit_client()
assert client1 is client2
finally:
images_mod._image_edit_client = saved
def test_flows_build_no_new_clients_or_sessions(
mock_ctx: MagicMock,
temp_db_path: str,
) -> None:
"""A full !doodlebob + chat turn constructs no new clients or sessions.
Every shared client and the embedding session are built once at
"startup"; running the whole image flow and a chat turn through the real
singletons (with only HTTP mocked) must not construct another
AsyncOpenAI client or requests.Session.
"""
import asyncio
import openai
import requests
from vibe_bot import llm_client
from vibe_bot.database import ChatDatabase
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.image_service import ImageService
# Startup: build every shared client and the embedding session once.
chat_client = llm_client.get_chat_client()
gen_client = llm_client.get_image_gen_client()
edit_client = llm_client.get_image_edit_client()
llm_client.get_embedding_session()
counts = {"async_openai": 0, "session": 0}
real_session = requests.Session
class CountingAsyncOpenAI(openai.AsyncOpenAI):
def __init__(self, **kwargs: Any) -> None:
counts["async_openai"] += 1
super().__init__(**kwargs)
def counting_session() -> requests.Session:
counts["session"] += 1
return real_session()
# layout, image prompt, verify verdict, chat reply — in call order.
completions_create = AsyncMock(
side_effect=[
_make_response("square", None),
_make_response("a detailed prompt", None),
_make_response("PASS", None),
_make_response("a chat reply", None),
]
)
image_response = MagicMock()
image_response.data = [MagicMock(b64_json=base64.b64encode(b"fake image").decode())]
images_generate = AsyncMock(return_value=image_response)
registry = MagicMock()
registry.to_openai_tools.return_value = []
db = ChatDatabase(db_path=temp_db_path)
with (
patch.object(openai, "AsyncOpenAI", CountingAsyncOpenAI),
patch.object(requests, "Session", counting_session),
patch.object(chat_client.chat.completions, "create", completions_create),
patch.object(gen_client.images, "generate", images_generate),
patch("vibe_bot.llm_client.embedding", return_value=[0.25] * 32),
):
asyncio.run(
ImageService(db, MagicMock()).generate(mock_ctx, message="a centaur")
)
asyncio.run(
ChatService(db, registry).handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
assert counts["async_openai"] == 0
assert counts["session"] == 0
assert llm_client.get_chat_client() is chat_client
assert llm_client.get_image_gen_client() is gen_client
assert llm_client.get_image_edit_client() is edit_client
def _make_response(content: str | None, tool_calls: list[object] | None) -> MagicMock:
"""Build a mock chat completion response with the given message fields."""
message = MagicMock()
message.content = content
message.tool_calls = tool_calls
return MagicMock(choices=[MagicMock(message=message)])
def test_chat_complete_skips_non_function_tool_call() -> None:
"""A tool call that is not of type 'function' is skipped, not executed."""
import asyncio
from vibe_bot.llm_client import chat_complete
called = {"n": 0}
def tool_executor(name: str, args: dict[str, str]) -> str:
called["n"] += 1
return f"executed:{name}"
custom_tool_call = MagicMock()
custom_tool_call.type = "custom"
mock_client = MagicMock()
mock_client.chat.completions.create = AsyncMock(
side_effect=[
_make_response(content=None, tool_calls=[custom_tool_call]),
_make_response(content="final answer", tool_calls=None),
]
)
with patch("vibe_bot.llm.chat.get_chat_client", return_value=mock_client):
result = asyncio.run(
chat_complete(
[{"role": "user", "content": "hi"}],
model="m",
max_tokens=10,
tool_executor=tool_executor,
)
)
assert result == "final answer"
assert called["n"] == 0
def test_tool_registry_dispatch_and_unknown_tool() -> None:
"""The registry renders schemas, dispatches known tools, and names unknowns."""
from vibe_bot.llm_client import ToolRegistry
def echo_tool(name: str, args: dict[str, str], **kwargs: object) -> str:
return f"echo:{args.get('text', '')}"
registry = ToolRegistry()
registry.register(
"echo",
"Echoes the text argument back.",
{"type": "object", "properties": {"text": {"type": "string"}}},
echo_tool,
)
tools = registry.to_openai_tools()
first = tools[0]
assert first["type"] == "function"
function_def = cast("dict[str, object]", first["function"])
assert function_def["name"] == "echo"
assert function_def["description"] == "Echoes the text argument back."
assert registry.execute("echo", {"text": "hi"}) == "echo:hi"
assert registry.execute("does_not_exist", {}) == "Unknown tool: does_not_exist"
vibe_bot/tests/test_logging.py
+186
View File
@@ -0,0 +1,186 @@
"""No-content logging tests.
A recognizable secret string is seeded into message content, image prompts,
and bot personalities; the logs captured at DEBUG must never contain it,
while metadata (bot name, user id, message id, counts) must.
"""
from __future__ import annotations
import asyncio
import base64
import logging
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.image_service import ImageService
from vibe_bot.tests._helpers import invoke
SECRET = "SECRET-DO-NOT-LOG-12345"
def _registry() -> MagicMock:
"""A mock ToolRegistry."""
reg = MagicMock()
reg.to_openai_tools.return_value = []
reg.execute.return_value = "tool result"
return reg
def _assert_no_secret(
caplog: pytest.LogCaptureFixture,
*metadata: str,
) -> None:
"""The secret must be absent from every record; the metadata must be present."""
for record in caplog.records:
assert (
SECRET not in record.getMessage()
), f"Secret leaked in log record: {record.getMessage()!r}"
assert SECRET not in caplog.text
for token in metadata:
assert token in caplog.text, f"Expected metadata {token!r} missing from logs"
def test_chat_service_logs_metadata_not_content(
caplog: pytest.LogCaptureFixture,
chat_db: Any,
mock_ctx: MagicMock,
) -> None:
"""A full chat turn (RAG + persist + reply) never logs the message content."""
caplog.set_level(logging.DEBUG)
svc = ChatService(chat_db, _registry())
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="A perfectly fine response."),
):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message=f"please remember {SECRET} forever",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
_assert_no_secret(caplog, "alfred", "12345")
def test_add_message_logs_metadata_not_content(
caplog: pytest.LogCaptureFixture,
chat_db: Any,
) -> None:
"""add_message (with embedding) never logs the stored content."""
caplog.set_level(logging.DEBUG)
assert chat_db.add_message(
message_id="msg-1",
user_id="12345",
username="testuser",
content=f"User: tell me about {SECRET}",
bot_name="alfred",
channel_id="channel-1",
guild_id="guild-1",
)
assert chat_db.add_message(
message_id="msg-1_response",
user_id="bot-123",
username="test-bot",
content=f"the bot knew {SECRET}",
bot_name="alfred",
channel_id="channel-1",
guild_id="guild-1",
role="assistant",
embed=False,
)
_assert_no_secret(caplog, "msg-1", "12345")
def test_history_lookups_log_metadata_not_content(
caplog: pytest.LogCaptureFixture,
chat_db: Any,
) -> None:
"""get_user_history / get_bot_history never log message or response content."""
caplog.set_level(logging.DEBUG)
chat_db.add_message(
message_id="msg-2",
user_id="12345",
username="testuser",
content=f"User: what is {SECRET}",
bot_name="alfred",
channel_id="channel-1",
guild_id="guild-1",
)
chat_db.add_message(
message_id="msg-2_response",
user_id="bot-123",
username="test-bot",
content=f"it is {SECRET}, clearly",
bot_name="alfred",
channel_id="channel-1",
guild_id="guild-1",
role="assistant",
embed=False,
)
user_history = chat_db.get_user_history("12345", limit=5)
bot_history = chat_db.get_bot_history("alfred", limit=5)
assert len(user_history) == 1
assert len(bot_history) == 1
_assert_no_secret(caplog, "msg-2", "alfred")
def test_doodlebob_logs_metadata_not_content(
caplog: pytest.LogCaptureFixture,
mock_ctx: MagicMock,
) -> None:
"""A doodlebob generation never logs the prompt or the derived image prompt."""
caplog.set_level(logging.DEBUG)
db = MagicMock()
db.get_image_generation_time_estimate.return_value = 12.0
svc = ImageService(db, MagicMock())
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(
side_effect=[
"portrait",
f"a very detailed painting of {SECRET}",
"pass",
]
),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=base64.b64encode(b"img").decode()),
),
):
asyncio.run(svc.generate(mock_ctx, message=f"draw {SECRET} on the moon"))
_assert_no_secret(caplog, "Doodlebob", "12345")
def test_custom_bot_creation_logs_metadata_not_personality(
caplog: pytest.LogCaptureFixture,
bot: Any,
mock_ctx: MagicMock,
) -> None:
"""!custom-bot never logs the personality text."""
caplog.set_level(logging.DEBUG)
invoke(
bot,
"custom-bot",
mock_ctx,
"secretbot",
personality=f"a butler who knows {SECRET}",
)
_assert_no_secret(caplog, "secretbot", "12345")
vibe_bot/tests/test_main.py
-1321
View File
File diff suppressed because it is too large Load Diff
vibe_bot/tests/test_prompts.py
+176
View File
@@ -0,0 +1,176 @@
"""Tests for the prompts module (prompt constants, layout parsing, user info)."""
from __future__ import annotations
from datetime import UTC, datetime
from unittest.mock import MagicMock
import pytest
from vibe_bot.config import (
IMAGE_GEN_SIZE_LANDSCAPE,
IMAGE_GEN_SIZE_PORTRAIT,
IMAGE_GEN_SIZE_SQUARE,
)
from vibe_bot.prompts import (
IMAGE_PROMPT_SYSTEM_PROMPT_TEMPLATE,
LAYOUT_SIZES,
RESPONSE_LENGTH_HINT,
build_system_prompt,
get_user_info,
parse_image_layout,
)
@pytest.fixture
def mock_ctx() -> MagicMock:
"""A minimal Discord user for get_user_info."""
author = MagicMock()
author.name = "testuser"
author.id = "12345"
author.global_name = "Test User"
author.nick = "tester"
author.top_role.name = "@everyone"
author.activities = []
author.joined_at = None
author.created_at = None
return author
@pytest.fixture
def mock_author_with_member_data() -> MagicMock:
"""A Discord user with full member data (role + activity + timestamps)."""
author = MagicMock()
author.name = "testuser"
author.id = "12345"
author.global_name = "Test User"
author.nick = "tester"
author.top_role.name = "Admin"
activity = MagicMock()
activity.name = "Chess"
author.activities = [activity]
author.joined_at = datetime(2024, 1, 15, tzinfo=UTC)
author.created_at = datetime(2023, 6, 1, tzinfo=UTC)
return author
def test_build_system_prompt_contains_personality_hint_and_user_info() -> None:
"""build_system_prompt assembles personality + length hint + user info block."""
result = build_system_prompt("you are a butler", "Username: alice")
assert result.startswith("you are a butler")
assert RESPONSE_LENGTH_HINT in result
assert "User Information:\nUsername: alice" in result
def test_layout_sizes_map_to_config() -> None:
"""LAYOUT_SIZES maps each layout to its configured canvas size."""
assert LAYOUT_SIZES["portrait"] == IMAGE_GEN_SIZE_PORTRAIT
assert LAYOUT_SIZES["landscape"] == IMAGE_GEN_SIZE_LANDSCAPE
assert LAYOUT_SIZES["square"] == IMAGE_GEN_SIZE_SQUARE
@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:
"""Valid LLM layout responses parse to the right layout."""
assert 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:
"""Empty or malformed responses fall back to square."""
assert 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."""
prompt = 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
assert "exactly as written" in lowered
assert "fountain pen wearing pants" in lowered
assert "centaur" in lowered
assert "not a person riding a horse" in lowered
def test_get_user_info_minimal(mock_ctx: MagicMock) -> None:
"""get_user_info with minimal member data includes the core identity lines."""
result = get_user_info(mock_ctx)
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_author_with_member_data: MagicMock,
) -> None:
"""get_user_info with full member data includes roles, activity, timestamps."""
result = get_user_info(mock_author_with_member_data)
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:
"""Optional fields are omitted when they are empty."""
mock_ctx.global_name = None
mock_ctx.nick = None
mock_ctx.top_role.name = "@everyone"
mock_ctx.activities = []
result = get_user_info(mock_ctx)
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_author_with_member_data: MagicMock,
) -> None:
"""Top role is included when it is not @everyone."""
result = get_user_info(mock_author_with_member_data)
assert "Top Role: Admin" in result
def test_get_user_info_no_activities(mock_ctx: MagicMock) -> None:
"""The activities line is omitted when there are none."""
mock_ctx.activities = []
result = get_user_info(mock_ctx)
assert "Activities:" not in result
vibe_bot/tests/test_services.py
+953
View File
@@ -0,0 +1,953 @@
"""Service-layer tests: chat, image, speech, and conversation services.
These exercise the LLM-backed logic directly (constructing each service with
mock dependencies) rather than going through the thin Discord command wrappers.
"""
from __future__ import annotations
import asyncio
import base64
from io import BytesIO
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import requests
from vibe_bot.config import TTS_VOICE
from vibe_bot.services.chat_service import ChatService
from vibe_bot.services.conversation_service import (
MAX_TOPIC_LENGTH,
ConversationService,
flip_counter,
)
from vibe_bot.services.image_service import (
MAX_IMAGE_DOWNLOAD_BYTES,
MAX_IMAGE_PROMPT_LENGTH,
ImageService,
_allowed_image_url,
_download_image_bytes,
select_image_layout,
verify_image_prompt,
)
from vibe_bot.services.speech_service import (
MAX_SPEAK_LENGTH,
SpeechService,
parse_voice_flag,
)
@pytest.fixture
def mock_ctx() -> MagicMock:
"""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
def _file_factory() -> MagicMock:
"""A File factory that records (buffer, filename) as a tuple."""
factory = MagicMock()
def make_file(buf: BytesIO, name: str) -> tuple[str, BytesIO, str]:
return ("FILE", buf, name)
factory.side_effect = make_file
return factory
def _sent_texts(ctx: MagicMock) -> list[str]:
"""All positional text messages sent through ctx.send."""
return [c.args[0] for c in ctx.send.call_args_list if c.args]
def _registry() -> MagicMock:
"""A mock ToolRegistry."""
reg = MagicMock()
reg.to_openai_tools.return_value = []
reg.execute.return_value = "tool result"
return reg
def _fake_bot(name: str) -> tuple[str, str, str, str]:
"""A stand-in custom bot tuple for manager.get_custom_bot."""
return (name, "a personality", "user-123", "2024-01-01")
# ---------------------------------------------------------------------------
# ChatService
# ---------------------------------------------------------------------------
def test_chat_success(mock_ctx: MagicMock) -> None:
"""A normal turn persists the exchange and sends the reply."""
db = MagicMock()
db.get_conversation_context.return_value = []
svc = ChatService(db, _registry())
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="This is a bot response"),
):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
db.add_message.assert_called()
assert mock_ctx.send.call_count >= 2
def test_chat_turn_embedding_budget(
mock_ctx: MagicMock,
temp_db_path: str,
) -> None:
"""One chat turn embeds exactly twice: the RAG query and the user row.
The assistant row is persisted with embed=False, so it costs no
embedding call and stores no embedding row.
"""
import sqlite3
from vibe_bot.database import ChatDatabase
db = ChatDatabase(db_path=temp_db_path)
svc = ChatService(db, _registry())
with (
patch(
"vibe_bot.llm_client.embedding",
return_value=[0.25] * 32,
) as mock_embedding,
patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="This is a bot response"),
),
):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
assert mock_embedding.call_count == 2
conn = sqlite3.connect(temp_db_path)
embedding_rows = conn.execute("SELECT COUNT(*) FROM message_embeddings").fetchone()
conn.close()
assert embedding_rows[0] == 1
def test_chat_error(mock_ctx: MagicMock) -> None:
"""An LLM error surfaces a friendly message."""
db = MagicMock()
db.get_conversation_context.return_value = []
svc = ChatService(db, _registry())
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(side_effect=Exception("API error")),
):
asyncio.run(
svc.handle(
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()
db.add_message.assert_not_called()
def test_chat_long_response_chunked(mock_ctx: MagicMock) -> None:
"""Long responses are split into multiple sends."""
db = MagicMock()
db.get_conversation_context.return_value = []
svc = ChatService(db, _registry())
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="x" * 2500),
):
asyncio.run(
svc.handle(
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_chat_includes_user_info(mock_ctx: MagicMock) -> None:
"""The system prompt sent to the LLM includes the requester's info."""
db = MagicMock()
db.get_conversation_context.return_value = []
svc = ChatService(db, _registry())
mock_llm = AsyncMock(return_value="resp")
with patch("vibe_bot.llm_client.chat_completion_with_tools", new=mock_llm):
asyncio.run(
svc.handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
system_prompt = mock_llm.call_args.kwargs["system_prompt"]
assert "testuser" in system_prompt
def test_chat_with_context_and_tools(mock_ctx: MagicMock) -> None:
"""Prior RAG context is prepended and tool calls reach the registry."""
db = MagicMock()
db.get_conversation_context.return_value = [
{"role": "user", "content": "old question"},
{"role": "assistant", "content": "old answer"},
]
registry = _registry()
captured: dict[str, Any] = {}
async def fake_llm(**kwargs: Any) -> str:
captured.update(kwargs)
kwargs["tool_executor"]("get_channel_members", {})
await kwargs["tool_call_notifier"]("get_channel_members", {})
return "resp"
with patch("vibe_bot.llm_client.chat_completion_with_tools", new=fake_llm):
asyncio.run(
ChatService(db, registry).handle(
mock_ctx,
bot_name="alfred",
message="hello",
system_prompt="you are a butler",
response_prefix="alfred response",
)
)
prompts = captured["prompts"]
assert prompts[0] == {"role": "user", "content": "old question"}
assert prompts[-1] == {"role": "user", "content": "hello"}
registry.execute.assert_called_once_with(
"get_channel_members", {}, channel=mock_ctx.channel
)
assert any("looking at the channel members" in t for t in _sent_texts(mock_ctx))
# ---------------------------------------------------------------------------
# SpeechService
# ---------------------------------------------------------------------------
def _speech_service(
tts: MagicMock | None,
manager: MagicMock | None = None,
make_file: MagicMock | None = None,
) -> SpeechService:
return SpeechService(
MagicMock(), manager or MagicMock(), tts, make_file or _file_factory()
)
def test_speak_tts_not_initialized(mock_ctx: MagicMock) -> None:
"""No TTS engine means a clear error, no LLM or TTS calls."""
svc = _speech_service(None)
asyncio.run(svc.speak(mock_ctx, message="hello world"))
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) -> None:
"""Empty text is rejected before any TTS work."""
svc = _speech_service(MagicMock())
asyncio.run(svc.speak(mock_ctx, message=""))
call_args = mock_ctx.send.call_args[0][0]
assert "Please provide text" in call_args
def test_speak_too_long(mock_ctx: MagicMock) -> None:
"""Oversized text is rejected without calling the TTS engine."""
tts = MagicMock()
svc = _speech_service(tts)
asyncio.run(svc.speak(mock_ctx, message="a" * (MAX_SPEAK_LENGTH + 1)))
tts.generate_audio.assert_not_called()
call_args = mock_ctx.send.call_args[0][0]
assert "Text too long to speak" in call_args
def test_speak_partial_audio_warns(mock_ctx: MagicMock) -> None:
"""Partial audio triggers a warning line."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(
audio=MagicMock(), partial=True, failed_chunks=1
)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
asyncio.run(svc.speak(mock_ctx, message="hello world"))
assert any("audio may be incomplete" in t for t in _sent_texts(mock_ctx))
def test_speak_plain_text(mock_ctx: MagicMock) -> None:
"""Plain text is spoken and the audio file is sent."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
asyncio.run(svc.speak(mock_ctx, message="hello world"))
tts.generate_audio.assert_called_once()
assert mock_ctx.send.call_count >= 2
def test_speak_with_custom_bot(mock_ctx: MagicMock) -> None:
"""A bot prefix routes through the LLM, then speaks the response."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = [
("alfred", "british butler", "user-123"),
]
manager.get_custom_bot.return_value = (
"alfred",
"british butler",
"user-123",
"2024-01-01",
)
svc = _speech_service(tts, manager=manager)
with patch(
"vibe_bot.llm_client.chat_completion_with_tools",
new=AsyncMock(return_value="The time is 3pm"),
):
asyncio.run(svc.speak(mock_ctx, message="alfred what time is it"))
tts.generate_audio.assert_called_once()
assert any("**alfred**:" in t for t in _sent_texts(mock_ctx))
def test_speak_uses_requested_voice(mock_ctx: MagicMock) -> None:
"""A trailing --voice flag selects that voice for the TTS call."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
asyncio.run(svc.speak(mock_ctx, message="hello world --voice af_bella"))
assert tts.generate_audio.call_args.kwargs["voice"] == "af_bella"
def test_speak_mid_text_voice_flag_spoken_verbatim(mock_ctx: MagicMock) -> None:
"""A --voice mid-message is preserved as speech and the default voice is used."""
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
message = "hello --voice af_bella world"
asyncio.run(svc.speak(mock_ctx, message=message))
assert tts.generate_audio.call_args.args[0] == message
assert tts.generate_audio.call_args.kwargs["voice"] == TTS_VOICE
def test_speak_unknown_voice(mock_ctx: MagicMock) -> None:
"""An unknown voice is rejected before any TTS call."""
tts = MagicMock()
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
asyncio.run(svc.speak(mock_ctx, message="hello --voice not_a_real_voice"))
tts.generate_audio.assert_not_called()
call_args = mock_ctx.send.call_args[0][0]
assert "Unknown voice" in call_args
def test_speak_language_lookup_uses_precomputed_dict(
mock_ctx: MagicMock,
) -> None:
"""The speak hot path resolves the language via VOICE_LANGUAGES.get().
The dict is built once at import from VOICES_LIST (covering every
catalog voice) and the per-speak lookup is a single dict get — no
per-call scan of the category list.
"""
from vibe_bot.config import VOICES_LIST
from vibe_bot.services import speech_service
class LanguageLookupSpy:
"""Counts .get() lookups on the voice->language mapping."""
def __init__(self, data: dict[str, str]) -> None:
self.data = data
self.lookups = 0
def get(self, key: str, default: str | None = None) -> str | None:
self.lookups += 1
return self.data.get(key, default)
def __contains__(self, key: object) -> bool:
return key in self.data
counting = LanguageLookupSpy(speech_service.VOICE_LANGUAGES)
assert set(counting.data) == {
voice for category in VOICES_LIST.values() for voice in category["voices"]
}
tts = MagicMock()
tts.generate_audio.return_value = MagicMock(audio=MagicMock(), partial=False)
manager = MagicMock()
manager.list_custom_bots.return_value = []
svc = _speech_service(tts, manager=manager)
with patch.object(speech_service, "VOICE_LANGUAGES", counting):
asyncio.run(svc.speak(mock_ctx, message="hello --voice bf_alice"))
assert counting.lookups == 1
assert tts.generate_audio.call_args.kwargs["lang"] == "en-gb"
# ---------------------------------------------------------------------------
# ConversationService
# ---------------------------------------------------------------------------
def _conversation_service(
manager: MagicMock | None = None,
) -> ConversationService:
return ConversationService(manager or MagicMock())
def test_flip_counter() -> None:
"""flip_counter toggles between 0 and 1."""
assert flip_counter(0) == 1
assert flip_counter(1) == 0
def test_talkforme_topic_too_long(mock_ctx: MagicMock) -> None:
"""Oversized topics are rejected before any LLM call."""
svc = _conversation_service()
asyncio.run(svc.run(mock_ctx, "a", "b", "3", "x" * (MAX_TOPIC_LENGTH + 1)))
call_args = mock_ctx.send.call_args[0][0]
assert "Topic too long" in call_args
def test_talkforme_bot1_not_found(mock_ctx: MagicMock) -> None:
"""A missing first bot is reported and the run stops."""
manager = MagicMock()
manager.get_custom_bot.return_value = None
svc = _conversation_service(manager=manager)
asyncio.run(svc.run(mock_ctx, "ghost", "alfred", "3", "cats"))
call_args = mock_ctx.send.call_args[0][0]
assert "ghost is not a real bot" in call_args
def test_talkforme_invalid_limit(mock_ctx: MagicMock) -> None:
"""A non-integer limit is rejected after both bots are found."""
manager = MagicMock()
manager.get_custom_bot.side_effect = _fake_bot
svc = _conversation_service(manager=manager)
asyncio.run(svc.run(mock_ctx, "a", "b", "abc", "cats"))
call_args = mock_ctx.send.call_args[0][0]
assert "Message limit must be an integer" in call_args
def test_talkforme_first_reply_chunked(mock_ctx: MagicMock) -> None:
"""Long first replies are sent in multiple chunks."""
manager = MagicMock()
manager.get_custom_bot.side_effect = _fake_bot
svc = _conversation_service(manager=manager)
with patch(
"vibe_bot.llm_client.chat_completion_with_history",
new=AsyncMock(return_value="y" * 2500),
):
asyncio.run(svc.run(mock_ctx, "a", "b", "1", "cats"))
assert mock_ctx.send.call_count >= 3
# ---------------------------------------------------------------------------
# ImageService (doodlebob / retcon)
# ---------------------------------------------------------------------------
def _image_service(db: MagicMock | None = None) -> ImageService:
return ImageService(db or MagicMock(), _file_factory())
def test_doodlebob_prompt_too_long(mock_ctx: MagicMock) -> None:
"""Oversized prompts are rejected before any LLM call."""
svc = _image_service()
asyncio.run(svc.generate(mock_ctx, message="a" * (MAX_IMAGE_PROMPT_LENGTH + 1)))
call_args = mock_ctx.send.call_args[0][0]
assert "Prompt too long" in call_args
def test_doodlebob_generate_success(mock_ctx: MagicMock) -> None:
"""A full generate flow ends with an image file and a completion line."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
b64 = base64.b64encode(b"fake image").decode()
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=b64),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur in a field"))
assert db.record_image_generation_time.called
assert any("Strike complete" in t for t in _sent_texts(mock_ctx))
def test_doodlebob_failed_generation(mock_ctx: MagicMock) -> None:
"""An empty image-generation response is reported as a failure."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=""),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
assert any("Failed to generate image" in t for t in _sent_texts(mock_ctx))
assert not db.record_image_generation_time.called
def test_doodlebob_reports_estimate(mock_ctx: MagicMock) -> None:
"""A prior-history estimate produces a Drone ETA line."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = 12.5
svc = _image_service(db)
b64 = base64.b64encode(b"fake image").decode()
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=b64),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
assert any("Drone ETA" in t for t in _sent_texts(mock_ctx))
def test_doodlebob_no_estimate_without_history(mock_ctx: MagicMock) -> None:
"""No estimate means no Drone ETA line."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
b64 = base64.b64encode(b"fake image").decode()
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value=b64),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
assert not any("Drone ETA" in t for t in _sent_texts(mock_ctx))
def test_doodlebob_empty_prompt_stops(mock_ctx: MagicMock) -> None:
"""An empty image-prompt response stops the flow without generating."""
db = MagicMock()
svc = _image_service(db)
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value=""),
),
patch("vibe_bot.llm_client.image_generation") as mock_gen,
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
mock_gen.assert_not_called()
assert not db.record_image_generation_time.called
def test_doodlebob_decode_failure(mock_ctx: MagicMock) -> None:
"""Invalid base64 from the generation API is reported as a failure."""
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
with (
patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(side_effect=["square", "a detailed prompt", "PASS"]),
),
patch(
"vibe_bot.llm_client.image_generation",
new=AsyncMock(return_value="abcde!!!"),
),
):
asyncio.run(svc.generate(mock_ctx, message="a centaur"))
assert any(
"Failed to process the generated image" in t for t in _sent_texts(mock_ctx)
)
@pytest.mark.parametrize(
("response", "expected"),
[
("portrait", "portrait"),
("landscape", "landscape"),
("square", "square"),
("PORTRAIT", "portrait"),
("I think landscape", "landscape"),
],
)
def test_select_image_layout_returns_parsed(
response: str,
expected: str,
) -> None:
"""select_image_layout parses the LLM's layout choice."""
with patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value=response),
):
result = asyncio.run(select_image_layout("a tall tree"))
assert result == expected
def test_select_image_layout_uses_minimal_token_budget() -> None:
"""Layout selection is a one-word answer, so max_tokens is 2."""
mock_llm = AsyncMock(return_value="square")
with patch("vibe_bot.llm_client.chat_completion_instruct", new=mock_llm):
assert asyncio.run(select_image_layout("a tall tree")) == "square"
assert mock_llm.call_args.kwargs["max_tokens"] == 2
def test_doodlebob_latency_within_budget(mock_ctx: MagicMock) -> None:
"""End-to-end doodlebob latency with 50ms simulated per LLM/image call.
Hermetic latency figure: four mocked calls (layout, prompt, verify,
generate) at 50ms each must dominate the wall time; the overhead on
top of the simulated 200ms stays far below the budget.
"""
import time
db = MagicMock()
db.get_image_generation_time_estimate.return_value = None
svc = _image_service(db)
b64 = base64.b64encode(b"fake image").decode()
responses = ["square", "a detailed prompt", "PASS"]
async def slow_instruct(**_kwargs: Any) -> str:
await asyncio.sleep(0.05)
return responses.pop(0)
async def slow_generate(**_kwargs: Any) -> str:
await asyncio.sleep(0.05)
return b64
with (
patch("vibe_bot.llm_client.chat_completion_instruct", new=slow_instruct),
patch("vibe_bot.llm_client.image_generation", new=slow_generate),
):
start = time.monotonic()
asyncio.run(svc.generate(mock_ctx, message="a centaur in a field"))
elapsed = time.monotonic() - start
assert elapsed >= 0.2
assert elapsed < 5.0
assert any("Strike complete" in t for t in _sent_texts(mock_ctx))
def test_verify_image_prompt_pass_keeps_prompt() -> None:
"""A PASS verdict keeps the original prompt unchanged."""
with patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value="PASS"),
):
result = asyncio.run(verify_image_prompt("a centaur", "a detailed prompt"))
assert result == "a detailed prompt"
def test_verify_image_prompt_correction_replaces() -> None:
"""A non-passing verdict is used as the corrected prompt."""
correction = "a rewritten prompt that is definitely long enough to be a fix"
with patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value=correction),
):
result = asyncio.run(verify_image_prompt("a centaur", "a detailed prompt"))
assert result == correction
def test_verify_image_prompt_empty_falls_back() -> None:
"""An empty verdict falls back to the original prompt."""
with patch(
"vibe_bot.llm_client.chat_completion_instruct",
new=AsyncMock(return_value=""),
):
result = asyncio.run(verify_image_prompt("a centaur", "a detailed prompt"))
assert result == "a detailed prompt"
def test_retcon_no_attachments(mock_ctx: MagicMock) -> None:
"""retcon with no attachments asks the user to attach an image."""
svc = _image_service()
mock_ctx.message.attachments = []
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
call_args = mock_ctx.send.call_args[0][0]
assert "Please attach an image" in call_args
def test_retcon_rejected_url_not_downloaded(mock_ctx: MagicMock) -> None:
"""A non-Discord attachment URL is refused before any download happens."""
svc = _image_service()
attachment = MagicMock()
attachment.url = "https://evil.example.com/img.png"
mock_ctx.message.attachments = [attachment]
mock_edit = AsyncMock(return_value="")
with (
patch("vibe_bot.llm_client.image_edit", new=mock_edit),
patch("requests.get") as mock_get,
):
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
mock_get.assert_not_called()
mock_edit.assert_not_called()
assert any("Please attach an image" in t for t in _sent_texts(mock_ctx))
def test_retcon_prompt_too_long(mock_ctx: MagicMock) -> None:
"""Oversized retcon prompts are rejected before any download."""
svc = _image_service()
asyncio.run(svc.edit(mock_ctx, message="a" * (MAX_IMAGE_PROMPT_LENGTH + 1)))
call_args = mock_ctx.send.call_args[0][0]
assert "Prompt too long" in call_args
def test_retcon_image_edit_empty(mock_ctx: MagicMock) -> None:
"""An empty edit response is reported as a failure."""
svc = _image_service()
attachment = MagicMock()
attachment.url = "https://cdn.discordapp.com/attachments/1/2/3/img.png"
mock_ctx.message.attachments = [attachment]
with (
patch(
"vibe_bot.services.image_service._download_image_bytes",
return_value=b"fake image bytes",
),
patch(
"vibe_bot.llm_client.image_edit",
new=AsyncMock(return_value=""),
),
):
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
call_args = mock_ctx.send.call_args[0][0]
assert "Failed to edit the image" in call_args
def test_retcon_success(mock_ctx: MagicMock) -> None:
"""A successful edit sends the edited image file."""
svc = _image_service()
attachment = MagicMock()
attachment.url = "https://cdn.discordapp.com/attachments/1/2/3/img.png"
mock_ctx.message.attachments = [attachment]
b64 = base64.b64encode(b"edited").decode()
with (
patch(
"vibe_bot.services.image_service._download_image_bytes",
return_value=b"fake image bytes",
),
patch(
"vibe_bot.llm_client.image_edit",
new=AsyncMock(return_value=b64),
),
):
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
assert any("Rewriting history" in t for t in _sent_texts(mock_ctx))
def test_retcon_edit_decode_failure(mock_ctx: MagicMock) -> None:
"""Invalid base64 from the edit API is reported as a processing failure."""
svc = _image_service()
attachment = MagicMock()
attachment.url = "https://cdn.discordapp.com/attachments/1/2/3/img.png"
mock_ctx.message.attachments = [attachment]
with (
patch(
"vibe_bot.services.image_service._download_image_bytes",
return_value=b"fake image bytes",
),
patch(
"vibe_bot.llm_client.image_edit",
new=AsyncMock(return_value="abcde!!!"),
),
):
asyncio.run(svc.edit(mock_ctx, message="make it blue"))
call_args = mock_ctx.send.call_args[0][0]
assert "Failed to process the edited image" in call_args
def test_allowed_image_url() -> None:
"""Only Discord CDN hosts are allowed for retcon downloads."""
assert _allowed_image_url("https://cdn.discordapp.com/a/b/c.png")
assert _allowed_image_url("https://media.discordapp.net/a/b/c.png")
assert not _allowed_image_url("https://example.com/a/b/c.png")
assert not _allowed_image_url("https://evilcdn.com/a/b/c.png")
assert not _allowed_image_url("http://[::1")
def test_download_image_bytes_success() -> None:
"""An allowed Discord URL is downloaded and its chunks joined."""
response = MagicMock()
response.headers = {}
response.iter_content.return_value = [b"abc", b"", b"def"]
response.raise_for_status.return_value = None
with patch(
"vibe_bot.services.image_service.requests.get",
return_value=response,
) as mock_get:
data = _download_image_bytes("https://cdn.discordapp.com/a/b/c.png")
mock_get.assert_called_once()
assert data == b"abcdef"
def test_download_image_bytes_request_failure() -> None:
"""A failing download returns None instead of raising."""
with patch(
"vibe_bot.services.image_service.requests.get",
side_effect=requests.RequestException("boom"),
):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") is None
def test_download_image_bytes_content_length_cap() -> None:
"""A declared Content-Length above the cap is refused before streaming."""
response = MagicMock()
response.headers = {"Content-Length": str(MAX_IMAGE_DOWNLOAD_BYTES + 1)}
response.iter_content = MagicMock()
with patch("vibe_bot.services.image_service.requests.get", return_value=response):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") is None
response.iter_content.assert_not_called()
def test_download_image_bytes_streaming_cap() -> None:
"""Streaming past the size cap is refused even without a Content-Length."""
response = MagicMock()
response.headers = {}
response.iter_content.return_value = [b"a" * (MAX_IMAGE_DOWNLOAD_BYTES + 1)]
with patch("vibe_bot.services.image_service.requests.get", return_value=response):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") is None
def test_download_image_bytes_bad_content_length() -> None:
"""A non-numeric Content-Length is ignored and streaming proceeds."""
response = MagicMock()
response.headers = {"Content-Length": "not-a-number"}
response.iter_content.return_value = [b"xyz"]
response.raise_for_status.return_value = None
with patch("vibe_bot.services.image_service.requests.get", return_value=response):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") == b"xyz"
def test_download_image_bytes_stream_error() -> None:
"""A mid-stream request error returns None instead of raising."""
response = MagicMock()
response.headers = {}
response.iter_content.side_effect = requests.RequestException("stream died")
with patch("vibe_bot.services.image_service.requests.get", return_value=response):
assert _download_image_bytes("https://cdn.discordapp.com/a/b/c.png") is None
# ---------------------------------------------------------------------------
# parse_voice_flag
# ---------------------------------------------------------------------------
def test_parse_voice_flag_no_flag() -> None:
"""Without a trailing flag the message is returned unchanged."""
assert parse_voice_flag("hello world") == ("hello world", None)
def test_parse_voice_flag_trailing() -> None:
"""A trailing --voice flag is split off."""
assert parse_voice_flag("hello world --voice af_bella") == (
"hello world",
"af_bella",
)
def test_parse_voice_flag_mid_text_preserved() -> None:
"""A --voice that is not at the end is treated as speech, not a flag."""
assert parse_voice_flag("hello --voice af_bella world") == (
"hello --voice af_bella world",
None,
)
def test_parse_voice_flag_missing_value() -> None:
"""A --voice with no value (or only trailing spaces) is not a flag."""
assert parse_voice_flag("hello --voice") == ("hello --voice", None)
assert parse_voice_flag("hello --voice ") == ("hello --voice ", None)
vibe_bot/tests/test_textutil.py
+109
View File
@@ -0,0 +1,109 @@
"""Tests for the textutil module."""
from __future__ import annotations
import pytest
from vibe_bot.textutil import split_message
def test_split_message_empty() -> None:
assert split_message("") == []
def test_split_message_short_text_single_chunk() -> None:
assert split_message("hello") == ["hello"]
def test_split_message_exact_limit() -> None:
text = "a" * 1900
chunks = split_message(text)
assert chunks == [text]
assert len(chunks) == 1
def test_split_message_just_over_limit_no_newline() -> None:
text = "a" * 1901
chunks = split_message(text)
assert "".join(chunks) == text
assert all(len(c) <= 1900 for c in chunks)
assert chunks == ["a" * 1900, "a"]
def test_split_message_no_newlines_long() -> None:
text = "x" * 5000
chunks = split_message(text)
assert "".join(chunks) == text
assert all(len(c) <= 1900 for c in chunks)
def test_split_message_respects_newlines() -> None:
# Many short lines spanning several chunks: boundaries fall on newlines.
text = "\n".join(f"line {i}" for i in range(1, 501))
chunks = split_message(text)
assert "".join(chunks) == text
assert all(len(c) <= 1900 for c in chunks)
assert len(chunks) > 1
# Every chunk except the last ends on a newline (split at a line boundary).
for chunk in chunks[:-1]:
assert chunk.endswith("\n")
def test_split_message_long_line_hard_split() -> None:
# A single line longer than the limit must be hard-split.
text = "a" * 5000 + "\n" + "short"
chunks = split_message(text)
assert "".join(chunks) == text
assert all(len(c) <= 1900 for c in chunks)
def test_split_message_emoji_round_trip() -> None:
# Multi-codepoint emoji at split boundaries must not corrupt on join.
text = "hello 🌍 world 👨‍👩‍👧‍👦 end " * 500
for limit in (1, 10, 100, 1900):
chunks = split_message(text, limit)
assert "".join(chunks) == text
assert all(len(c) <= limit for c in chunks)
def test_split_message_property_round_trip_corpus() -> None:
corpus = [
"",
"a",
"a" * 1900,
"a" * 1901,
"line1\nline2\nline3",
"para one\n\npara two\n\npara three",
"no newline " * 1000,
"emoji 🎉 " * 1000,
"👨‍👩‍👧‍👦" * 2000,
"mixed 🌍 text and 日本語 and emoji 🎊 here",
# NFD combining marks (e + U+0301, a + U+0308) straddle split points.
"cafe\u0301 " * 1000,
"a\u0308\u0301 b\u0327\u0301 c\u0308 " * 1000,
# Code spans with backticks and spaces.
"`inline code` and `x = 1` spans " * 500,
"``double backtick`` `single` " * 500,
]
for text in corpus:
for limit in (1, 5, 100, 1900):
chunks = split_message(text, limit)
assert "".join(chunks) == text, f"round-trip failed for limit={limit}"
assert all(len(c) <= limit for c in chunks)
def test_split_message_exact_multiple_of_limit() -> None:
"""Input that is an exact multiple of the limit splits into full chunks."""
for limit in (1, 5, 100, 1900):
text = "z" * (limit * 4)
chunks = split_message(text, limit)
assert "".join(chunks) == text
assert len(chunks) == 4
assert all(len(c) == limit for c in chunks)
def test_split_message_limit_must_be_positive() -> None:
with pytest.raises(ValueError):
split_message("hello", 0)
with pytest.raises(ValueError):
split_message("hello", -5)
vibe_bot/tests/test_tts.py
+21 -11
View File
@@ -22,7 +22,7 @@ def test_tts_engine_init(mock_kokoro_tts: MagicMock) -> None:
def test_generate_audio(mock_kokoro_tts: MagicMock) -> None:
"""Test audio generation returns a BytesIO object."""
"""Test audio generation returns a full (non-partial) AudioResult."""
from io import BytesIO
from vibe_bot.tts import TTSEngine
@@ -30,9 +30,11 @@ def test_generate_audio(mock_kokoro_tts: MagicMock) -> None:
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("hello world this is a test")
assert isinstance(result, BytesIO)
result.seek(0)
data = result.read()
assert isinstance(result.audio, BytesIO)
assert result.partial is False
assert result.failed_chunks == 0
result.audio.seek(0)
data = result.audio.read()
assert len(data) > 0
@@ -57,7 +59,8 @@ def test_generate_audio_single_chunk(mock_kokoro_tts: MagicMock) -> None:
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("single chunk text")
assert isinstance(result, BytesIO)
assert isinstance(result.audio, BytesIO)
assert result.partial is False
mock_kokoro_tts["process_chunk_sequential"].assert_called_once()
@@ -77,7 +80,9 @@ def test_generate_audio_multiple_chunks(mock_kokoro_tts: MagicMock) -> None:
"this text is long enough to be split into multiple chunks",
)
assert isinstance(result, BytesIO)
assert isinstance(result.audio, BytesIO)
assert result.partial is False
assert result.failed_chunks == 0
assert mock_kokoro_tts["process_chunk_sequential"].call_count == 3
@@ -108,7 +113,12 @@ def test_generate_audio_chunk_failure(mock_kokoro_tts: MagicMock) -> None:
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("good chunk bad chunk another good")
assert isinstance(result, BytesIO)
# Audio is still produced for the good chunks, but flagged as partial.
assert isinstance(result.audio, BytesIO)
assert result.partial is True
assert result.failed_chunks == 1
result.audio.seek(0)
assert len(result.audio.read()) > 0
def test_generate_audio_all_chunks_fail(mock_kokoro_tts: MagicMock) -> None:
@@ -145,13 +155,13 @@ def test_generate_audio_returns_seekable(mock_kokoro_tts: MagicMock) -> None:
engine = TTSEngine("/tmp/test-model.onnx", "/tmp/test-voices.bin")
result = engine.generate_audio("hello world")
result.seek(0)
data = result.read()
result.audio.seek(0)
data = result.audio.read()
assert len(data) > 0
# Should be able to seek and read again
result.seek(0)
data2 = result.read()
result.audio.seek(0)
data2 = result.audio.read()
assert data == data2