wehub-resource-sync
188 lines
6.5 KiB
Python
188 lines
6.5 KiB
Python
"""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() == ()
|