202 lines
10 KiB
Python
202 lines
10 KiB
Python
|
import unittest
|
||
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
|
import asyncio
|
||
|
|
from telegram_helper import TelegramHelper, LogicResult
|
||
|
|
from telegram import Update
|
||
|
|
from telegram.ext import ContextTypes
|
||
|
|
import os
|
||
|
|
|
||
|
|
class TestTelegramHelper(unittest.IsolatedAsyncioTestCase):
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
self.mock_bot = AsyncMock()
|
||
|
|
self.telegram_helper = TelegramHelper(self.mock_bot)
|
||
|
|
|
||
|
|
async def test_start_logic(self):
|
||
|
|
self.mock_bot.start.return_value = None
|
||
|
|
result = await self.telegram_helper._start_logic()
|
||
|
|
self.mock_bot.start.assert_called_once()
|
||
|
|
self.assertEqual(result, "Hello! I'm your AI assistant. How can I help you today?")
|
||
|
|
|
||
|
|
async def test_clear_logic(self):
|
||
|
|
user_id = 123
|
||
|
|
self.mock_bot.clear_conversation_history.return_value = None
|
||
|
|
result = await self.telegram_helper._clear_logic(user_id)
|
||
|
|
self.mock_bot.clear_conversation_history.assert_called_once_with(user_id)
|
||
|
|
self.assertEqual(result, "Conversation history cleared. Let's start fresh!")
|
||
|
|
|
||
|
|
async def test_status_logic(self):
|
||
|
|
self.mock_bot.get_bot_status.return_value = "Bot is operational."
|
||
|
|
result = await self.telegram_helper._status_logic()
|
||
|
|
self.mock_bot.get_bot_status.assert_called_once()
|
||
|
|
self.assertEqual(result, "Bot is operational.")
|
||
|
|
|
||
|
|
async def test_switch_logic_supported(self):
|
||
|
|
self.mock_bot.switch_model.return_value = "Model switched successfully."
|
||
|
|
result = await self.telegram_helper._switch_logic()
|
||
|
|
self.mock_bot.switch_model.assert_called_once()
|
||
|
|
self.assertEqual(result, "Model switched successfully.")
|
||
|
|
|
||
|
|
async def test_switch_logic_not_supported(self):
|
||
|
|
# Remove switch_model from mock_bot to simulate not supported
|
||
|
|
delattr(self.mock_bot, 'switch_model')
|
||
|
|
result = await self.telegram_helper._switch_logic()
|
||
|
|
self.assertFalse(hasattr(self.mock_bot, 'switch_model'))
|
||
|
|
self.assertEqual(result, "Model switching is not supported for this bot.")
|
||
|
|
|
||
|
|
async def test_handle_message_logic_success(self):
|
||
|
|
user_id = 123
|
||
|
|
user_message = "Test message"
|
||
|
|
self.mock_bot.handle_message.return_value = "Bot response <think>thought</think>"
|
||
|
|
result = await self.telegram_helper._handle_message_logic(user_id, user_message)
|
||
|
|
self.mock_bot.handle_message.assert_called_once_with(user_id, user_message)
|
||
|
|
self.assertTrue(result["success"])
|
||
|
|
self.assertEqual(result["response_text"], f"Bot response {TelegramHelper.HTML_QUOTE_BLOCK_START}thought{TelegramHelper.HTML_QUOTE_BLOCK_END}")
|
||
|
|
self.assertIsNone(result["error_message"])
|
||
|
|
|
||
|
|
async def test_handle_message_logic_exception(self):
|
||
|
|
user_id = 123
|
||
|
|
user_message = "Test message"
|
||
|
|
self.mock_bot.handle_message.side_effect = Exception("Bot error")
|
||
|
|
result = await self.telegram_helper._handle_message_logic(user_id, user_message)
|
||
|
|
self.assertFalse(result["success"])
|
||
|
|
self.assertIsNone(result["response_text"])
|
||
|
|
self.assertEqual(result["error_message"], "Bot error")
|
||
|
|
|
||
|
|
async def test_abort_processing_logic(self):
|
||
|
|
user_id = 123
|
||
|
|
self.mock_bot.abort_processing.return_value = "Processing aborted."
|
||
|
|
result = await self.telegram_helper._abort_processing_logic(user_id)
|
||
|
|
self.mock_bot.abort_processing.assert_called_once_with(user_id)
|
||
|
|
self.assertEqual(result, "Processing aborted.")
|
||
|
|
|
||
|
|
@patch.object(TelegramHelper, '_start_logic', new_callable=AsyncMock)
|
||
|
|
async def test_start_command(self, mock_start_logic):
|
||
|
|
mock_update = MagicMock(spec=Update)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
mock_start_logic.return_value = "Hello!"
|
||
|
|
await self.telegram_helper.start(mock_update, mock_context)
|
||
|
|
mock_update.message.reply_text.assert_called_once_with("Hello!")
|
||
|
|
mock_start_logic.assert_called_once()
|
||
|
|
|
||
|
|
@patch.object(TelegramHelper, '_clear_logic', new_callable=AsyncMock)
|
||
|
|
async def test_clear_command(self, mock_clear_logic):
|
||
|
|
mock_update = MagicMock(spec=Update)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
mock_update.effective_user.id = 123
|
||
|
|
mock_clear_logic.return_value = "Cleared!"
|
||
|
|
await self.telegram_helper.clear(mock_update, mock_context)
|
||
|
|
mock_update.message.reply_text.assert_called_once_with("Cleared!")
|
||
|
|
mock_clear_logic.assert_called_once_with(123)
|
||
|
|
|
||
|
|
@patch.object(TelegramHelper, '_status_logic', new_callable=AsyncMock)
|
||
|
|
async def test_status_command(self, mock_status_logic):
|
||
|
|
mock_update = MagicMock(spec=Update)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
mock_status_logic.return_value = "Status OK"
|
||
|
|
await self.telegram_helper.status(mock_update, mock_context)
|
||
|
|
mock_update.message.reply_text.assert_called_once_with("Status OK")
|
||
|
|
mock_status_logic.assert_called_once()
|
||
|
|
|
||
|
|
@patch.object(TelegramHelper, '_switch_logic', new_callable=AsyncMock)
|
||
|
|
async def test_switch_command(self, mock_switch_logic):
|
||
|
|
mock_update = MagicMock(spec=Update)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
mock_switch_logic.return_value = "Switched"
|
||
|
|
await self.telegram_helper.switch(mock_update, mock_context)
|
||
|
|
mock_update.message.reply_text.assert_called_once_with("Switched")
|
||
|
|
mock_switch_logic.assert_called_once()
|
||
|
|
|
||
|
|
@patch.object(TelegramHelper, '_handle_message_logic', new_callable=AsyncMock)
|
||
|
|
async def test_handle_message_command_success_short_message(self, mock_handle_message_logic):
|
||
|
|
mock_update = MagicMock(spec=Update)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
mock_update.effective_user.id = 123
|
||
|
|
mock_update.message.text = "User message"
|
||
|
|
mock_update.effective_chat.id = 456
|
||
|
|
mock_update.message.reply_text.return_value = AsyncMock(message_id=789)
|
||
|
|
|
||
|
|
mock_handle_message_logic.return_value = LogicResult(success=True, response_text="Short response", error_message=None)
|
||
|
|
|
||
|
|
await self.telegram_helper.handle_message(mock_update, mock_context)
|
||
|
|
|
||
|
|
mock_update.message.reply_text.assert_any_call("Processing your request...", reply_markup=unittest.mock.ANY)
|
||
|
|
self.mock_bot.set_processing_status.assert_called_once_with(123, 789)
|
||
|
|
mock_handle_message_logic.assert_called_once_with(123, "User message")
|
||
|
|
mock_context.bot.delete_message.assert_called_once_with(chat_id=456, message_id=789)
|
||
|
|
self.mock_bot.clear_processing_status.assert_called_once_with(123)
|
||
|
|
mock_update.message.reply_text.assert_any_call("Short response")
|
||
|
|
|
||
|
|
@patch.object(TelegramHelper, '_handle_message_logic', new_callable=AsyncMock)
|
||
|
|
async def test_handle_message_command_success_long_message(self, mock_handle_message_logic):
|
||
|
|
mock_update = MagicMock(spec=Update)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
mock_update.effective_user.id = 123
|
||
|
|
mock_update.message.text = "User message"
|
||
|
|
mock_update.effective_chat.id = 456
|
||
|
|
mock_update.message.reply_text.return_value = AsyncMock(message_id=789)
|
||
|
|
|
||
|
|
long_response = "a" * 5000 # Longer than 4096
|
||
|
|
mock_handle_message_logic.return_value = LogicResult(success=True, response_text=long_response, error_message=None)
|
||
|
|
|
||
|
|
with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
|
||
|
|
await self.telegram_helper.handle_message(mock_update, mock_context)
|
||
|
|
self.assertEqual(mock_update.message.reply_text.call_count, 1 + 2) # Initial + two chunks
|
||
|
|
self.assertEqual(mock_sleep.call_count, 1) # One sleep for the second chunk
|
||
|
|
|
||
|
|
@patch.object(TelegramHelper, '_handle_message_logic', new_callable=AsyncMock)
|
||
|
|
async def test_handle_message_command_logic_failure(self, mock_handle_message_logic):
|
||
|
|
mock_update = MagicMock(spec=Update)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
mock_update.effective_user.id = 123
|
||
|
|
mock_update.message.text = "User message"
|
||
|
|
mock_update.effective_chat.id = 456
|
||
|
|
mock_update.message.reply_text.return_value = AsyncMock(message_id=789)
|
||
|
|
|
||
|
|
mock_handle_message_logic.return_value = LogicResult(success=False, response_text=None, error_message="Logic error")
|
||
|
|
|
||
|
|
await self.telegram_helper.handle_message(mock_update, mock_context)
|
||
|
|
|
||
|
|
mock_context.bot.delete_message.assert_called_once()
|
||
|
|
self.mock_bot.clear_processing_status.assert_called_once()
|
||
|
|
mock_update.message.reply_text.assert_any_call("Sorry, an error occurred while processing your request.")
|
||
|
|
|
||
|
|
@patch.object(TelegramHelper, '_abort_processing_logic', new_callable=AsyncMock)
|
||
|
|
async def test_abort_processing_callback(self, mock_abort_processing_logic):
|
||
|
|
mock_query = MagicMock()
|
||
|
|
mock_query.from_user.id = 123
|
||
|
|
mock_update = MagicMock(callback_query=mock_query)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
|
||
|
|
mock_abort_processing_logic.return_value = "Aborted!"
|
||
|
|
|
||
|
|
await self.telegram_helper.abort_processing(mock_update, mock_context)
|
||
|
|
|
||
|
|
mock_query.answer.assert_called_once()
|
||
|
|
mock_abort_processing_logic.assert_called_once_with(123)
|
||
|
|
mock_query.edit_message_text.assert_called_once_with(text="Aborted!")
|
||
|
|
|
||
|
|
@patch('telegram_helper.browse_command', new_callable=AsyncMock)
|
||
|
|
async def test_browse_command_handler(self, mock_browse_command):
|
||
|
|
mock_update = MagicMock(spec=Update)
|
||
|
|
mock_context = MagicMock(spec=ContextTypes.DEFAULT_TYPE)
|
||
|
|
|
||
|
|
await self.telegram_helper.browse(mock_update, mock_context)
|
||
|
|
mock_browse_command.assert_called_once_with(mock_update, mock_context, self.mock_bot)
|
||
|
|
|
||
|
|
@patch.dict(os.environ, {'TELEGRAM_BOT_TOKEN': 'test_token'})
|
||
|
|
@patch('telegram_helper.Application.builder')
|
||
|
|
def test_run_method(self, mock_builder):
|
||
|
|
mock_app_builder = mock_builder.return_value
|
||
|
|
mock_app = mock_app_builder.token.return_value.build.return_value
|
||
|
|
|
||
|
|
self.telegram_helper.run()
|
||
|
|
|
||
|
|
mock_builder.assert_called_once()
|
||
|
|
mock_app_builder.token.assert_called_once_with('test_token')
|
||
|
|
mock_app.add_handler.assert_any_call(unittest.mock.ANY) # Check if handlers are added
|
||
|
|
mock_app.run_polling.assert_called_once()
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|