tests/messaging/test_tree_repository.py

from unittest.mock import MagicMock

import pytest

from messaging.models import IncomingMessage
from messaging.trees.data import MessageNode, MessageState, MessageTree
from messaging.trees.repository import TreeRepository


@pytest.fixture
def repository():
    return TreeRepository()


@pytest.fixture
def sample_tree():
    incoming = IncomingMessage(
        text="root",
        chat_id="c1",
        user_id="u1",
        message_id="root_id",
        platform="telegram",
    )
    node = MessageNode(node_id="root_id", incoming=incoming, status_message_id="s1")
    return MessageTree(node)


def test_add_and_get_tree(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)

    assert repository.get_tree("root_id") == sample_tree
    assert repository.get_tree_for_node("root_id") == sample_tree
    assert repository.has_node("root_id")


def test_get_tree_nonexistent(repository):
    assert repository.get_tree("none") is None
    assert repository.get_tree_for_node("none") is None


def test_register_node(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)
    repository.register_node("child_id", "root_id")

    assert repository.get_tree_for_node("child_id") == sample_tree
    assert repository.has_node("child_id")


def test_get_node(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)
    node = repository.get_node("root_id")

    assert node is not None
    assert node.node_id == "root_id"
    assert repository.get_node("none") is None


def test_is_tree_busy(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)
    assert repository.is_tree_busy("root_id") is False

    sample_tree._is_processing = True
    assert repository.is_tree_busy("root_id") is True
    assert repository.is_node_tree_busy("root_id") is True


def test_get_queue_size(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)
    assert repository.get_queue_size("root_id") == 0

    # We can't easily put items in asyncio.Queue without async,
    # but we can mock it for this unit test if needed, or just skip if it's too complex.
    # Actually, we can use a mock queue since this is a unit test of the repository wrapper.
    sample_tree._queue = MagicMock()
    sample_tree._queue.qsize.return_value = 5

    assert repository.get_queue_size("root_id") == 5


def test_resolve_parent_node_id(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)
    repository.register_node("s1", "root_id")

    # 1. Direct node match
    assert repository.resolve_parent_node_id("root_id") == "root_id"

    # 2. Status message match
    # find_node_by_status_message is used inside resolve_parent_node_id
    # sample_tree has root_id node with status_message_id "s1"
    assert repository.resolve_parent_node_id("s1") == "root_id"

    # 3. No match
    assert repository.resolve_parent_node_id("unknown") is None


def test_get_pending_children(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)

    # Create a child node
    child_incoming = IncomingMessage(
        text="child",
        chat_id="c1",
        user_id="u1",
        message_id="child_id",
        platform="telegram",
    )
    child_node = MessageNode(
        node_id="child_id",
        incoming=child_incoming,
        status_message_id="s2",
        parent_id="root_id",
        state=MessageState.PENDING,
    )

    sample_tree._nodes["child_id"] = child_node
    sample_tree.get_node("root_id").children_ids.append("child_id")
    repository.register_node("child_id", "root_id")

    pending = repository.get_pending_children("root_id")
    assert len(pending) == 1
    assert pending[0].node_id == "child_id"


def test_to_from_dict(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)
    data = repository.to_dict()

    assert "trees" in data
    assert "root_id" in data["trees"]
    assert "node_to_tree" in data
    assert data["node_to_tree"]["root_id"] == "root_id"

    new_repo = TreeRepository.from_dict(data)
    tree = new_repo.get_tree("root_id")
    assert tree is not None
    assert tree.root_id == "root_id"
    assert new_repo.get_tree_for_node("root_id") == tree


def test_all_trees(repository, sample_tree):
    repository.add_tree("root_id", sample_tree)
    assert len(repository.all_trees()) == 1
    assert repository.tree_ids() == ["root_id"]