Jack1808/Claude_Code
0
1from unittest.mock import AsyncMock, MagicMock, patch2 3import pytest4 5from messaging.handler import ClaudeMessageHandler6from messaging.models import IncomingMessage7from messaging.trees.data import MessageNode, MessageTree8from messaging.trees.queue_manager import MessageState9 10 11@pytest.fixture12def handler(mock_platform, mock_cli_manager, mock_session_store):13 return ClaudeMessageHandler(mock_platform, mock_cli_manager, mock_session_store)14 15 16def test_get_initial_status_new_conversation(handler):17 """New conversation always returns launching message."""18 result = handler._get_initial_status(None, None)19 assert "Launching" in result20 21 22def test_get_initial_status_reply_tree_busy_queued(handler):23 """Reply to tree when busy returns queued message."""24 mock_queue = MagicMock()25 mock_queue.is_node_tree_busy.return_value = True26 mock_queue.get_queue_size.return_value = 227 handler.replace_tree_queue(mock_queue)28 result = handler._get_initial_status(MagicMock(), "parent_1")29 assert "Queued" in result30 assert "position 3" in result31 32 33def test_get_initial_status_reply_tree_not_busy_continuing(handler):34 """Reply to tree when not busy returns continuing message."""35 mock_queue = MagicMock()36 mock_queue.is_node_tree_busy.return_value = False37 handler.replace_tree_queue(mock_queue)38 result = handler._get_initial_status(MagicMock(), "parent_1")39 assert "Continuing" in result40 41 42@pytest.mark.asyncio43async def test_handle_message_stop_command(44 handler, mock_platform, incoming_message_factory45):46 incoming = incoming_message_factory(text="/stop")47 48 # Mock stop_all_tasks49 handler.stop_all_tasks = AsyncMock(return_value=5)50 51 await handler.handle_message(incoming)52 53 handler.stop_all_tasks.assert_called_once()54 mock_platform.queue_send_message.assert_called_once_with(55 incoming.chat_id,56 "⏹ *Stopped\\.* Cancelled 5 pending or active requests\\.",57 fire_and_forget=False,58 message_thread_id=None,59 )60 61 62@pytest.mark.asyncio63async def test_handle_message_stop_command_reply_stops_only_target_node(64 handler, mock_platform, mock_cli_manager, incoming_message_factory65):66 # Create a tree with a root node and register its status message ID mapping.67 root_incoming = incoming_message_factory(68 text="do something", message_id="root_msg", reply_to_message_id=None69 )70 tree = await handler.tree_queue.create_tree(71 node_id="root_msg",72 incoming=root_incoming,73 status_message_id="status_root",74 )75 handler.tree_queue.register_node("status_root", tree.root_id)76 77 # Reply "/stop" to the status message; should stop only that node.78 incoming = incoming_message_factory(79 text="/stop",80 message_id="stop_msg",81 reply_to_message_id="status_root",82 )83 84 handler.stop_all_tasks = AsyncMock(return_value=999)85 86 await handler.handle_message(incoming)87 88 handler.stop_all_tasks.assert_not_called()89 mock_cli_manager.stop_all.assert_not_called()90 assert tree.get_node("root_msg").state == MessageState.ERROR91 mock_platform.queue_send_message.assert_called_once_with(92 incoming.chat_id,93 "⏹ *Stopped\\.* Cancelled 1 request\\.",94 fire_and_forget=False,95 message_thread_id=None,96 )97 98 99@pytest.mark.asyncio100async def test_handle_message_stop_command_reply_unknown_does_not_stop_all(101 handler, mock_platform, mock_cli_manager, incoming_message_factory102):103 incoming = incoming_message_factory(104 text="/stop",105 message_id="stop_msg",106 reply_to_message_id="unknown_msg",107 )108 109 handler.stop_all_tasks = AsyncMock(return_value=5)110 111 await handler.handle_message(incoming)112 113 handler.stop_all_tasks.assert_not_called()114 mock_cli_manager.stop_all.assert_not_called()115 mock_platform.queue_send_message.assert_called_once_with(116 incoming.chat_id,117 "⏹ *Stopped\\.* Nothing to stop for that message\\.",118 fire_and_forget=False,119 message_thread_id=None,120 )121 122 123@pytest.mark.asyncio124async def test_handle_message_stats_command(125 handler, mock_platform, mock_cli_manager, incoming_message_factory126):127 incoming = incoming_message_factory(text="/stats")128 mock_cli_manager.get_stats.return_value = {"active_sessions": 2}129 130 await handler.handle_message(incoming)131 132 mock_platform.queue_send_message.assert_called_once()133 args, kwargs = mock_platform.queue_send_message.call_args134 assert "Active CLI: 2" in args[1]135 assert kwargs["fire_and_forget"] is False136 assert kwargs.get("message_thread_id") is None137 138 139@pytest.mark.asyncio140async def test_handle_message_filters_status_messages(141 handler, mock_platform, incoming_message_factory142):143 incoming = incoming_message_factory(text="⏳ Thinking...")144 145 await handler.handle_message(incoming)146 147 mock_platform.queue_send_message.assert_not_called()148 149 150@pytest.mark.asyncio151async def test_handle_message_new_conversation(152 handler, mock_platform, mock_session_store, incoming_message_factory153):154 incoming = incoming_message_factory(text="hello")155 mock_platform.queue_send_message.return_value = "status_123"156 157 # We need to mock tree_queue methods158 with (159 patch.object(handler.tree_queue, "create_tree", AsyncMock()) as mock_create,160 patch.object(161 handler.tree_queue, "enqueue", AsyncMock(return_value=False)162 ) as mock_enqueue,163 ):164 mock_tree = MagicMock()165 mock_tree.root_id = "root_1"166 mock_tree.to_dict.return_value = {"data": "tree"}167 mock_create.return_value = mock_tree168 169 await handler.handle_message(incoming)170 171 mock_create.assert_called_once()172 mock_enqueue.assert_called_once()173 mock_session_store.save_tree.assert_called_once_with("root_1", {"data": "tree"})174 175 176@pytest.mark.asyncio177async def test_handle_message_queued(handler, mock_platform, incoming_message_factory):178 incoming = incoming_message_factory(text="hello", message_id="msg_1")179 mock_platform.queue_send_message.return_value = "status_123"180 181 with (182 patch.object(handler.tree_queue, "create_tree", AsyncMock()) as mock_create,183 patch.object(handler.tree_queue, "enqueue", AsyncMock(return_value=True)),184 patch.object(handler.tree_queue, "get_queue_size", MagicMock(return_value=3)),185 ):186 mock_tree = MagicMock()187 mock_tree.root_id = "root_1"188 mock_tree.to_dict.return_value = {}189 mock_create.return_value = mock_tree190 191 await handler.handle_message(incoming)192 193 mock_platform.queue_edit_message.assert_called_once_with(194 incoming.chat_id,195 "status_123",196 "📋 *Queued* \\(position 3\\) \\- waiting\\.\\.\\.",197 parse_mode="MarkdownV2",198 )199 200 201@pytest.mark.asyncio202async def test_update_queue_positions(handler, mock_platform):203 root_incoming = IncomingMessage(204 text="Root",205 chat_id="chat_1",206 user_id="user_1",207 message_id="root",208 platform="telegram",209 )210 root = MessageNode(211 node_id="root",212 incoming=root_incoming,213 status_message_id="status_root",214 )215 tree = MessageTree(root)216 217 child_incoming_1 = IncomingMessage(218 text="Child 1",219 chat_id="chat_1",220 user_id="user_1",221 message_id="child_1",222 platform="telegram",223 reply_to_message_id="root",224 )225 child_incoming_2 = IncomingMessage(226 text="Child 2",227 chat_id="chat_1",228 user_id="user_1",229 message_id="child_2",230 platform="telegram",231 reply_to_message_id="root",232 )233 234 await tree.add_node(235 node_id="child_1",236 incoming=child_incoming_1,237 status_message_id="status_1",238 parent_id="root",239 )240 await tree.add_node(241 node_id="child_2",242 incoming=child_incoming_2,243 status_message_id="status_2",244 parent_id="root",245 )246 247 await tree.enqueue("child_1")248 await tree.enqueue("child_2")249 250 await handler.update_queue_positions(tree)251 252 calls = mock_platform.queue_edit_message.call_args_list253 assert len(calls) == 2254 assert calls[0][0][0] == "chat_1"255 assert calls[0][0][1] == "status_1"256 assert "position 1" in calls[0][0][2]257 assert calls[1][0][0] == "chat_1"258 assert calls[1][0][1] == "status_2"259 assert "position 2" in calls[1][0][2]260 261 262@pytest.mark.asyncio263async def test_mark_node_processing(handler, mock_platform):264 root_incoming = IncomingMessage(265 text="Root",266 chat_id="chat_1",267 user_id="user_1",268 message_id="root",269 platform="telegram",270 )271 root = MessageNode(272 node_id="root",273 incoming=root_incoming,274 status_message_id="status_root",275 )276 tree = MessageTree(root)277 278 child_incoming = IncomingMessage(279 text="Child",280 chat_id="chat_1",281 user_id="user_1",282 message_id="child",283 platform="telegram",284 reply_to_message_id="root",285 )286 287 await tree.add_node(288 node_id="child",289 incoming=child_incoming,290 status_message_id="status_child",291 parent_id="root",292 )293 294 await handler.mark_node_processing(tree, "child")295 296 mock_platform.queue_edit_message.assert_called_once()297 args, kwargs = mock_platform.queue_edit_message.call_args298 assert args[0] == "chat_1"299 assert args[1] == "status_child"300 assert "Processing" in args[2]301 assert kwargs["parse_mode"] == "MarkdownV2"302 303 304@pytest.mark.asyncio305async def test_stop_all_tasks(handler, mock_cli_manager, mock_platform):306 mock_node = MagicMock()307 mock_node.incoming.chat_id = "chat_1"308 mock_node.status_message_id = "status_1"309 310 with patch.object(311 handler.tree_queue, "cancel_all", AsyncMock(return_value=[mock_node])312 ):313 count = await handler.stop_all_tasks()314 315 assert count == 1316 mock_cli_manager.stop_all.assert_called_once()317 mock_platform.fire_and_forget.assert_called_once()318 319 320async def mock_async_gen(events):321 for e in events:322 yield e323 324 325@pytest.mark.asyncio326async def test_process_node_success_flow(handler, mock_cli_manager, mock_platform):327 # Setup328 node_id = "node_1"329 mock_node = MagicMock()330 mock_node.incoming.chat_id = "chat_1"331 mock_node.incoming.text = "hello"332 mock_node.status_message_id = "status_1"333 mock_node.parent_id = None334 335 mock_session = MagicMock()336 # Mock start_task to return our async generator337 events = [338 {339 "type": "assistant",340 "message": {"content": [{"type": "thinking", "thinking": "Let me think"}]},341 },342 {343 "type": "assistant",344 "message": {"content": [{"type": "text", "text": "Hello world"}]},345 },346 {"type": "exit", "code": 0},347 ]348 mock_session.start_task.return_value = mock_async_gen(events)349 350 mock_cli_manager.get_or_create_session.return_value = (351 mock_session,352 "session_1",353 False,354 )355 356 mock_tree = MagicMock()357 mock_tree.update_state = AsyncMock()358 mock_tree.root_id = "root_1"359 mock_tree.to_dict.return_value = {}360 361 with patch.object(362 handler.tree_queue, "get_tree_for_node", MagicMock(return_value=mock_tree)363 ):364 await handler._process_node(node_id, mock_node)365 366 # Verify state updates367 mock_tree.update_state.assert_any_call(node_id, MessageState.IN_PROGRESS)368 mock_tree.update_state.assert_any_call(369 node_id, MessageState.COMPLETED, session_id="session_1"370 )371 372 # Verify UI updates (at least the final one)373 # Note: update_ui is debounced, but COMPLETED/ERROR/CANCELLED are forced374 mock_platform.queue_edit_message.assert_called()375 last_call = mock_platform.queue_edit_message.call_args_list[-1]376 assert "✅ *Complete*" in last_call[0][2]377 assert "Hello world" in last_call[0][2]378 379 380@pytest.mark.asyncio381async def test_process_node_error_flow(handler, mock_cli_manager, mock_platform):382 node_id = "node_1"383 mock_node = MagicMock()384 mock_node.incoming.chat_id = "chat_1"385 mock_node.incoming.text = "hello"386 mock_node.status_message_id = "status_1"387 388 mock_session = MagicMock()389 events = [{"type": "error", "error": {"message": "CLI crashed"}}]390 mock_session.start_task.return_value = mock_async_gen(events)391 mock_cli_manager.get_or_create_session.return_value = (392 mock_session,393 "session_1",394 False,395 )396 397 mock_tree = MagicMock()398 mock_tree.update_state = AsyncMock()399 400 with (401 patch.object(402 handler.tree_queue, "get_tree_for_node", MagicMock(return_value=mock_tree)403 ),404 patch.object(405 handler.tree_queue, "mark_node_error", AsyncMock(return_value=[mock_node])406 ),407 ):408 await handler._process_node(node_id, mock_node)409 410 handler.tree_queue.mark_node_error.assert_called_once_with(411 node_id, "CLI crashed", propagate_to_children=True412 )413 414 last_call = mock_platform.queue_edit_message.call_args_list[-1]415 assert "❌ *Error*" in last_call[0][2]416 assert "CLI crashed" in last_call[0][2]417 418 419@pytest.mark.asyncio420async def test_handle_message_clear_command_stops_deletes_and_wipes_state(421 handler, mock_platform, mock_session_store, incoming_message_factory422):423 # Create some tracked messages across two chats. /clear should only delete424 # messages for the current chat.425 root_1 = incoming_message_factory(426 text="do something",427 chat_id="chat_1",428 message_id="100",429 reply_to_message_id=None,430 )431 await handler.tree_queue.create_tree(432 node_id="100",433 incoming=root_1,434 status_message_id="101",435 )436 437 root_2 = incoming_message_factory(438 text="other chat",439 chat_id="chat_2",440 message_id="200",441 reply_to_message_id=None,442 )443 await handler.tree_queue.create_tree(444 node_id="200",445 incoming=root_2,446 status_message_id="201",447 )448 449 events = []450 451 async def _stop():452 events.append("stop")453 return 0454 455 async def _del(chat_id, message_id, fire_and_forget=True):456 events.append(f"del:{chat_id}:{message_id}:{fire_and_forget}")457 458 handler.stop_all_tasks = AsyncMock(side_effect=_stop)459 mock_platform.queue_delete_message = AsyncMock(side_effect=_del)460 461 incoming = incoming_message_factory(462 text="/clear", chat_id="chat_1", message_id="150"463 )464 await handler.handle_message(incoming)465 466 assert events and events[0] == "stop"467 deleted_ids = {e.split(":")[2] for e in events[1:]}468 assert deleted_ids == {"100", "101", "150"}469 assert all(e.endswith(":False") for e in events[1:])470 471 mock_session_store.clear_all.assert_called_once()472 assert handler.tree_queue.get_tree_count() == 0473 mock_platform.queue_send_message.assert_not_called()474 475 476@pytest.mark.asyncio477async def test_handle_message_clear_command_with_mention(478 handler, mock_platform, mock_session_store, incoming_message_factory479):480 handler.stop_all_tasks = AsyncMock(return_value=0)481 482 incoming = incoming_message_factory(483 text="/clear@MyBot", chat_id="chat_1", message_id="10"484 )485 await handler.handle_message(incoming)486 487 handler.stop_all_tasks.assert_called_once()488 mock_platform.queue_delete_message.assert_called_once_with(489 "chat_1",490 "10",491 fire_and_forget=False,492 )493 mock_session_store.clear_all.assert_called_once()494 495 496@pytest.mark.asyncio497async def test_handle_message_clear_command_deletes_message_log_ids(498 handler, mock_platform, mock_session_store, incoming_message_factory499):500 handler.stop_all_tasks = AsyncMock(return_value=0)501 mock_session_store.get_message_ids_for_chat.return_value = ["42", "43"]502 503 incoming = incoming_message_factory(504 text="/clear", chat_id="chat_1", message_id="150"505 )506 await handler.handle_message(incoming)507 508 deleted = {c.args[1] for c in mock_platform.queue_delete_message.call_args_list}509 assert deleted == {"42", "43", "150"}510 511 512@pytest.mark.asyncio513async def test_handle_message_clear_command_reply_clears_branch(514 handler, mock_platform, mock_session_store, incoming_message_factory515):516 """Reply /clear to a message clears only that branch."""517 root_incoming = incoming_message_factory(518 text="root", chat_id="chat_1", message_id="100", reply_to_message_id=None519 )520 tree = await handler.tree_queue.create_tree(521 node_id="100", incoming=root_incoming, status_message_id="101"522 )523 handler.tree_queue.register_node("101", tree.root_id)524 525 child_incoming = incoming_message_factory(526 text="child",527 chat_id="chat_1",528 message_id="102",529 reply_to_message_id="100",530 )531 await handler.tree_queue.add_to_tree(532 parent_node_id="100",533 node_id="102",534 incoming=child_incoming,535 status_message_id="103",536 )537 538 deleted_ids = []539 540 async def _capture_delete(chat_id, message_id, fire_and_forget=True):541 deleted_ids.append(message_id)542 543 mock_platform.queue_delete_message = AsyncMock(side_effect=_capture_delete)544 545 incoming = incoming_message_factory(546 text="/clear",547 chat_id="chat_1",548 message_id="150",549 reply_to_message_id="102",550 )551 await handler.handle_message(incoming)552 553 assert set(deleted_ids) == {"102", "103", "150"}554 assert "100" not in deleted_ids555 assert "101" not in deleted_ids556 mock_session_store.remove_node_mappings.assert_called()557 assert handler.tree_queue.get_tree_for_node("102") is None558 assert handler.tree_queue.get_tree_for_node("100") is not None559 560 561@pytest.mark.asyncio562async def test_handle_message_clear_command_reply_unknown_sends_nothing(563 handler, mock_platform, mock_session_store, incoming_message_factory564):565 """Reply /clear to unknown message sends 'Nothing to clear'."""566 incoming = incoming_message_factory(567 text="/clear",568 chat_id="chat_1",569 message_id="150",570 reply_to_message_id="999",571 )572 await handler.handle_message(incoming)573 574 mock_platform.queue_send_message.assert_called_once()575 call_args = mock_platform.queue_send_message.call_args[0]576 assert "Nothing to clear" in call_args[1]577 mock_session_store.clear_all.assert_not_called()578 579 580@pytest.mark.asyncio581async def test_handle_message_clear_command_reply_to_root_clears_tree(582 handler, mock_platform, mock_session_store, incoming_message_factory583):584 """Reply /clear to root message clears entire tree."""585 root_incoming = incoming_message_factory(586 text="root", chat_id="chat_1", message_id="100", reply_to_message_id=None587 )588 await handler.tree_queue.create_tree(589 node_id="100", incoming=root_incoming, status_message_id="101"590 )591 592 deleted_ids = []593 594 async def _capture_delete(chat_id, message_id, fire_and_forget=True):595 deleted_ids.append(message_id)596 597 mock_platform.queue_delete_message = AsyncMock(side_effect=_capture_delete)598 599 incoming = incoming_message_factory(600 text="/clear",601 chat_id="chat_1",602 message_id="150",603 reply_to_message_id="100",604 )605 await handler.handle_message(incoming)606 607 assert set(deleted_ids) == {"100", "101", "150"}608 mock_session_store.remove_tree.assert_called_once_with("100")609 assert handler.tree_queue.get_tree_count() == 0610 611 612@pytest.mark.asyncio613async def test_handle_message_clear_command_reply_pending_voice_cancels(614 handler, mock_platform, mock_session_store, incoming_message_factory615):616 """Reply /clear to a voice note during transcription cancels it."""617 618 async def cancel_pending(chat_id, reply_id):619 if reply_id == "100":620 return ("100", "101")621 return None622 623 mock_platform.cancel_pending_voice = AsyncMock(side_effect=cancel_pending)624 mock_platform.queue_delete_message = AsyncMock()625 deleted_ids = []626 627 async def _capture_delete(chat_id, message_id, fire_and_forget=True):628 deleted_ids.append(message_id)629 630 mock_platform.queue_delete_message = AsyncMock(side_effect=_capture_delete)631 632 incoming = incoming_message_factory(633 text="/clear",634 chat_id="chat_1",635 message_id="150",636 reply_to_message_id="100",637 )638 await handler.handle_message(incoming)639 640 mock_platform.cancel_pending_voice.assert_called_once_with("chat_1", "100")641 assert set(deleted_ids) == {"100", "101", "150"}642 call_args = mock_platform.queue_send_message.call_args[0]643 assert "Voice note cancelled" in call_args[1]644