Files
wehub-resource-sync 5296d0e97c
CI / Ban suppressions and legacy annotations (push) Has been cancelled
CI / pytest (push) Has been cancelled
CI / ruff-check (push) Has been cancelled
CI / ruff-format (push) Has been cancelled
CI / ty (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:35:44 +08:00

259 lines
8.2 KiB
Python

import asyncio
from dataclasses import FrozenInstanceError
import pytest
from free_claude_code.messaging.models import MessageScope
from free_claude_code.messaging.trees.node import MessageNode, MessageState
from free_claude_code.messaging.trees.runtime import MessageTree
_SCOPE = MessageScope(platform="telegram", chat_id="chat")
def _tree() -> MessageTree:
return MessageTree(
MessageNode(
node_id="root",
scope=_SCOPE,
prompt="prompt root",
status_message_id="status-root",
)
)
async def _add(
tree: MessageTree,
node_id: str,
status_message_id: str,
parent_id: str = "root",
parent_reference_id: str | None = None,
):
return await tree.add_and_enqueue(
node_id,
_SCOPE,
f"prompt {node_id}",
status_message_id,
parent_id,
parent_reference_id or parent_id,
)
@pytest.mark.asyncio
async def test_enqueue_finish_and_claim_next_are_atomic_and_fifo() -> None:
tree = _tree()
root = await tree.enqueue_or_claim("root")
child_a = await _add(tree, "child-a", "status-a")
child_b = await _add(tree, "child-b", "status-b")
assert root.accepted and root.claim is not None and root.position is None
assert child_a.claim is None and child_a.position == 1
assert child_b.claim is None and child_b.position == 2
completion = await tree.finish_and_claim_next(root.claim.claim_id)
assert completion.next_claim is not None
assert completion.next_claim.node.node_id == "child-a"
assert [(entry.node.node_id, entry.position) for entry in completion.queue] == [
("child-b", 1)
]
@pytest.mark.asyncio
async def test_duplicate_unknown_and_terminal_admission_are_rejected_without_wedge() -> (
None
):
tree = _tree()
first = await tree.enqueue_or_claim("root")
duplicate = await tree.enqueue_or_claim("root")
unknown = await tree.enqueue_or_claim("status-root")
assert first.claim is not None
assert duplicate.accepted is False
assert unknown.accepted is False
assert duplicate.snapshot is None
assert unknown.snapshot is None
await tree.complete_claim(first.claim.claim_id, "session")
await tree.finish_and_claim_next(first.claim.claim_id)
terminal = await tree.enqueue_or_claim("root")
assert terminal.accepted is False
assert terminal.snapshot is None
@pytest.mark.asyncio
async def test_stale_claim_id_cannot_clear_a_newer_claim() -> None:
tree = _tree()
first = await tree.enqueue_or_claim("root")
assert first.claim is not None
await tree.finish_and_claim_next(first.claim.claim_id)
second = await _add(tree, "child", "status-child")
assert second.claim is not None
stale = await tree.finish_and_claim_next(first.claim.claim_id)
assert stale.next_claim is None
assert await tree.record_session(second.claim.claim_id, "session") is not None
await tree.finish_and_claim_next(second.claim.claim_id)
third = await _add(tree, "third", "status-third")
assert third.claim is not None
@pytest.mark.asyncio
async def test_cancellation_keeps_active_identity_until_matching_finish() -> None:
tree = _tree()
root = await tree.enqueue_or_claim("root")
queued = await _add(tree, "queued", "status-queued")
assert root.claim is not None and queued.position == 1
cancelled = await tree.cancel_node("root")
assert cancelled.active_claim == root.claim
assert [node.node_id for node in cancelled.nodes] == ["root"]
completion = await tree.finish_and_claim_next(root.claim.claim_id)
assert completion.next_claim is not None
assert completion.next_claim.node.node_id == "queued"
@pytest.mark.asyncio
@pytest.mark.parametrize("propagate", [False, True])
async def test_cancelled_claim_rejects_late_failure_and_descendant_propagation(
propagate: bool,
) -> None:
tree = _tree()
root = await tree.enqueue_or_claim("root")
queued = await _add(tree, "queued", "status-queued")
assert root.claim is not None and queued.position == 1
await tree.cancel_node("root")
late_failure = await tree.fail_claim(
root.claim.claim_id,
propagate=propagate,
)
assert late_failure.snapshot is None
assert late_failure.affected == ()
assert late_failure.queue_update is None
queued_view = await tree.node_view("queued")
assert queued_view is not None
assert queued_view.state is MessageState.PENDING
completion = await tree.finish_and_claim_next(root.claim.claim_id)
assert completion.next_claim is not None
assert completion.next_claim.node.node_id == "queued"
@pytest.mark.asyncio
async def test_cancelled_queue_member_is_never_claimed() -> None:
tree = _tree()
root = await tree.enqueue_or_claim("root")
await _add(tree, "a", "status-a")
await _add(tree, "b", "status-b")
assert root.claim is not None
cancelled = await tree.cancel_node("a")
assert [node.node_id for node in cancelled.nodes] == ["a"]
completion = await tree.finish_and_claim_next(root.claim.claim_id)
assert completion.next_claim is not None
assert completion.next_claim.node.node_id == "b"
@pytest.mark.asyncio
async def test_prompt_subtree_removal_returns_every_platform_message() -> None:
tree = _tree()
await _add(tree, "child", "status-child")
removed = await tree.remove_message_subtree("root")
assert removed.removed_message_ids == frozenset(
{"root", "status-root", "child", "status-child"}
)
@pytest.mark.asyncio
async def test_status_subtree_preserves_prompt_siblings_and_invalidates_resume() -> (
None
):
tree = _tree()
root = await tree.enqueue_or_claim("root")
assert root.claim is not None
await tree.record_session(root.claim.claim_id, "session-root")
direct_sibling = await _add(tree, "sibling", "status-sibling")
status_descendant = await _add(
tree,
"descendant",
"status-descendant",
parent_reference_id="status-root",
)
assert direct_sibling.position == 1
assert status_descendant.position == 2
removed = await tree.remove_message_subtree("status-root")
assert removed.removed_message_ids == frozenset(
{"status-root", "descendant", "status-descendant"}
)
root_view = await tree.node_view("root")
sibling_view = await tree.node_view("sibling")
assert root_view is not None
assert root_view.state is MessageState.ERROR
assert root_view.session_id is None
assert sibling_view is not None
assert sibling_view.state is MessageState.PENDING
assert await tree.node_view("descendant") is None
assert await tree.resolve_reply("status-root") is None
completion = await tree.finish_and_claim_next(root.claim.claim_id)
assert completion.next_claim is not None
assert completion.next_claim.node.node_id == "sibling"
assert completion.next_claim.parent_session_id is None
@pytest.mark.asyncio
async def test_concurrent_duplicate_enqueue_produces_one_claim_only() -> None:
tree = _tree()
gate = asyncio.Event()
async def enqueue():
await gate.wait()
return await tree.enqueue_or_claim("root")
tasks = [asyncio.create_task(enqueue()) for _ in range(20)]
gate.set()
results = await asyncio.gather(*tasks)
assert sum(result.claim is not None for result in results) == 1
assert sum(result.accepted for result in results) == 1
@pytest.mark.asyncio
async def test_results_are_frozen_and_do_not_expose_mutable_nodes() -> None:
tree = _tree()
decision = await tree.enqueue_or_claim("root")
assert decision.claim is not None
with pytest.raises(FrozenInstanceError):
decision.claim.__setattr__("claim_id", "replacement")
assert not hasattr(decision.claim, "__dict__")
assert not isinstance(decision.claim.node, MessageNode)
def test_message_tree_has_no_lock_or_partial_mutation_escape_hatches() -> None:
forbidden = {
"with_lock",
"enqueue",
"dequeue",
"put_queue_unlocked",
"remove_from_queue",
"set_processing_state",
"clear_current_node",
"set_current_task",
"cancel_current_task",
"set_node_error_sync",
"drain_queue_and_mark_cancelled",
"reset_processing_state",
}
assert forbidden.isdisjoint(vars(MessageTree))