Jack1808/Claude_Code

1import asyncio2import contextlib3from unittest.mock import AsyncMock, MagicMock4 5import pytest6 7from messaging.models import IncomingMessage8from messaging.trees.data import MessageNode, MessageState, MessageTree9from messaging.trees.processor import TreeQueueProcessor10 11 12@pytest.fixture13def tree_processor():14    return TreeQueueProcessor()15 16 17@pytest.fixture18def sample_incoming():19    return IncomingMessage(20        text="test message",21        chat_id="chat123",22        user_id="user456",23        message_id="msg789",24        platform="telegram",25    )26 27 28@pytest.fixture29def sample_node(sample_incoming):30    return MessageNode(31        node_id="msg789", incoming=sample_incoming, status_message_id="status123"32    )33 34 35@pytest.fixture36def sample_tree(sample_node):37    return MessageTree(sample_node)38 39 40@pytest.mark.asyncio41async def test_process_node_success(tree_processor, sample_tree, sample_node):42    processor = AsyncMock()43 44    await tree_processor.process_node(sample_tree, sample_node, processor)45 46    processor.assert_called_once_with(sample_node.node_id, sample_node)47    assert sample_tree._current_node_id is None48 49 50@pytest.mark.asyncio51async def test_process_node_cancelled(tree_processor, sample_tree, sample_node):52    processor = AsyncMock(side_effect=asyncio.CancelledError)53 54    with pytest.raises(asyncio.CancelledError):55        await tree_processor.process_node(sample_tree, sample_node, processor)56 57    assert sample_tree._current_node_id is None58 59 60@pytest.mark.asyncio61async def test_process_node_exception(tree_processor, sample_tree, sample_node):62    processor = AsyncMock(side_effect=Exception("Test error"))63 64    # We need to mock update_state to verify it was called65    sample_tree.update_state = AsyncMock()66 67    await tree_processor.process_node(sample_tree, sample_node, processor)68 69    sample_tree.update_state.assert_called_once_with(70        sample_node.node_id, MessageState.ERROR, error_message="Test error"71    )72    assert sample_tree._current_node_id is None73 74 75@pytest.mark.asyncio76async def test_enqueue_and_start_when_free(tree_processor, sample_tree):77    processor = AsyncMock()78    node_id = "node1"79 80    # Mock get_node to return a node81    node = MagicMock(spec=MessageNode)82    sample_tree.get_node = MagicMock(return_value=node)83 84    was_queued = await tree_processor.enqueue_and_start(sample_tree, node_id, processor)85 86    assert was_queued is False87    assert sample_tree._is_processing is True88    assert sample_tree._current_node_id == node_id89    assert sample_tree._current_task is not None90 91    # Clean up task92    sample_tree._current_task.cancel()93    with contextlib.suppress(asyncio.CancelledError):94        await sample_tree._current_task95 96 97@pytest.mark.asyncio98async def test_enqueue_and_start_when_busy(tree_processor, sample_tree):99    processor = AsyncMock()100    sample_tree._is_processing = True101    node_id = "node1"102 103    was_queued = await tree_processor.enqueue_and_start(sample_tree, node_id, processor)104 105    assert was_queued is True106    assert sample_tree._queue.qsize() == 1107    assert sample_tree._queue.get_nowait() == node_id108 109 110def test_cancel_current_task(tree_processor, sample_tree):111    mock_task = MagicMock(spec=asyncio.Task)112    mock_task.done.return_value = False113    sample_tree._current_task = mock_task114 115    cancelled = tree_processor.cancel_current(sample_tree)116 117    assert cancelled is True118    mock_task.cancel.assert_called_once()119 120 121def test_cancel_current_task_already_done(tree_processor, sample_tree):122    mock_task = MagicMock(spec=asyncio.Task)123    mock_task.done.return_value = True124    sample_tree._current_task = mock_task125 126    cancelled = tree_processor.cancel_current(sample_tree)127 128    assert cancelled is False129    mock_task.cancel.assert_not_called()130 131 132@pytest.mark.asyncio133async def test_process_next_queue_empty(tree_processor, sample_tree):134    processor = AsyncMock()135    sample_tree._is_processing = True136 137    await tree_processor._process_next(sample_tree, processor)138 139    assert sample_tree._is_processing is False140 141 142@pytest.mark.asyncio143async def test_process_next_with_item(tree_processor, sample_tree):144    processor = AsyncMock()145    await sample_tree._queue.put("next_node")146 147    node = MagicMock(spec=MessageNode)148    sample_tree.get_node = MagicMock(return_value=node)149 150    await tree_processor._process_next(sample_tree, processor)151 152    assert sample_tree._current_node_id == "next_node"153    assert sample_tree._current_task is not None154 155    # Clean up156    sample_tree._current_task.cancel()157    with contextlib.suppress(asyncio.CancelledError):158        await sample_tree._current_task159 160 161@pytest.mark.asyncio162async def test_process_next_triggers_queue_update(sample_tree):163    callback = AsyncMock()164    processor = TreeQueueProcessor(queue_update_callback=callback)165 166    await sample_tree._queue.put("next_node")167    sample_tree.get_node = MagicMock(return_value=None)168 169    await processor._process_next(sample_tree, AsyncMock())170 171    callback.assert_awaited_once_with(sample_tree)172 173 174@pytest.mark.asyncio175async def test_process_next_triggers_node_started(sample_tree):176    node_started = AsyncMock()177    processor = TreeQueueProcessor(node_started_callback=node_started)178 179    await sample_tree._queue.put("next_node")180    sample_tree.get_node = MagicMock(return_value=None)181 182    await processor._process_next(sample_tree, AsyncMock())183 184    node_started.assert_awaited_once_with(sample_tree, "next_node")185