250 lines
11 KiB
Python
250 lines
11 KiB
Python
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
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("🏠 Лобби"), "/lobby")
|
|
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_mode_actions_map_to_slash_commands(self):
|
|
expected = {
|
|
"🎨 Мой стиль": "/style",
|
|
"🗑 Сбросить данные": "/clear",
|
|
"✍️ Написать пост": "/post",
|
|
"💡 Идея": "/idea",
|
|
"👤 Статус подписки": "/balance",
|
|
"⭐ Оплатить подписку": "/subscribe",
|
|
}
|
|
for label, command in expected.items():
|
|
self.assertEqual(keyboards.resolve_button_text(label), command)
|
|
|
|
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_every_mode_keyboard_has_all_mode_switches(self):
|
|
expected = set(keyboards.MODE_BUTTONS)
|
|
for mode in bot.MODE_STATES:
|
|
labels = {b.text for row in keyboards.mode_keyboard(mode).keyboard for b in row}
|
|
self.assertTrue(expected.issubset(labels), mode)
|
|
|
|
def test_actions_are_limited_to_the_current_mode(self):
|
|
training = {b.text for row in keyboards.mode_keyboard("training").keyboard for b in row}
|
|
generation = {b.text for row in keyboards.mode_keyboard("generation").keyboard for b in row}
|
|
subscription = {b.text for row in keyboards.mode_keyboard("subscription").keyboard for b in row}
|
|
self.assertIn("🗑 Сбросить данные", training)
|
|
self.assertNotIn("🗑 Сбросить данные", generation)
|
|
self.assertIn("💡 Идея", generation)
|
|
self.assertNotIn("💡 Идея", training)
|
|
self.assertIn("⭐ Оплатить подписку", subscription)
|
|
self.assertNotIn("⭐ Оплатить подписку", generation)
|
|
|
|
def test_every_button_is_resolvable(self):
|
|
for mode in bot.MODE_STATES:
|
|
for row in keyboards.mode_keyboard(mode).keyboard:
|
|
for button in row:
|
|
self.assertIsNotNone(keyboards.resolve_button_text(button.text), button.text)
|
|
|
|
def test_keyboards_are_resize_keyboard(self):
|
|
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_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_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_idea_button_generates_only_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_post_button_generates_only_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_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_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("training"))
|
|
notify.assert_not_awaited()
|
|
process_voice.assert_not_awaited()
|
|
self.assertIn("только в режиме «Генерация»", message.answers[0][0])
|
|
|
|
|
|
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 ("lobby", "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()
|