admin/cyclop
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Refactored gemini, openai and claude into one file and removed logic from the base class, also made helper class definable from command line

Browse Source
This commit is contained in:
Jonathan Lucas
2025-06-03 13:04:42 -05:00
parent bd0ce3e340
commit f15228fa58
36 changed files with 487 additions and 3847 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_telegram_helper.py
-356
View File
@@ -1,356 +0,0 @@
import unittest
from unittest.mock import MagicMock, patch, mock_open, AsyncMock
import asyncio
import os
import sys
# Assuming telegram_helper.py is in the parent directory or PYTHONPATH is set
from telegram_helper import TelegramHelper, MessageHandlerLogicResult
# Mock for the bot passed to TelegramHelper
class MockBot:
def __init__(self):
self.start = AsyncMock()
self.clear_conversation_history = MagicMock()
self.get_bot_status = AsyncMock(return_value="Bot Status OK")
self.switch_model = AsyncMock(return_value="Model Switched OK")
self.handle_message = AsyncMock() # Needs to return a string
self.abort_processing = AsyncMock(return_value="Abort OK")
self.set_processing_status = MagicMock()
self.clear_processing_status = MagicMock()
self.processing_status = {} # Add the attribute
# Mock for telegram.Update and related objects
def create_mock_update(message_text=None, user_id=123, chat_id=456, message_id=789, callback_query_data=None):
update = MagicMock()
update.effective_user.id = user_id
update.effective_chat.id = chat_id
if message_text:
update.message.text = message_text
update.message.reply_text = AsyncMock(return_value=MagicMock(message_id=message_id)) # reply_text returns a Message obj
if callback_query_data:
update.callback_query.data = callback_query_data
update.callback_query.from_user.id = user_id
update.callback_query.answer = AsyncMock()
update.callback_query.edit_message_text = AsyncMock()
return update
# Mock for telegram.ext.ContextTypes.DEFAULT_TYPE
def create_mock_context():
context = MagicMock()
context.bot.delete_message = AsyncMock()
context.bot.edit_message_text = AsyncMock() # For update_status_message
return context
class TestTelegramHelper(unittest.IsolatedAsyncioTestCase): # Use IsolatedAsyncioTestCase for async methods
def setUp(self):
self.mock_bot = MockBot()
# Default paths for reboot files, can be overridden in tests
self.reboot_claude_file = ".test_reboot_claude"
self.reboot_file = ".test_doreboot"
self.helper = TelegramHelper(
self.mock_bot,
reboot_claude_file_path=self.reboot_claude_file,
reboot_file_path=self.reboot_file,
chunk_message_sleep_duration=0.001 # Faster sleep for tests
)
# Clean up any potential leftover reboot files from previous runs
if os.path.exists(self.reboot_claude_file):
os.remove(self.reboot_claude_file)
if os.path.exists(self.reboot_file):
os.remove(self.reboot_file)
def tearDown(self):
# Clean up reboot files created during tests
if os.path.exists(self.reboot_claude_file):
os.remove(self.reboot_claude_file)
if os.path.exists(self.reboot_file):
os.remove(self.reboot_file)
async def test_start_logic(self):
response = await self.helper._start_logic()
self.mock_bot.start.assert_called_once()
self.assertEqual(response, "Hello! I\'m your AI assistant. How can I help you today?")
async def test_start_command(self):
mock_update = create_mock_update(message_text="/start")
mock_context = create_mock_context()
with patch.object(self.helper, \'_start_logic\', new_callable=AsyncMock) as mock_logic:
mock_logic.return_value = "Start Logic Response"
await self.helper.start(mock_update, mock_context)
mock_logic.assert_called_once()
mock_update.message.reply_text.assert_called_once_with("Start Logic Response")
async def test_clear_logic(self):
user_id = 123
response = await self.helper._clear_logic(user_id) # _clear_logic is async after refactor
self.mock_bot.clear_conversation_history.assert_called_once_with(user_id)
self.assertEqual(response, "Conversation history cleared. Let\'s start fresh!")
async def test_clear_command(self):
mock_update = create_mock_update(message_text="/clear", user_id=123)
mock_context = create_mock_context()
with patch.object(self.helper, \'_clear_logic\', new_callable=AsyncMock) as mock_logic:
mock_logic.return_value = "Clear Logic Response"
await self.helper.clear(mock_update, mock_context)
mock_logic.assert_called_once_with(123)
mock_update.message.reply_text.assert_called_once_with("Clear Logic Response")
async def test_status_logic(self):
self.mock_bot.get_bot_status.return_value = "Test Status"
response = await self.helper._status_logic()
self.mock_bot.get_bot_status.assert_called_once()
self.assertEqual(response, "Test Status")
async def test_switch_logic_supported(self):
self.mock_bot.switch_model.return_value = "Switched to Large Model"
response = await self.helper._switch_logic()
self.mock_bot.switch_model.assert_called_once()
self.assertEqual(response, "Switched to Large Model")
async def test_switch_logic_not_supported(self):
del self.mock_bot.switch_model # Simulate bot not having the attribute
response = await self.helper._switch_logic()
self.assertEqual(response, "Model switching is not supported for this bot.")
async def test_handle_message_logic_success(self):
user_id = 100
user_message = "Hello bot"
bot_response = "Hello user <think>Thinking hard</think> Done."
expected_processed_response = f"Hello user {self.helper.HTML_QUOTE_BLOCK_START}Thinking hard{self.helper.HTML_QUOTE_BLOCK_END} Done."
self.mock_bot.handle_message.return_value = bot_response
result = await self.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"], expected_processed_response)
self.assertIsNone(result["error_message"])
async def test_handle_message_logic_bot_exception(self):
user_id = 101
user_message = "Trigger error"
self.mock_bot.handle_message.side_effect = Exception("Bot Error")
result = await self.helper._handle_message_logic(user_id, user_message)
self.assertFalse(result["success"])
self.assertIsNone(result["response_text"])
self.assertEqual(result["error_message"], "Bot Error")
@patch(\'logging.error\')
async def test_handle_message_command_success_short_message(self, mock_logging_error):
mock_update = create_mock_update(message_text="Hi", user_id=200, chat_id=201, message_id=202)
mock_context = create_mock_context()
logic_result = MessageHandlerLogicResult(success=True, response_text="Short response", error_message=None)
with patch.object(self.helper, \'_handle_message_logic\', new_callable=AsyncMock) as mock_message_logic:
mock_message_logic.return_value = logic_result
await self.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(200, 202) # user_id, status_message_id
mock_message_logic.assert_called_once_with(200, "Hi")
mock_context.bot.delete_message.assert_called_once_with(chat_id=201, message_id=202)
self.mock_bot.clear_processing_status.assert_called_once_with(200)
mock_update.message.reply_text.assert_any_call("Short response") # Final response
self.assertEqual(mock_update.message.reply_text.call_count, 2) # Processing + final
@patch(\'logging.error\')
async def test_handle_message_command_success_long_message_chunks(self, mock_logging_error):
mock_update = create_mock_update(message_text="Long text", user_id=200, chat_id=201, message_id=202)
mock_context = create_mock_context()
long_response_text = "a" * 5000 # Longer than 4096
chunk1 = long_response_text[:4096]
chunk2 = long_response_text[4096:]
logic_result = MessageHandlerLogicResult(success=True, response_text=long_response_text, error_message=None)
with patch.object(self.helper, \'_handle_message_logic\', new_callable=AsyncMock) as mock_message_logic, \
patch(\'asyncio.sleep\', new_callable=AsyncMock) as mock_sleep: # Mock sleep
mock_message_logic.return_value = logic_result
await self.helper.handle_message(mock_update, mock_context)
mock_update.message.reply_text.assert_any_call(chunk1)
mock_update.message.reply_text.assert_any_call(chunk2)
mock_sleep.assert_called_once_with(self.helper.chunk_message_sleep_duration)
self.assertEqual(mock_update.message.reply_text.call_count, 3) # Processing + 2 chunks
@patch(\'logging.error\')
async def test_handle_message_command_logic_fails(self, mock_logging_error):
mock_update = create_mock_update(message_text="Cause error in logic", user_id=200)
mock_context = create_mock_context()
logic_result = MessageHandlerLogicResult(success=False, response_text=None, error_message="Logic Failed")
with patch.object(self.helper, \'_handle_message_logic\', new_callable=AsyncMock) as mock_message_logic:
mock_message_logic.return_value = logic_result
await self.helper.handle_message(mock_update, mock_context)
mock_update.message.reply_text.assert_any_call("Sorry, an error occurred while processing your request.")
self.assertEqual(mock_update.message.reply_text.call_count, 2) # Processing + error message
@patch(\'logging.error\')
async def test_handle_message_command_telegram_exception_after_logic(self, mock_logging_error):
mock_update = create_mock_update(message_text="Test", user_id=200)
mock_context = create_mock_context()
logic_result = MessageHandlerLogicResult(success=True, response_text="OK", error_message=None)
# Make sending the final reply fail
mock_update.message.reply_text.side_effect = [
MagicMock(message_id=202), # For "Processing..."
Exception("Telegram API Error") # For the actual response
]
with patch.object(self.helper, \'_handle_message_logic\', new_callable=AsyncMock) as mock_message_logic:
mock_message_logic.return_value = logic_result
await self.helper.handle_message(mock_update, mock_context)
# Check if the generic error message was attempted
# This is tricky because reply_text is already mocked with side_effect.
# We\'d expect logs. Let\'s check logs or if processing status was cleared.
self.mock_bot.clear_processing_status.assert_called_once_with(200)
mock_logging_error.assert_any_call(unittest.mock.string_containing("Outer error in handle_message"))
async def test_abort_processing_logic(self):
user_id = 300
self.mock_bot.abort_processing.return_value = "Aborted by bot"
response = await self.helper._abort_processing_logic(user_id)
self.mock_bot.abort_processing.assert_called_once_with(user_id)
self.assertEqual(response, "Aborted by bot")
async def test_abort_processing_command(self):
mock_update = create_mock_update(callback_query_data=\'abort\', user_id=301)
mock_context = create_mock_context()
with patch.object(self.helper, \'_abort_processing_logic\', new_callable=AsyncMock) as mock_logic:
mock_logic.return_value = "Abort Logic Done"
await self.helper.abort_processing(mock_update, mock_context)
mock_update.callback_query.answer.assert_called_once()
mock_logic.assert_called_once_with(301)
mock_update.callback_query.edit_message_text.assert_called_once_with(text="Abort Logic Done")
def test_reboot_logic_claude_and_main(self):
user_message_parts = ["/reboot", "claude"]
chat_id_to_write = "12345"
with patch("builtins.open", mock_open()) as mock_file:
self.helper._reboot_logic(user_message_parts, chat_id_to_write)
# Check claude reboot file
mock_file.assert_any_call(self.reboot_claude_file, \'w\')
# Check main doreboot file
mock_file.assert_any_call(self.reboot_file, \'w\')
handle_claude = mock_file.return_value
handle_main = mock_file.return_value # mock_open reuses the handle for multiple calls
# Check if write was called for claude file (empty write)
# This part of assertion is tricky with single mock_file. Better to use different mocks if possible
# or check the sequence of calls if the mock supports it well.
# For now, assert_any_call ensures it was opened.
# Check content for main reboot file
# Need to ensure the write for self.reboot_file had chat_id_to_write
# This requires more sophisticated mock_open or patching os.path.exists and multiple open calls
# Simpler check: was open(self.reboot_file, \'w\') called? Yes, via assert_any_call.
# And was open(self.reboot_claude_file, \'w\') called? Yes.
# Verify files were created (mock_open doesn\'t actually create them)
# This test relies on mock_open\'s behavior. To test file content, need more setup.
# For now, assume open was called correctly.
def test_reboot_logic_main_only(self):
user_message_parts = ["/reboot"]
chat_id_to_write = "67890"
with patch("builtins.open", mock_open()) as mock_file:
self.helper._reboot_logic(user_message_parts, chat_id_to_write)
# Ensure claude file was NOT opened for writing.
# This requires asserting that a specific call didn\'t happen, or checking call_args_list
claude_call = unittest.mock.call(self.reboot_claude_file, \'w\')
self.assertNotIn(claude_call, mock_file.call_args_list)
mock_file.assert_any_call(self.reboot_file, \'w\')
@patch(\'sys.exit\') # Mock sys.exit to prevent test runner from exiting
async def test_reboot_command(self, mock_sys_exit):
mock_update = create_mock_update(message_text="/reboot claude", chat_id="chat1")
mock_context = create_mock_context()
with patch.object(self.helper, \'_reboot_logic\') as mock_reboot_file_logic:
await self.helper.reboot(mock_update, mock_context)
mock_reboot_file_logic.assert_called_once_with(["/reboot", "claude"], "chat1")
mock_update.message.reply_text.assert_called_once_with("Rebooting the bot...")
mock_sys_exit.assert_called_once_with(0)
@patch(\'os.path.exists\')
@patch(\'builtins.open\', new_callable=mock_open)
@patch(\'os.remove\')
async def test_check_doreboot_file_logic_file_exists(self, mock_os_remove, mock_file_open, mock_os_path_exists):
mock_os_path_exists.return_value = True
mock_file_open.return_value.read.return_value.strip.return_value = "chat123"
chat_id = await self.helper._check_doreboot_file_logic()
mock_os_path_exists.assert_called_once_with(self.reboot_file)
mock_file_open.assert_called_once_with(self.reboot_file, \'r\')
mock_os_remove.assert_called_once_with(self.reboot_file)
self.assertEqual(chat_id, "chat123")
@patch(\'os.path.exists\', return_value=False)
async def test_check_doreboot_file_logic_file_not_exists(self, mock_os_path_exists):
chat_id = await self.helper._check_doreboot_file_logic()
mock_os_path_exists.assert_called_once_with(self.reboot_file)
self.assertIsNone(chat_id)
@patch(\'logging.error\')
@patch(\'os.path.exists\', return_value=True)
@patch(\'builtins.open\', side_effect=IOError("Read error"))
@patch(\'os.remove\') # To check if remove is called even on read error
async def test_check_doreboot_file_logic_read_error(self, mock_os_remove, mock_file_open, mock_os_path_exists, mock_log_error):
chat_id = await self.helper._check_doreboot_file_logic()
self.assertIsNone(chat_id)
mock_log_error.assert_any_call(unittest.mock.string_containing("Error reading reboot file"))
# Check if os.remove was attempted even after read error
mock_os_remove.assert_called_once_with(self.reboot_file)
async def test_check_doreboot_file_command_sends_message(self):
mock_application = MagicMock()
mock_application.bot.send_message = AsyncMock()
with patch.object(self.helper, \'_check_doreboot_file_logic\', new_callable=AsyncMock) as mock_logic:
mock_logic.return_value = "chat789" # Simulate chat_id found
await self.helper.check_doreboot_file(mock_application)
mock_logic.assert_called_once()
mock_application.bot.send_message.assert_called_once_with(
chat_id="chat789", text="The application has finished initializing."
)
async def test_check_doreboot_file_command_no_chat_id(self):
mock_application = MagicMock()
mock_application.bot.send_message = AsyncMock()
with patch.object(self.helper, \'_check_doreboot_file_logic\', new_callable=AsyncMock) as mock_logic:
mock_logic.return_value = None # Simulate no chat_id found
await self.helper.check_doreboot_file(mock_application)
mock_logic.assert_called_once()
mock_application.bot.send_message.assert_not_called()
# Note: Testing the run() method itself is more of an integration test,
# as it involves setting up the full Application and polling loop.
# Unit tests here focus on the helper\'s own logic methods.
if __name__ == \'__main__\':
unittest.main()