Jack1808/Claude_Code
0
1from unittest.mock import AsyncMock, MagicMock, patch2 3import pytest4 5from messaging.handler import ClaudeMessageHandler6from messaging.models import IncomingMessage7from messaging.rendering.telegram_markdown import render_markdown_to_mdv28from messaging.trees.data import MessageNode, MessageState9 10 11def test_render_markdown_to_mdv2_empty_returns_empty():12 assert render_markdown_to_mdv2("") == ""13 14 15def test_render_markdown_to_mdv2_covers_common_structures():16 md = (17 "# Heading\n\n"18 "Text with *em* and **strong** and ~~strike~~ and `code`.\n\n"19 "- item1\n"20 "- item2\n\n"21 "3. third\n\n"22 "> quote\n\n"23 "[link](http://example.com/a\\)b)\n\n"24 "\n\n"25 "```python\nprint('x')\n```\n"26 )27 out = render_markdown_to_mdv2(md)28 assert "*Heading*" in out29 assert "_em_" in out30 assert "*strong*" in out31 assert "~strike~" in out32 assert "`code`" in out33 assert "\\- item1" in out34 assert "3\\." in out35 assert "> quote" in out36 assert "[link]" in out37 assert "alt (http://example.com/img.png)" in out38 assert "```" in out39 40 41def test_render_markdown_to_mdv2_renders_table_as_code_block():42 md = "| a | b |\n|---|---|\n| 1 | 2 |\n| 3 | 4 |\n\nAfter.\n"43 out = render_markdown_to_mdv2(md)44 assert "```" in out45 assert "| a" in out46 assert "| b" in out47 assert "| ---" in out48 assert "After" in out49 50 51def test_render_markdown_to_mdv2_table_without_blank_line_still_renders():52 md = "Here's a table:\n| a | b |\n|---|---|\n| 1 | 2 |\n"53 out = render_markdown_to_mdv2(md)54 assert "Here's a table" in out55 assert "```" in out56 assert "| a" in out57 assert "| ---" in out58 59 60def test_render_markdown_to_mdv2_table_escapes_backticks_and_backslashes_in_cells():61 md = "| a | b |\n|---|---|\n| \\\\ | `` ` `` |\n"62 out = render_markdown_to_mdv2(md)63 assert "```" in out64 # In Telegram code blocks we escape backslashes and backticks.65 assert "\\\\" in out # rendered cell backslash becomes double-backslash66 assert "\\`" in out # rendered cell backtick is escaped67 68 69def test_render_markdown_to_mdv2_table_inside_list_keeps_bullet_prefix():70 md = "-\n | a | b |\n |---|---|\n | 1 | 2 |\n"71 out = render_markdown_to_mdv2(md)72 assert "```" in out73 assert out.lstrip().startswith("\\-")74 assert out.find("\\-") < out.find("```")75 76 77def test_get_initial_status_branches():78 platform = MagicMock()79 cli_manager = MagicMock()80 session_store = MagicMock()81 handler = ClaudeMessageHandler(platform, cli_manager, session_store)82 83 with (84 patch.object(85 handler.tree_queue, "is_node_tree_busy", MagicMock(return_value=True)86 ),87 patch.object(handler.tree_queue, "get_queue_size", MagicMock(return_value=2)),88 ):89 s1 = handler._get_initial_status(tree=object(), parent_node_id="p")90 assert "Queued" in s191 assert "position 3" in s1 or "position 3" in s1.replace("\\", "")92 93 with patch.object(94 handler.tree_queue, "is_node_tree_busy", MagicMock(return_value=False)95 ):96 s2 = handler._get_initial_status(tree=object(), parent_node_id="p")97 assert "Continuing" in s298 99 s3 = handler._get_initial_status(tree=None, parent_node_id=None)100 assert "Launching" in s3101 102 103@pytest.mark.asyncio104async def test_update_queue_positions_handles_snapshot_error_and_skips_non_pending():105 platform = MagicMock()106 platform.queue_edit_message = AsyncMock()107 platform.fire_and_forget = MagicMock(108 side_effect=lambda c: getattr(c, "close", lambda: None)()109 )110 111 cli_manager = MagicMock()112 session_store = MagicMock()113 handler = ClaudeMessageHandler(platform, cli_manager, session_store)114 115 # Snapshot error is swallowed.116 tree = MagicMock()117 tree.get_queue_snapshot = AsyncMock(side_effect=RuntimeError("boom"))118 await handler.update_queue_positions(tree)119 platform.fire_and_forget.assert_not_called()120 121 # Normal path: only PENDING nodes get an update.122 node_pending = MagicMock()123 node_pending.state = MessageState.PENDING124 node_pending.incoming.chat_id = "c"125 node_pending.status_message_id = "s"126 127 node_done = MagicMock()128 node_done.state = MessageState.COMPLETED129 130 tree.get_queue_snapshot = AsyncMock(return_value=["n1", "n2"])131 tree.get_node = MagicMock(side_effect=[node_pending, node_done])132 133 await handler.update_queue_positions(tree)134 assert platform.fire_and_forget.call_count == 1135 136 137@pytest.mark.asyncio138async def test_process_node_session_limit_marks_error_and_updates_ui():139 platform = MagicMock()140 platform.queue_edit_message = AsyncMock()141 platform.fire_and_forget = MagicMock(142 side_effect=lambda c: getattr(c, "close", lambda: None)()143 )144 145 cli_manager = MagicMock()146 cli_manager.get_or_create_session = AsyncMock(side_effect=RuntimeError("limit"))147 cli_manager.get_stats.return_value = {"active_sessions": 0}148 149 session_store = MagicMock()150 handler = ClaudeMessageHandler(platform, cli_manager, session_store)151 152 fake_tree = MagicMock()153 fake_tree.update_state = AsyncMock()154 with patch.object(155 handler.tree_queue, "get_tree_for_node", MagicMock(return_value=fake_tree)156 ):157 incoming = IncomingMessage(158 text="hi",159 chat_id="c",160 user_id="u",161 message_id="n1",162 platform="telegram",163 )164 node = MessageNode(node_id="n1", incoming=incoming, status_message_id="s1")165 166 await handler._process_node("n1", node)167 assert platform.queue_edit_message.await_count >= 1168 fake_tree.update_state.assert_awaited()169 170 171@pytest.mark.asyncio172async def test_stop_all_tasks_saves_tree_for_cancelled_nodes():173 platform = MagicMock()174 platform.queue_edit_message = AsyncMock()175 platform.fire_and_forget = MagicMock(176 side_effect=lambda c: getattr(c, "close", lambda: None)()177 )178 179 cli_manager = MagicMock()180 cli_manager.stop_all = AsyncMock()181 cli_manager.get_stats.return_value = {"active_sessions": 0}182 183 session_store = MagicMock()184 handler = ClaudeMessageHandler(platform, cli_manager, session_store)185 186 incoming = IncomingMessage(187 text="hi",188 chat_id="c",189 user_id="u",190 message_id="n1",191 platform="telegram",192 )193 node = MessageNode(node_id="n1", incoming=incoming, status_message_id="s1")194 195 tree = MagicMock()196 tree.root_id = "root"197 tree.to_dict = MagicMock(return_value={"root": "ok"})198 with (199 patch.object(handler.tree_queue, "cancel_all", AsyncMock(return_value=[node])),200 patch.object(201 handler.tree_queue, "get_tree_for_node", MagicMock(return_value=tree)202 ),203 ):204 count = await handler.stop_all_tasks()205 assert count == 1206 cli_manager.stop_all.assert_awaited_once()207 session_store.save_tree.assert_called_once_with("root", {"root": "ok"})208 209 210@pytest.mark.asyncio211async def test_handle_message_reply_with_tree_but_no_parent_treated_as_new():212 platform = MagicMock()213 platform.queue_send_message = AsyncMock(return_value="status_1")214 platform.queue_edit_message = AsyncMock()215 216 cli_manager = MagicMock()217 cli_manager.get_stats.return_value = {"active_sessions": 0}218 219 session_store = MagicMock()220 handler = ClaudeMessageHandler(platform, cli_manager, session_store)221 222 # Force "tree exists but parent can't be resolved" branch.223 mock_queue = MagicMock()224 mock_queue.get_tree_for_node.return_value = object()225 mock_queue.resolve_parent_node_id.return_value = None226 mock_queue.create_tree = AsyncMock(227 return_value=MagicMock(root_id="root", to_dict=MagicMock(return_value={"t": 1}))228 )229 mock_queue.register_node = MagicMock()230 mock_queue.enqueue = AsyncMock(return_value=False)231 handler.replace_tree_queue(mock_queue)232 233 incoming = IncomingMessage(234 text="reply",235 chat_id="c",236 user_id="u",237 message_id="m1",238 platform="telegram",239 reply_to_message_id="some_reply",240 )241 242 await handler.handle_message(incoming)243 mock_queue.create_tree.assert_awaited_once()244 245 246@pytest.mark.asyncio247async def test_update_ui_handles_transcript_render_exception():248 """When transcript.render raises, update_ui catches and does not crash."""249 platform = MagicMock()250 platform.queue_edit_message = AsyncMock()251 platform.fire_and_forget = MagicMock(252 side_effect=lambda c: getattr(c, "close", lambda: None)()253 )254 255 cli_manager = MagicMock()256 session_store = MagicMock()257 258 async def _mock_start_task(*args, **kwargs):259 yield {260 "type": "content_block_delta",261 "index": 0,262 "delta": {"type": "text_delta", "text": "hi"},263 }264 yield {"type": "complete", "status": "success"}265 266 mock_session = MagicMock()267 mock_session.start_task = _mock_start_task268 cli_manager.get_or_create_session = AsyncMock(269 return_value=(mock_session, "s1", False)270 )271 cli_manager.remove_session = AsyncMock()272 cli_manager.get_stats.return_value = {"active_sessions": 0}273 274 handler = ClaudeMessageHandler(platform, cli_manager, session_store)275 mock_queue = MagicMock()276 mock_queue.get_tree_for_node.return_value = None277 handler.replace_tree_queue(mock_queue)278 279 incoming = IncomingMessage(280 text="hi",281 chat_id="c",282 user_id="u",283 message_id="n1",284 platform="telegram",285 )286 node = MessageNode(node_id="n1", incoming=incoming, status_message_id="s1")287 288 with patch.object(handler, "_create_transcript_and_render_ctx") as mock_create:289 transcript = MagicMock()290 transcript.render = MagicMock(side_effect=ValueError("render failed"))291 render_ctx = MagicMock()292 mock_create.return_value = (transcript, render_ctx)293 294 await handler._process_node("n1", node)295 296 assert transcript.render.call_count >= 1297 298 299@pytest.mark.asyncio300async def test_handle_message_incoming_text_none_safe():301 """handle_message does not crash when incoming.text is None (e.g. malformed adapter)."""302 platform = MagicMock()303 platform.queue_send_message = AsyncMock(return_value="status_1")304 platform.queue_edit_message = AsyncMock()305 306 cli_manager = MagicMock()307 cli_manager.get_stats.return_value = {"active_sessions": 0}308 309 session_store = MagicMock()310 handler = ClaudeMessageHandler(platform, cli_manager, session_store)311 mock_queue = MagicMock()312 mock_queue.get_tree_for_node.return_value = None313 mock_queue.resolve_parent_node_id.return_value = None314 mock_queue.create_tree = AsyncMock(315 return_value=MagicMock(root_id="root", to_dict=MagicMock(return_value={"t": 1}))316 )317 mock_queue.register_node = MagicMock()318 mock_queue.enqueue = AsyncMock(return_value=True)319 handler.replace_tree_queue(mock_queue)320 321 incoming = MagicMock()322 incoming.text = None323 incoming.chat_id = "c"324 incoming.user_id = "u"325 incoming.message_id = "m1"326 incoming.platform = "telegram"327 incoming.reply_to_message_id = None328 incoming.is_reply = MagicMock(return_value=False)329 330 await handler.handle_message(incoming)331 mock_queue.create_tree.assert_awaited_once()332 333 334@pytest.mark.asyncio335async def test_process_parsed_event_malformed_content_continues():336 """Malformed/unknown parsed event does not crash _process_parsed_event."""337 platform = MagicMock()338 platform.queue_edit_message = AsyncMock()339 340 cli_manager = MagicMock()341 session_store = MagicMock()342 handler = ClaudeMessageHandler(platform, cli_manager, session_store)343 344 transcript = MagicMock()345 update_ui = AsyncMock()346 347 last_status, had = await handler._process_parsed_event(348 parsed={"type": "unknown_type"},349 transcript=transcript,350 update_ui=update_ui,351 last_status=None,352 had_transcript_events=False,353 tree=None,354 node_id="n1",355 captured_session_id=None,356 )357 assert last_status is None358 assert had is False359 360 361@pytest.mark.asyncio362async def test_handler_update_ui_edit_failure_does_not_crash():363 """When queue_edit_message raises during streaming, _process_node continues and completes."""364 platform = MagicMock()365 platform.queue_edit_message = AsyncMock(366 side_effect=RuntimeError("Telegram API error")367 )368 platform.fire_and_forget = MagicMock(369 side_effect=lambda c: getattr(c, "close", lambda: None)()370 )371 372 async def _mock_start_task(*args, **kwargs):373 yield {374 "type": "content_block_delta",375 "index": 0,376 "delta": {"type": "text_delta", "text": "Hello"},377 }378 yield {379 "type": "content_block_delta",380 "index": 0,381 "delta": {"type": "text_delta", "text": " world"},382 }383 yield {"type": "complete", "status": "success"}384 385 mock_session = MagicMock()386 mock_session.start_task = _mock_start_task387 cli_manager = MagicMock()388 cli_manager.get_or_create_session = AsyncMock(389 return_value=(mock_session, "s1", False)390 )391 cli_manager.remove_session = AsyncMock()392 cli_manager.get_stats.return_value = {"active_sessions": 0}393 394 session_store = MagicMock()395 handler = ClaudeMessageHandler(platform, cli_manager, session_store)396 mock_queue = MagicMock()397 mock_queue.get_tree_for_node.return_value = None398 handler.replace_tree_queue(mock_queue)399 400 incoming = IncomingMessage(401 text="hi",402 chat_id="c",403 user_id="u",404 message_id="n1",405 platform="telegram",406 )407 node = MessageNode(node_id="n1", incoming=incoming, status_message_id="s1")408 409 await handler._process_node("n1", node)410 411 cli_manager.remove_session.assert_awaited_once()412 