Fix: Address test failures in telegram_helper.py tests

This commit is contained in:
cyclop-bot
2025-06-03 13:08:39 -05:00
parent 3208cc6cd3
commit 79e6f8bd18
+67 -61
View File
@@ -2,7 +2,7 @@ import unittest
from unittest.mock import AsyncMock, MagicMock, patch
import asyncio
from telegram_helper import TelegramHelper, LogicResult
from telegram import Update
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import ContextTypes
import os
@@ -10,6 +10,16 @@ class TestTelegramHelper(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.mock_bot = AsyncMock()
# Mocking Update and ContextTypes.DEFAULT_TYPE with AsyncMock for methods that are awaited
self.mock_update = AsyncMock(spec=Update)
self.mock_context = AsyncMock(spec=ContextTypes.DEFAULT_TYPE)
# Ensure message and effective_chat are also AsyncMock if their methods are awaited
self.mock_update.message = AsyncMock()
self.mock_update.effective_chat = AsyncMock()
self.mock_update.callback_query = AsyncMock() # For abort_processing
# Mock context.bot as an AsyncMock
self.mock_context.bot = AsyncMock()
self.telegram_helper = TelegramHelper(self.mock_bot)
async def test_start_logic(self):
@@ -72,130 +82,126 @@ class TestTelegramHelper(unittest.IsolatedAsyncioTestCase):
@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)
# Using self.mock_update and self.mock_context from setUp
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!")
await self.telegram_helper.start(self.mock_update, self.mock_context)
self.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
self.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!")
await self.telegram_helper.clear(self.mock_update, self.mock_context)
self.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")
await self.telegram_helper.status(self.mock_update, self.mock_context)
self.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")
await self.telegram_helper.switch(self.mock_update, self.mock_context)
self.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)
self.mock_update.effective_user.id = 123
self.mock_update.message.text = "User message"
self.mock_update.effective_chat.id = 456
# Mock the return value of the first reply_text call
self.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)
await self.telegram_helper.handle_message(self.mock_update, self.mock_context)
mock_update.message.reply_text.assert_any_call("Processing your request...", reply_markup=unittest.mock.ANY)
# Assert the first call to reply_text for "Processing your request..."
self.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_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")
# Assert the second call to reply_text for the actual response
self.mock_update.message.reply_text.assert_any_call("Short response")
# Ensure total calls are 2 (processing + actual response)
self.assertEqual(self.mock_update.message.reply_text.call_count, 2)
@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)
self.mock_update.effective_user.id = 123
self.mock_update.message.text = "User message"
self.mock_update.effective_chat.id = 456
self.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
await self.telegram_helper.handle_message(self.mock_update, self.mock_context)
# Initial call + two chunks
self.assertEqual(self.mock_update.message.reply_text.call_count, 1 + 2)
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)
self.mock_update.effective_user.id = 123
self.mock_update.message.text = "User message"
self.mock_update.effective_chat.id = 456
self.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)
await self.telegram_helper.handle_message(self.mock_update, self.mock_context)
mock_context.bot.delete_message.assert_called_once()
self.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.")
self.mock_update.message.reply_text.assert_any_call("Sorry, an error occurred while processing your request.")
# Also assert the initial message
self.mock_update.message.reply_text.assert_any_call("Processing your request...", reply_markup=unittest.mock.ANY)
self.assertEqual(self.mock_update.message.reply_text.call_count, 2)
@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)
# Use self.mock_update.callback_query which is an AsyncMock
self.mock_update.callback_query.from_user.id = 123
mock_abort_processing_logic.return_value = "Aborted!"
await self.telegram_helper.abort_processing(mock_update, mock_context)
await self.telegram_helper.abort_processing(self.mock_update, self.mock_context)
mock_query.answer.assert_called_once()
self.mock_update.callback_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!")
self.mock_update.callback_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(self.mock_update, self.mock_context)
mock_browse_command.assert_called_once_with(self.mock_update, self.mock_context, self.mock_bot)
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('os.getenv', return_value='test_token')
@patch('telegram_helper.Application.builder')
def test_run_method(self, mock_builder):
async def test_run_method(self, mock_builder, mock_getenv):
# Re-initialize telegram_helper after patching os.getenv
self.telegram_helper = TelegramHelper(self.mock_bot)
mock_app_builder = mock_builder.return_value
mock_app = mock_app_builder.token.return_value.build.return_value
self.telegram_helper.run()
mock_getenv.assert_called_once_with('TELEGRAM_BOT_TOKEN')
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.add_handler.assert_called() # Check if handlers are added at least once
mock_app.run_polling.assert_called_once()
if __name__ == '__main__':
unittest.main()
unittest.main()