Jack1808/Claude_Code

1import asyncio2from unittest.mock import MagicMock3 4import pytest5 6from messaging.handler import ClaudeMessageHandler7from messaging.trees.data import MessageState8 9 10@pytest.fixture11def handler_integration(mock_platform, mock_cli_manager, mock_session_store):12    # Use real TreeQueueManager13    handler = ClaudeMessageHandler(mock_platform, mock_cli_manager, mock_session_store)14    return handler15 16 17async def mock_async_gen(events):18    for e in events:19        yield e20 21 22@pytest.mark.asyncio23async def test_full_conversation_flow_single_user(24    handler_integration, mock_platform, mock_cli_manager, incoming_message_factory25):26    # 1. First message27    msg1 = incoming_message_factory(text="message 1", message_id="m1")28    mock_platform.queue_send_message.return_value = "s1"29 30    # Mock CLI session for m131    mock_session1 = MagicMock()32    mock_session1.start_task.return_value = mock_async_gen(33        [34            {"type": "session_info", "session_id": "sess1"},35            {36                "type": "assistant",37                "message": {"content": [{"type": "text", "text": "Reply 1"}]},38            },39            {"type": "exit", "code": 0, "stderr": None},40        ]41    )42    mock_cli_manager.get_or_create_session.return_value = (43        mock_session1,44        "pending_1",45        True,46    )47 48    await handler_integration.handle_message(msg1)49 50    # Wait for processing51    tree = handler_integration.tree_queue.get_tree_for_node("m1")52    for _ in range(10):53        if tree.get_node("m1").state.value == MessageState.COMPLETED.value:54            break55        await asyncio.sleep(0.01)56 57    assert tree.get_node("m1").state.value == MessageState.COMPLETED.value58    assert tree.get_node("m1").session_id == "sess1"59    mock_session1.start_task.assert_called_with(60        "message 1", session_id=None, fork_session=False61    )62 63    # 2. Reply to m164    msg2 = incoming_message_factory(65        text="message 2", message_id="m2", reply_to_message_id="m1"66    )67    mock_platform.queue_send_message.return_value = "s2"68 69    # Mock CLI session for m270    mock_session2 = MagicMock()71    mock_session2.start_task.return_value = mock_async_gen(72        [73            {"type": "session_info", "session_id": "sess2"},74            {75                "type": "assistant",76                "message": {"content": [{"type": "text", "text": "Reply 2"}]},77            },78            {"type": "exit", "code": 0, "stderr": None},79        ]80    )81    mock_cli_manager.get_or_create_session.reset_mock()82    mock_cli_manager.get_or_create_session.return_value = (83        mock_session2,84        "pending_2",85        True,86    )87 88    await handler_integration.handle_message(msg2)89 90    # Wait for processing91    for _ in range(10):92        if tree.get_node("m2").state.value == MessageState.COMPLETED.value:93            break94        await asyncio.sleep(0.01)95 96    assert tree.get_node("m2").state.value == MessageState.COMPLETED.value97    assert tree.get_node("m2").parent_id == "m1"98    mock_cli_manager.get_or_create_session.assert_called_with(session_id=None)99    mock_session2.start_task.assert_called_with(100        "message 2", session_id="sess1", fork_session=True101    )102 103 104@pytest.mark.asyncio105async def test_error_propagation_chain(106    handler_integration, mock_platform, mock_cli_manager, incoming_message_factory107):108    msg1 = incoming_message_factory(text="m1", message_id="m1")109    mock_platform.queue_send_message.return_value = "s1"110 111    mock_session1 = MagicMock()112    mock_session1.start_task.return_value = mock_async_gen(113        [{"type": "error", "error": {"message": "failed"}}]114    )115    mock_cli_manager.get_or_create_session.return_value = (116        mock_session1,117        "sess1",118        False,119    )120 121    await handler_integration.handle_message(msg1)122    tree = handler_integration.tree_queue.get_tree_for_node("m1")123 124    msg2 = incoming_message_factory(125        text="m2", message_id="m2", reply_to_message_id="m1"126    )127    await handler_integration.handle_message(msg2)128 129    # Wait for m1 to fail130    for _ in range(20):131        if tree.get_node("m1").state.value == MessageState.ERROR.value:132            break133        await asyncio.sleep(0.01)134 135    # Give a tiny bit of time for propagation and skipping in processor136    await asyncio.sleep(0.05)137 138    assert tree.get_node("m1").state.value == MessageState.ERROR.value139    assert tree.get_node("m2").state.value == MessageState.ERROR.value140    assert "Parent failed" in tree.get_node("m2").error_message141 142 143@pytest.mark.asyncio144async def test_concurrent_replies_to_different_trees(145    handler_integration, mock_platform, mock_cli_manager, incoming_message_factory146):147    msg1 = incoming_message_factory(text="t1", message_id="t1")148    msg2 = incoming_message_factory(text="t2", message_id="t2")149 150    mock_session1 = MagicMock()151    mock_session1.start_task.return_value = mock_async_gen(152        [{"type": "exit", "code": 0}]153    )154    mock_session2 = MagicMock()155    mock_session2.start_task.return_value = mock_async_gen(156        [{"type": "exit", "code": 0}]157    )158 159    mock_cli_manager.get_or_create_session.side_effect = [160        (mock_session1, "s1", False),161        (mock_session2, "s2", False),162    ]163 164    await handler_integration.handle_message(msg1)165    await handler_integration.handle_message(msg2)166 167    # Wait for both168    for _ in range(20):169        node1 = handler_integration.tree_queue.get_node("t1")170        node2 = handler_integration.tree_queue.get_node("t2")171        if (172            node1173            and node2174            and node1.state.value == MessageState.COMPLETED.value175            and node2.state.value == MessageState.COMPLETED.value176        ):177            break178        await asyncio.sleep(0.01)179 180    assert (181        handler_integration.tree_queue.get_node("t1").state.value182        == MessageState.COMPLETED.value183    )184    assert (185        handler_integration.tree_queue.get_node("t2").state.value186        == MessageState.COMPLETED.value187    )188