361 lines
17 KiB
Python
361 lines
17 KiB
Python
import unittest
|
|
from datetime import datetime
|
|
from types import SimpleNamespace
|
|
from typing import Any, cast
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from aiogram.types import Chat, Message
|
|
|
|
import bot
|
|
import keyboards
|
|
|
|
|
|
class _StubState:
|
|
def __init__(self, mode="lobby"):
|
|
self.value = bot.MODE_STATES[mode].state
|
|
|
|
async def get_state(self):
|
|
return self.value
|
|
|
|
async def set_state(self, state):
|
|
self.value = state.state
|
|
|
|
async def clear(self):
|
|
self.value = None
|
|
|
|
|
|
class ButtonTextTests(unittest.TestCase):
|
|
def test_mode_buttons_map_to_switch_commands(self):
|
|
self.assertEqual(keyboards.resolve_button_text("🎓 Обучение"), "/training")
|
|
self.assertEqual(keyboards.resolve_button_text("✍️ Генерация"), "/generate")
|
|
self.assertEqual(keyboards.resolve_button_text("💳 Подписка"), "/subscription")
|
|
self.assertEqual(keyboards.resolve_button_text("📖 Помощь"), "/help")
|
|
|
|
def test_main_menu_button_maps_to_internal_transition(self):
|
|
self.assertEqual(
|
|
keyboards.resolve_button_text(keyboards.MAIN_MENU_BUTTON),
|
|
keyboards.MAIN_MENU_VALUE,
|
|
)
|
|
|
|
def test_regular_text_is_not_treated_as_button(self):
|
|
self.assertIsNone(keyboards.resolve_button_text("Напиши пост про запуск"))
|
|
self.assertIsNone(keyboards.resolve_button_text("привет"))
|
|
self.assertIsNone(keyboards.resolve_button_text(None))
|
|
|
|
def test_button_lookup_tolerates_whitespace(self):
|
|
self.assertEqual(keyboards.resolve_button_text(" 📖 Помощь "), "/help")
|
|
|
|
|
|
class KeyboardLayoutTests(unittest.TestCase):
|
|
def test_main_menu_only_shows_user_facing_modes(self):
|
|
labels = {b.text for row in keyboards.main_menu_keyboard().keyboard for b in row}
|
|
self.assertEqual(labels, set(keyboards.MODE_BUTTONS))
|
|
self.assertNotIn("🏠 Лобби", labels)
|
|
|
|
def test_each_mode_keyboard_shows_expected_actions(self):
|
|
for mode in bot.MODE_STATES:
|
|
labels = [b.text for row in keyboards.mode_keyboard(mode).keyboard for b in row]
|
|
expected = [keyboards.MAIN_MENU_BUTTON]
|
|
if mode == "training":
|
|
expected.append(keyboards.TRAINING_RESET_BUTTON)
|
|
self.assertEqual(labels, expected, mode)
|
|
|
|
def test_mode_entry_copy_uses_short_paragraphs(self):
|
|
self.assertIn("\n", bot.MODE_COPY["training"][1])
|
|
self.assertIn("\n", bot.MODE_COPY["generation"][1])
|
|
self.assertIn("\n", bot.MODE_COPY["subscription"][1])
|
|
|
|
def test_training_reset_button_routes_to_confirmation_action(self):
|
|
self.assertEqual(
|
|
keyboards.resolve_button_text(keyboards.TRAINING_RESET_BUTTON),
|
|
keyboards.TRAINING_RESET_VALUE,
|
|
)
|
|
|
|
def test_every_button_is_resolvable(self):
|
|
markups = [keyboards.main_menu_keyboard()]
|
|
markups.extend(keyboards.mode_keyboard(mode) for mode in bot.MODE_STATES)
|
|
for markup in markups:
|
|
for row in markup.keyboard:
|
|
for button in row:
|
|
self.assertIsNotNone(keyboards.resolve_button_text(button.text), button.text)
|
|
|
|
def test_keyboards_are_resize_keyboard(self):
|
|
self.assertTrue(keyboards.main_menu_keyboard().resize_keyboard)
|
|
for mode in bot.MODE_STATES:
|
|
self.assertTrue(keyboards.mode_keyboard(mode).resize_keyboard)
|
|
|
|
|
|
class _StubMessage:
|
|
def __init__(self, text=None, user_id=42, chat_id=555):
|
|
self.text = text
|
|
self.from_user = SimpleNamespace(id=user_id)
|
|
self.chat = SimpleNamespace(id=chat_id)
|
|
self.message_thread_id = None
|
|
self.bot = AsyncMock()
|
|
self.answers: list[tuple[str, dict]] = []
|
|
|
|
async def answer(self, text, **kwargs):
|
|
self.answers.append((text, kwargs))
|
|
|
|
|
|
class StyleCommandTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_style_without_tone_explains_how_to_set_it(self):
|
|
message = _StubMessage()
|
|
with patch.object(bot.asyncio, "to_thread", new=AsyncMock(side_effect=["", [], []])):
|
|
await bot.style_command(message, _StubState("training"))
|
|
self.assertIn("скриншот", message.answers[0][0])
|
|
|
|
async def test_style_shows_saved_tone(self):
|
|
message = _StubMessage()
|
|
with patch.object(bot.asyncio, "to_thread", new=AsyncMock(side_effect=["Короткие фразы", ["a", "b"], []])):
|
|
await bot.style_command(message, _StubState("training"))
|
|
self.assertIn("Короткие фразы", message.answers[0][0])
|
|
self.assertIn("Накоплено постов: 2", message.answers[1][0])
|
|
|
|
async def test_style_shows_saved_account_profile(self):
|
|
message = _StubMessage("/style")
|
|
state = _StubState("training")
|
|
with patch.object(bot.asyncio, "to_thread", new=AsyncMock(side_effect=["Профиль аккаунта", ["Пост"], ["Заметка"]])):
|
|
await bot.style_command(message, state)
|
|
self.assertIn("Профиль аккаунта", message.answers[0][0])
|
|
self.assertIn("Накоплено постов: 1", message.answers[1][0])
|
|
|
|
async def test_style_is_rejected_outside_training(self):
|
|
message = _StubMessage()
|
|
with patch.object(bot, "get_tone_of_voice") as get_tone:
|
|
await bot.style_command(message, _StubState("generation"))
|
|
get_tone.assert_not_called()
|
|
self.assertIn("только в режиме", message.answers[0][0])
|
|
|
|
|
|
class ModeRoutingTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_mode_state_is_stored_by_aiogram_memory_storage(self):
|
|
from aiogram.fsm.context import FSMContext
|
|
from aiogram.fsm.storage.base import StorageKey
|
|
from aiogram.fsm.storage.memory import MemoryStorage
|
|
|
|
storage = MemoryStorage()
|
|
state = FSMContext(storage, StorageKey(bot_id=1, chat_id=2, user_id=3))
|
|
await bot.enter_mode(_StubMessage(), state, "generation")
|
|
self.assertEqual(await bot.current_mode(state), "generation")
|
|
await storage.close()
|
|
|
|
async def test_mode_button_switches_from_any_mode(self):
|
|
message = _StubMessage("🎓 Обучение")
|
|
state = _StubState("generation")
|
|
with patch.object(bot, "notify_development_deploy", new=AsyncMock()):
|
|
await bot.handle_text(message, state)
|
|
self.assertEqual(await state.get_state(), bot.BotMode.training.state)
|
|
|
|
async def test_main_menu_button_returns_to_mode_selection(self):
|
|
message = _StubMessage(keyboards.MAIN_MENU_BUTTON)
|
|
state = _StubState("training")
|
|
with patch.object(bot, "notify_development_deploy", new=AsyncMock()):
|
|
await bot.handle_text(message, state)
|
|
self.assertEqual(await state.get_state(), bot.BotMode.lobby.state)
|
|
labels = [
|
|
button.text
|
|
for row in message.answers[0][1]["reply_markup"].keyboard
|
|
for button in row
|
|
]
|
|
self.assertEqual(set(labels), set(keyboards.MODE_BUTTONS))
|
|
|
|
async def test_slash_switch_command_works_from_another_mode(self):
|
|
message = _StubMessage("/training")
|
|
state = _StubState("subscription")
|
|
await bot.handle_command_text(message, state)
|
|
self.assertEqual(await state.get_state(), bot.BotMode.training.state)
|
|
|
|
async def test_help_button_switches_to_help_mode(self):
|
|
message = _StubMessage("📖 Помощь")
|
|
state = _StubState("lobby")
|
|
with (
|
|
patch.object(bot, "notify_development_deploy", new=AsyncMock()),
|
|
patch.object(bot, "send_rich_message", new=AsyncMock()),
|
|
):
|
|
await bot.handle_text(message, state)
|
|
self.assertEqual(await state.get_state(), bot.BotMode.help.state)
|
|
|
|
async def test_subscription_entry_shows_status_and_renew_button(self):
|
|
message = _StubMessage("💳 Подписка")
|
|
state = _StubState("lobby")
|
|
with (
|
|
patch.object(bot, "notify_development_deploy", new=AsyncMock()),
|
|
patch.object(bot.asyncio, "to_thread", new=AsyncMock(return_value=None)),
|
|
):
|
|
await bot.handle_text(message, state)
|
|
|
|
self.assertEqual(await state.get_state(), bot.BotMode.subscription.state)
|
|
self.assertEqual(len(message.answers), 2)
|
|
mode_keyboard = message.answers[0][1]["reply_markup"]
|
|
self.assertEqual(message.answers[0][1]["reply_markup"].keyboard[0][0].text, keyboards.MAIN_MENU_BUTTON)
|
|
self.assertIn("Подписка не активирована", message.answers[1][0])
|
|
renew_keyboard = message.answers[1][1]["reply_markup"]
|
|
self.assertEqual(
|
|
renew_keyboard.inline_keyboard[0][0].callback_data,
|
|
"subscription:renew",
|
|
)
|
|
|
|
async def test_renew_callback_sends_invoice_without_intermediate_message(self):
|
|
message_bot = AsyncMock()
|
|
message = Message(
|
|
message_id=91,
|
|
date=datetime.now(),
|
|
chat=Chat(id=555, type="private"),
|
|
).as_(message_bot)
|
|
callback = SimpleNamespace(
|
|
message=message,
|
|
from_user=SimpleNamespace(id=42),
|
|
answer=AsyncMock(),
|
|
)
|
|
with patch.object(bot, "_send_subscription_invoice", new=AsyncMock()) as send_invoice:
|
|
await bot.subscription_renew(callback, _StubState("subscription"))
|
|
|
|
callback.answer.assert_awaited_once_with()
|
|
send_invoice.assert_awaited_once_with(message, 42)
|
|
message_bot.assert_not_awaited()
|
|
|
|
async def test_idea_command_generates_only_in_generation_mode(self):
|
|
message = _StubMessage("/idea")
|
|
state = _StubState("generation")
|
|
with patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify:
|
|
await bot.handle_command_text(message, state)
|
|
classify.assert_awaited_once_with(message, "предложи идею для нового поста", state=state)
|
|
|
|
async def test_post_command_generates_only_in_generation_mode(self):
|
|
message = _StubMessage("/post")
|
|
state = _StubState("generation")
|
|
with patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify:
|
|
await bot.handle_command_text(message, state)
|
|
classify.assert_awaited_once_with(message, "Напиши пост", state=state)
|
|
|
|
async def test_training_text_is_saved_as_note_without_entering_generation(self):
|
|
message = _StubMessage("Кофейня в Калининграде для семей с детьми")
|
|
with patch.object(bot, "_rebuild_account_profile", new=AsyncMock()) as rebuild:
|
|
await bot.handle_text(message, _StubState("training"))
|
|
call = cast(Any, rebuild.await_args)
|
|
self.assertEqual(call.args, (42, message))
|
|
self.assertEqual(call.kwargs["new_note"], "Кофейня в Калининграде для семей с детьми")
|
|
self.assertIsInstance(call.kwargs["status"], bot.OperationStatus)
|
|
|
|
async def test_training_profile_rebuild_does_not_send_profile_details_when_status_exists(self):
|
|
status = AsyncMock()
|
|
message = _StubMessage()
|
|
with (
|
|
patch.object(bot.asyncio, "to_thread", new=AsyncMock(side_effect=["old", ["note"], ["post"], None])),
|
|
patch.object(bot, "create_tone_profile", new=AsyncMock(return_value="new profile")),
|
|
patch.object(bot, "send_rich_message", new=AsyncMock()) as rich_send,
|
|
):
|
|
await bot._rebuild_account_profile(42, cast(Any, message), new_note="note", status=status)
|
|
|
|
status.finish.assert_awaited_once_with("✅ Обучение завершено")
|
|
rich_send.assert_not_awaited()
|
|
self.assertEqual(message.answers, [])
|
|
|
|
async def test_plain_text_is_only_sent_to_generator_in_generation_mode(self):
|
|
message = _StubMessage("Напиши пост про запуск продукта")
|
|
state = _StubState("generation")
|
|
with (
|
|
patch.object(bot, "notify_development_deploy", new=AsyncMock()),
|
|
patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify,
|
|
):
|
|
await bot.handle_text(message, state)
|
|
classify.assert_awaited_once_with(message, "Напиши пост про запуск продукта", state=state)
|
|
|
|
async def test_plain_text_in_lobby_does_not_trigger_generation(self):
|
|
message = _StubMessage("Напиши пост")
|
|
state = _StubState("lobby")
|
|
with (
|
|
patch.object(bot, "notify_development_deploy", new=AsyncMock()),
|
|
patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify,
|
|
):
|
|
await bot.handle_text(message, state)
|
|
classify.assert_not_awaited()
|
|
self.assertIn("Выбери режим", message.answers[0][0])
|
|
|
|
async def test_style_command_is_rejected_outside_training(self):
|
|
message = _StubMessage("/style")
|
|
state = _StubState("generation")
|
|
with patch.object(bot, "style_command", new=AsyncMock()) as style:
|
|
await bot.handle_command_text(message, state)
|
|
style.assert_awaited_once_with(message, state)
|
|
|
|
async def test_clear_command_is_rejected_outside_training(self):
|
|
message = _StubMessage("/clear")
|
|
state = _StubState("generation")
|
|
with patch.object(bot, "clear_user_data") as clear_data:
|
|
await bot.handle_command_text(message, state)
|
|
clear_data.assert_not_called()
|
|
self.assertIn("только в режиме «Обучение»", message.answers[0][0])
|
|
|
|
async def test_subscription_command_is_rejected_outside_subscription_mode(self):
|
|
message = _StubMessage("/balance")
|
|
state = _StubState("generation")
|
|
with patch.object(bot, "subscription_row") as get_subscription:
|
|
await bot.handle_command_text(message, state)
|
|
get_subscription.assert_not_called()
|
|
self.assertIn("только в режиме «Подписка»", message.answers[0][0])
|
|
|
|
async def test_image_is_rejected_outside_training(self):
|
|
message = _StubMessage()
|
|
with patch.object(bot, "_handle_image_message", new=AsyncMock()) as process_image:
|
|
await bot.update_tone_from_image(message, AsyncMock(), _StubState("generation"))
|
|
process_image.assert_not_awaited()
|
|
self.assertIn("только в режиме «Обучение»", message.answers[0][0])
|
|
|
|
async def test_voice_is_rejected_outside_learning_and_generation(self):
|
|
message = _StubMessage()
|
|
with (
|
|
patch.object(bot, "notify_development_deploy", new=AsyncMock()) as notify,
|
|
patch.object(bot, "_process_voice", new=AsyncMock()) as process_voice,
|
|
):
|
|
await bot.transcribe_voice(message, AsyncMock(), _StubState("subscription"))
|
|
notify.assert_not_awaited()
|
|
process_voice.assert_not_awaited()
|
|
self.assertIn("«Обучение» и «Генерация»", message.answers[0][0])
|
|
|
|
async def test_voice_in_training_transcribes_then_updates_profile_not_generation(self):
|
|
message = SimpleNamespace(
|
|
from_user=SimpleNamespace(id=42),
|
|
voice=SimpleNamespace(file_size=4, mime_type="audio/ogg"),
|
|
audio=None,
|
|
answer=AsyncMock(),
|
|
chat=SimpleNamespace(id=555),
|
|
)
|
|
bot_client = SimpleNamespace(download=AsyncMock(side_effect=lambda file, destination: destination.write(b"ogg")))
|
|
state = _StubState("training")
|
|
with (
|
|
patch.object(bot, "notify_development_deploy", new=AsyncMock()),
|
|
patch.dict("os.environ", {"APINEX_API_KEY": "not-logged"}, clear=False),
|
|
patch.object(bot, "transcribe_audio", new=AsyncMock(return_value="Мы — кофейня в Калининграде")) as stt,
|
|
patch.object(bot, "_rebuild_account_profile", new=AsyncMock()) as rebuild,
|
|
patch.object(bot, "_classify_and_respond", new=AsyncMock()) as classify,
|
|
patch.object(bot, "typing_status", return_value=AsyncMock()),
|
|
):
|
|
await bot._process_voice(message, bot_client, state)
|
|
stt.assert_awaited_once()
|
|
call = cast(Any, rebuild.await_args)
|
|
self.assertEqual(call.args, (42, message))
|
|
self.assertEqual(call.kwargs["new_note"], "Мы — кофейня в Калининграде")
|
|
self.assertIsNone(call.kwargs["status"])
|
|
classify.assert_not_awaited()
|
|
|
|
|
|
|
|
class MenuPublishTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_commands_are_published_to_telegram(self):
|
|
fake_bot = AsyncMock()
|
|
await bot.set_main_menu(fake_bot)
|
|
names = [c.command for c in fake_bot.set_my_commands.await_args.kwargs["commands"]]
|
|
for expected in ("start", "training", "generate", "subscription", "help", "clear", "post"):
|
|
self.assertIn(expected, names)
|
|
|
|
async def test_publish_failure_does_not_crash_startup(self):
|
|
fake_bot = AsyncMock()
|
|
fake_bot.set_my_commands.side_effect = RuntimeError("boom")
|
|
await bot.set_main_menu(fake_bot)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|