Jack1808/Claude_Code

1from unittest.mock import AsyncMock, MagicMock2 3import pytest4 5from messaging.handler import ClaudeMessageHandler6from messaging.trees.data import MessageState7 8 9async def _gen_session(events):10    for e in events:11        yield e12 13 14@pytest.fixture15def handler(mock_platform, mock_cli_manager, mock_session_store):16    return ClaudeMessageHandler(mock_platform, mock_cli_manager, mock_session_store)17 18 19@pytest.mark.asyncio20async def test_sibling_replies_fork_from_parent_session_id(21    handler, mock_cli_manager, incoming_message_factory22):23    # Root node A with a known session_id.24    root_incoming = incoming_message_factory(text="A", message_id="A")25    tree = await handler.tree_queue.create_tree(26        node_id="A", incoming=root_incoming, status_message_id="status_A"27    )28    await tree.update_state("A", MessageState.COMPLETED, session_id="sess_A")29 30    # Add two sibling replies R1 and R2 under A.31    r1_incoming = incoming_message_factory(32        text="R1", message_id="R1", reply_to_message_id="A"33    )34    r2_incoming = incoming_message_factory(35        text="R2", message_id="R2", reply_to_message_id="A"36    )37    _, r1_node = await handler.tree_queue.add_to_tree(38        "A", "R1", r1_incoming, "status_R1"39    )40    _, r2_node = await handler.tree_queue.add_to_tree(41        "A", "R2", r2_incoming, "status_R2"42    )43 44    # Mock a fresh cli_session per node.45    calls = []46 47    async def _get_or_create_session(session_id=None):48        cli_session = MagicMock()49 50        async def _start_task(prompt, session_id=None, fork_session=False):51            calls.append((prompt, session_id, fork_session))52            child_sid = f"sess_{prompt}"53            async for ev in _gen_session(54                [55                    {"type": "session_info", "session_id": child_sid},56                    {"type": "exit", "code": 0, "stderr": None},57                ]58            ):59                yield ev60 61        cli_session.start_task = _start_task62        return cli_session, f"pending_{len(calls) + 1}", True63 64    mock_cli_manager.get_or_create_session = AsyncMock(65        side_effect=_get_or_create_session66    )67 68    await handler._process_node("R1", r1_node)69    await handler._process_node("R2", r2_node)70 71    # Both siblings must resume from the same parent session and fork.72    assert calls[0][0] == "R1"73    assert calls[0][1] == "sess_A"74    assert calls[0][2] is True75 76    assert calls[1][0] == "R2"77    assert calls[1][1] == "sess_A"78    assert calls[1][2] is True79 80 81@pytest.mark.asyncio82async def test_grandchild_reply_forks_from_branch_session(83    handler, mock_cli_manager, incoming_message_factory84):85    root_incoming = incoming_message_factory(text="A", message_id="A")86    tree = await handler.tree_queue.create_tree(87        node_id="A", incoming=root_incoming, status_message_id="status_A"88    )89    await tree.update_state("A", MessageState.COMPLETED, session_id="sess_A")90 91    r1_incoming = incoming_message_factory(92        text="R1", message_id="R1", reply_to_message_id="A"93    )94    _, r1_node = await handler.tree_queue.add_to_tree(95        "A", "R1", r1_incoming, "status_R1"96    )97 98    calls = []99 100    async def _get_or_create_session(session_id=None):101        cli_session = MagicMock()102 103        async def _start_task(prompt, session_id=None, fork_session=False):104            calls.append((prompt, session_id, fork_session))105            # R1 gets its own forked session id.106            child_sid = "sess_R1"107            async for ev in _gen_session(108                [109                    {"type": "session_info", "session_id": child_sid},110                    {"type": "exit", "code": 0, "stderr": None},111                ]112            ):113                yield ev114 115        cli_session.start_task = _start_task116        return cli_session, "pending_R1", True117 118    mock_cli_manager.get_or_create_session = AsyncMock(119        side_effect=_get_or_create_session120    )121 122    await handler._process_node("R1", r1_node)123    assert r1_node.session_id == "sess_R1"124 125    # Grandchild C1 replies to R1 and must fork from sess_R1, not sess_A.126    c1_incoming = incoming_message_factory(127        text="C1", message_id="C1", reply_to_message_id="R1"128    )129    _, c1_node = await handler.tree_queue.add_to_tree(130        "R1", "C1", c1_incoming, "status_C1"131    )132 133    async def _get_or_create_session_c1(session_id=None):134        cli_session = MagicMock()135 136        async def _start_task(prompt, session_id=None, fork_session=False):137            calls.append((prompt, session_id, fork_session))138            async for ev in _gen_session(139                [140                    {"type": "session_info", "session_id": "sess_C1"},141                    {"type": "exit", "code": 0, "stderr": None},142                ]143            ):144                yield ev145 146        cli_session.start_task = _start_task147        return cli_session, "pending_C1", True148 149    mock_cli_manager.get_or_create_session = AsyncMock(150        side_effect=_get_or_create_session_c1151    )152 153    await handler._process_node("C1", c1_node)154 155    # The last call should be for C1 and must resume from sess_R1.156    assert calls[-1][0] == "C1"157    assert calls[-1][1] == "sess_R1"158    assert calls[-1][2] is True159