Jack1808/Claude_Code

1from unittest.mock import MagicMock2 3import pytest4 5from messaging.models import IncomingMessage6from messaging.trees.data import MessageNode, MessageState, MessageTree7from messaging.trees.repository import TreeRepository8 9 10@pytest.fixture11def repository():12    return TreeRepository()13 14 15@pytest.fixture16def sample_tree():17    incoming = IncomingMessage(18        text="root",19        chat_id="c1",20        user_id="u1",21        message_id="root_id",22        platform="telegram",23    )24    node = MessageNode(node_id="root_id", incoming=incoming, status_message_id="s1")25    return MessageTree(node)26 27 28def test_add_and_get_tree(repository, sample_tree):29    repository.add_tree("root_id", sample_tree)30 31    assert repository.get_tree("root_id") == sample_tree32    assert repository.get_tree_for_node("root_id") == sample_tree33    assert repository.has_node("root_id")34 35 36def test_get_tree_nonexistent(repository):37    assert repository.get_tree("none") is None38    assert repository.get_tree_for_node("none") is None39 40 41def test_register_node(repository, sample_tree):42    repository.add_tree("root_id", sample_tree)43    repository.register_node("child_id", "root_id")44 45    assert repository.get_tree_for_node("child_id") == sample_tree46    assert repository.has_node("child_id")47 48 49def test_get_node(repository, sample_tree):50    repository.add_tree("root_id", sample_tree)51    node = repository.get_node("root_id")52 53    assert node is not None54    assert node.node_id == "root_id"55    assert repository.get_node("none") is None56 57 58def test_is_tree_busy(repository, sample_tree):59    repository.add_tree("root_id", sample_tree)60    assert repository.is_tree_busy("root_id") is False61 62    sample_tree._is_processing = True63    assert repository.is_tree_busy("root_id") is True64    assert repository.is_node_tree_busy("root_id") is True65 66 67def test_get_queue_size(repository, sample_tree):68    repository.add_tree("root_id", sample_tree)69    assert repository.get_queue_size("root_id") == 070 71    # We can't easily put items in asyncio.Queue without async,72    # but we can mock it for this unit test if needed, or just skip if it's too complex.73    # Actually, we can use a mock queue since this is a unit test of the repository wrapper.74    sample_tree._queue = MagicMock()75    sample_tree._queue.qsize.return_value = 576 77    assert repository.get_queue_size("root_id") == 578 79 80def test_resolve_parent_node_id(repository, sample_tree):81    repository.add_tree("root_id", sample_tree)82    repository.register_node("s1", "root_id")83 84    # 1. Direct node match85    assert repository.resolve_parent_node_id("root_id") == "root_id"86 87    # 2. Status message match88    # find_node_by_status_message is used inside resolve_parent_node_id89    # sample_tree has root_id node with status_message_id "s1"90    assert repository.resolve_parent_node_id("s1") == "root_id"91 92    # 3. No match93    assert repository.resolve_parent_node_id("unknown") is None94 95 96def test_get_pending_children(repository, sample_tree):97    repository.add_tree("root_id", sample_tree)98 99    # Create a child node100    child_incoming = IncomingMessage(101        text="child",102        chat_id="c1",103        user_id="u1",104        message_id="child_id",105        platform="telegram",106    )107    child_node = MessageNode(108        node_id="child_id",109        incoming=child_incoming,110        status_message_id="s2",111        parent_id="root_id",112        state=MessageState.PENDING,113    )114 115    sample_tree._nodes["child_id"] = child_node116    sample_tree.get_node("root_id").children_ids.append("child_id")117    repository.register_node("child_id", "root_id")118 119    pending = repository.get_pending_children("root_id")120    assert len(pending) == 1121    assert pending[0].node_id == "child_id"122 123 124def test_to_from_dict(repository, sample_tree):125    repository.add_tree("root_id", sample_tree)126    data = repository.to_dict()127 128    assert "trees" in data129    assert "root_id" in data["trees"]130    assert "node_to_tree" in data131    assert data["node_to_tree"]["root_id"] == "root_id"132 133    new_repo = TreeRepository.from_dict(data)134    tree = new_repo.get_tree("root_id")135    assert tree is not None136    assert tree.root_id == "root_id"137    assert new_repo.get_tree_for_node("root_id") == tree138 139 140def test_all_trees(repository, sample_tree):141    repository.add_tree("root_id", sample_tree)142    assert len(repository.all_trees()) == 1143    assert repository.tree_ids() == ["root_id"]144