tests/messaging/test_tree_queue.py

"""Focused tests for the tree package's private graph and queue values."""

import pytest

from free_claude_code.messaging.models import MessageScope
from free_claude_code.messaging.trees.graph import MessageTreeGraph
from free_claude_code.messaging.trees.node import MessageNode, MessageState
from free_claude_code.messaging.trees.queue import MessageNodeQueue
from free_claude_code.messaging.trees.snapshot import (
    TreeSnapshot,
    node_from_snapshot,
    node_to_snapshot,
)

_SCOPE = MessageScope(platform="telegram", chat_id="chat")


def _root() -> MessageNode:
    return MessageNode(
        node_id="root",
        scope=_SCOPE,
        prompt="prompt root",
        status_message_id="status-root",
    )


def _add(
    graph: MessageTreeGraph,
    node_id: str,
    status_message_id: str,
    parent_id: str,
) -> MessageNode:
    return graph.add_node(
        node_id=node_id,
        scope=_SCOPE,
        prompt=f"prompt {node_id}",
        status_message_id=status_message_id,
        parent_id=parent_id,
        parent_reference_id=parent_id,
    )


def test_node_snapshot_round_trip_preserves_execution_state_only() -> None:
    node = _root()
    node.children_ids.extend(["child-a", "child-b"])
    node.update_state(MessageState.ERROR)
    node.update_state(MessageState.COMPLETED, session_id="session-root")

    snapshot = node_to_snapshot(node)
    restored = node_from_snapshot(snapshot, _SCOPE)

    assert restored.node_id == "root"
    assert restored.scope == _SCOPE
    assert restored.prompt == ""
    assert restored.status_message_id == "status-root"
    assert restored.state is MessageState.COMPLETED
    assert restored.session_id == "session-root"
    assert restored.parent_reference_id is None
    assert restored.children_ids == []
    assert "children_ids" not in snapshot
    assert "created_at" not in snapshot
    assert "completed_at" not in snapshot
    assert "error_message" not in snapshot


def test_graph_snapshot_round_trip_preserves_links_and_status_lookup() -> None:
    graph = MessageTreeGraph(_root())
    child = _add(graph, "child", "status-child", "root")
    _add(graph, "grandchild", "status-grandchild", "child")
    child.update_state(MessageState.COMPLETED, session_id="session-child")

    restored = MessageTreeGraph.from_snapshot(graph.snapshot())
    parent = restored.get_parent("grandchild")
    status_child = restored.find_node_by_status_message("status-child")

    assert restored.root_id == "root"
    assert parent is not None
    assert parent.node_id == "child"
    assert restored.get_parent_session_id("grandchild") == "session-child"
    assert status_child is not None
    assert status_child.node_id == "child"
    assert restored.get_descendants("root") == ["root", "child", "grandchild"]


def test_graph_restore_normalizes_numeric_ids_to_string_references() -> None:
    graph = MessageTreeGraph(_root())
    _add(graph, "child", "status-child", "root")
    snapshot = graph.snapshot()
    snapshot.nodes["child"]["node_id"] = 2
    snapshot.nodes["child"]["status_message_id"] = 123
    snapshot.nodes["child"]["parent_id"] = "root"
    snapshot.nodes["2"] = snapshot.nodes.pop("child")

    restored = MessageTreeGraph.from_snapshot(snapshot)

    child = restored.find_node_by_status_message("123")
    assert child is not None and child.node_id == "2"


def test_graph_restore_normalizes_legacy_parent_to_prompt_reference() -> None:
    graph = MessageTreeGraph(_root())
    _add(graph, "child", "status-child", "root")
    snapshot = graph.snapshot()
    snapshot.nodes["child"].pop("parent_reference_id")

    restored = MessageTreeGraph.from_snapshot(snapshot)

    child = restored.get_node("child")
    assert child is not None
    assert child.parent_reference_id == "root"
    assert restored.snapshot().nodes["child"]["parent_reference_id"] == "root"


def test_cleared_status_round_trip_preserves_prompt_anchor() -> None:
    graph = MessageTreeGraph(_root())
    graph.clear_status("root")

    restored = MessageTreeGraph.from_snapshot(graph.snapshot())

    root = restored.get_root()
    assert root.status_message_id is None
    assert root.state is MessageState.ERROR
    assert restored.resolve_reference("root") is not None
    assert restored.resolve_reference("status-root") is None


def test_snapshot_rejects_runnable_node_without_status() -> None:
    snapshot = MessageTreeGraph(_root()).snapshot()
    snapshot.nodes["root"]["status_message_id"] = None

    with pytest.raises(ValueError, match="requires a status"):
        MessageTreeGraph.from_snapshot(snapshot)


def test_graph_rejects_duplicate_node_and_status_identity() -> None:
    graph = MessageTreeGraph(_root())
    _add(graph, "child", "status-child", "root")

    with pytest.raises(ValueError, match="already exists"):
        _add(graph, "child", "status-other", "root")
    with pytest.raises(ValueError, match="already exists"):
        _add(graph, "other", "status-child", "root")
    with pytest.raises(ValueError, match="must be distinct"):
        _add(graph, "same", "same", "root")


def test_graph_reference_subtree_can_remove_exact_nodes_and_status_lookups() -> None:
    graph = MessageTreeGraph(_root())
    _add(graph, "branch", "status-branch", "root")
    _add(graph, "leaf", "status-leaf", "branch")
    _add(graph, "sibling", "status-sibling", "root")

    references = graph.get_reference_descendants("branch")
    graph.remove_nodes(
        {reference for reference in references if not reference.startswith("status-")}
    )

    assert graph.get_node("branch") is None
    assert graph.get_node("leaf") is None
    assert graph.find_node_by_status_message("status-branch") is None
    assert graph.find_node_by_status_message("status-leaf") is None
    assert graph.get_descendants("root") == ["root", "sibling"]


def test_tree_snapshot_rejects_invalid_wire_shapes() -> None:
    assert TreeSnapshot.from_json(None) is None
    assert TreeSnapshot.from_json({"root_id": "root", "nodes": []}) is None
    assert TreeSnapshot.from_json({"nodes": {}}) is None


def test_node_queue_is_unique_fifo_and_supports_atomic_removal() -> None:
    queue = MessageNodeQueue()

    assert queue.put("a") is True
    assert queue.put("b") is True
    assert queue.put("a") is False
    assert queue.items() == ("a", "b")
    assert queue.remove("a") is True
    assert queue.remove("a") is False
    assert queue.items() == ("b",)
    assert queue.pop() == "b"
    assert queue.pop() is None

    assert queue.put("c") is True
    assert queue.put("d") is True
    assert queue.drain() == ("c", "d")
    assert queue.items() == ()