Spaces:
Sleeping
Sleeping
| """Concurrency and race condition tests for tree data structures and queue manager.""" | |
| import asyncio | |
| import pytest | |
| from messaging.models import IncomingMessage | |
| from messaging.trees.data import MessageNode, MessageState, MessageTree | |
| from messaging.trees.queue_manager import TreeQueueManager | |
| def _make_incoming(text: str = "hello", msg_id: str = "m1") -> IncomingMessage: | |
| """Create a minimal IncomingMessage for testing.""" | |
| return IncomingMessage( | |
| text=text, | |
| chat_id="chat1", | |
| user_id="user1", | |
| message_id=msg_id, | |
| platform="test", | |
| ) | |
| def _make_tree(root_id: str = "root") -> MessageTree: | |
| """Create a tree with a single root node.""" | |
| root = MessageNode( | |
| node_id=root_id, | |
| incoming=_make_incoming(msg_id=root_id), | |
| status_message_id=f"status_{root_id}", | |
| state=MessageState.PENDING, | |
| ) | |
| return MessageTree(root) | |
| class TestMessageTreeConcurrency: | |
| """Concurrency tests for MessageTree operations.""" | |
| async def test_concurrent_add_node_serialized(self): | |
| """Concurrent add_node calls should all succeed via lock serialization.""" | |
| tree = _make_tree("root") | |
| count = 10 | |
| async def add(i: int): | |
| return await tree.add_node( | |
| node_id=f"child_{i}", | |
| incoming=_make_incoming(msg_id=f"child_{i}"), | |
| status_message_id=f"status_{i}", | |
| parent_id="root", | |
| ) | |
| results = await asyncio.gather(*[add(i) for i in range(count)]) | |
| assert len(results) == count | |
| # All nodes plus root | |
| assert len(tree.all_nodes()) == count + 1 | |
| # Root should have all children | |
| root = tree.get_root() | |
| assert len(root.children_ids) == count | |
| async def test_concurrent_enqueue_dequeue_no_loss(self): | |
| """Concurrent enqueue/dequeue should not lose items.""" | |
| tree = _make_tree("root") | |
| # Add nodes first | |
| for i in range(10): | |
| await tree.add_node( | |
| node_id=f"n{i}", | |
| incoming=_make_incoming(msg_id=f"n{i}"), | |
| status_message_id=f"s{i}", | |
| parent_id="root", | |
| ) | |
| # Enqueue all concurrently | |
| await asyncio.gather(*[tree.enqueue(f"n{i}") for i in range(10)]) | |
| assert tree.get_queue_size() == 10 | |
| # Dequeue all | |
| dequeued = [] | |
| for _ in range(10): | |
| nid = await tree.dequeue() | |
| if nid: | |
| dequeued.append(nid) | |
| assert len(dequeued) == 10 | |
| assert set(dequeued) == {f"n{i}" for i in range(10)} | |
| assert tree.get_queue_size() == 0 | |
| async def test_dequeue_empty_returns_none(self): | |
| """Dequeue on empty queue returns None.""" | |
| tree = _make_tree("root") | |
| result = await tree.dequeue() | |
| assert result is None | |
| async def test_concurrent_update_state(self): | |
| """Concurrent state updates should all apply (last writer wins).""" | |
| tree = _make_tree("root") | |
| for i in range(5): | |
| await tree.add_node( | |
| node_id=f"n{i}", | |
| incoming=_make_incoming(msg_id=f"n{i}"), | |
| status_message_id=f"s{i}", | |
| parent_id="root", | |
| ) | |
| # Update all nodes concurrently | |
| await asyncio.gather( | |
| *[tree.update_state(f"n{i}", MessageState.IN_PROGRESS) for i in range(5)] | |
| ) | |
| for i in range(5): | |
| node = tree.get_node(f"n{i}") | |
| assert node is not None | |
| assert node.state == MessageState.IN_PROGRESS | |
| async def test_update_state_nonexistent_node(self): | |
| """Updating state of a nonexistent node should not raise.""" | |
| tree = _make_tree("root") | |
| # Should just log a warning, not raise | |
| await tree.update_state("nonexistent", MessageState.ERROR) | |
| async def test_add_node_invalid_parent_raises(self): | |
| """Adding a node with nonexistent parent should raise ValueError.""" | |
| tree = _make_tree("root") | |
| with pytest.raises(ValueError, match="not found in tree"): | |
| await tree.add_node( | |
| node_id="child", | |
| incoming=_make_incoming(), | |
| status_message_id="s1", | |
| parent_id="nonexistent", | |
| ) | |
| async def test_queue_snapshot_matches_enqueue_order(self): | |
| """Queue snapshot should return items in FIFO order.""" | |
| tree = _make_tree("root") | |
| for i in range(5): | |
| await tree.add_node( | |
| node_id=f"n{i}", | |
| incoming=_make_incoming(msg_id=f"n{i}"), | |
| status_message_id=f"s{i}", | |
| parent_id="root", | |
| ) | |
| for i in range(5): | |
| await tree.enqueue(f"n{i}") | |
| snapshot = await tree.get_queue_snapshot() | |
| assert snapshot == [f"n{i}" for i in range(5)] | |
| async def test_enqueue_returns_position(self): | |
| """Enqueue should return 1-indexed position.""" | |
| tree = _make_tree("root") | |
| for i in range(3): | |
| await tree.add_node( | |
| node_id=f"n{i}", | |
| incoming=_make_incoming(msg_id=f"n{i}"), | |
| status_message_id=f"s{i}", | |
| parent_id="root", | |
| ) | |
| pos1 = await tree.enqueue("n0") | |
| pos2 = await tree.enqueue("n1") | |
| pos3 = await tree.enqueue("n2") | |
| assert pos1 == 1 | |
| assert pos2 == 2 | |
| assert pos3 == 3 | |
| class TestMessageTreeNavigation: | |
| """Tests for tree navigation methods.""" | |
| async def test_get_children(self): | |
| """get_children returns child nodes.""" | |
| tree = _make_tree("root") | |
| await tree.add_node("c1", _make_incoming(msg_id="c1"), "s1", "root") | |
| await tree.add_node("c2", _make_incoming(msg_id="c2"), "s2", "root") | |
| children = tree.get_children("root") | |
| assert len(children) == 2 | |
| assert {c.node_id for c in children} == {"c1", "c2"} | |
| def test_get_children_nonexistent(self): | |
| """get_children for nonexistent node returns empty list.""" | |
| tree = _make_tree("root") | |
| assert tree.get_children("nonexistent") == [] | |
| def test_get_parent_root(self): | |
| """Root node has no parent.""" | |
| tree = _make_tree("root") | |
| assert tree.get_parent("root") is None | |
| async def test_get_parent_child(self): | |
| """Child node's parent is the root.""" | |
| tree = _make_tree("root") | |
| await tree.add_node("c1", _make_incoming(msg_id="c1"), "s1", "root") | |
| parent = tree.get_parent("c1") | |
| assert parent is not None | |
| assert parent.node_id == "root" | |
| async def test_get_parent_session_id(self): | |
| """get_parent_session_id returns parent's session_id.""" | |
| tree = _make_tree("root") | |
| await tree.update_state("root", MessageState.COMPLETED, session_id="sess_abc") | |
| await tree.add_node("c1", _make_incoming(msg_id="c1"), "s1", "root") | |
| session_id = tree.get_parent_session_id("c1") | |
| assert session_id == "sess_abc" | |
| def test_get_parent_session_id_root(self): | |
| """Root node has no parent session.""" | |
| tree = _make_tree("root") | |
| assert tree.get_parent_session_id("root") is None | |
| def test_has_node(self): | |
| """has_node returns True for existing nodes.""" | |
| tree = _make_tree("root") | |
| assert tree.has_node("root") is True | |
| assert tree.has_node("nonexistent") is False | |
| async def test_find_node_by_status_message(self): | |
| """find_node_by_status_message finds the right node.""" | |
| tree = _make_tree("root") | |
| await tree.add_node("c1", _make_incoming(msg_id="c1"), "status_c1", "root") | |
| found = tree.find_node_by_status_message("status_c1") | |
| assert found is not None | |
| assert found.node_id == "c1" | |
| def test_find_node_by_status_message_not_found(self): | |
| """find_node_by_status_message returns None if not found.""" | |
| tree = _make_tree("root") | |
| assert tree.find_node_by_status_message("nonexistent") is None | |
| class TestMessageTreeSerialization: | |
| """Tests for tree serialization/deserialization.""" | |
| async def test_round_trip(self): | |
| """Tree should survive serialization round-trip.""" | |
| tree = _make_tree("root") | |
| await tree.add_node("c1", _make_incoming(msg_id="c1"), "s1", "root") | |
| await tree.add_node("c2", _make_incoming(msg_id="c2"), "s2", "root") | |
| await tree.update_state("root", MessageState.COMPLETED, session_id="sess1") | |
| data = tree.to_dict() | |
| restored = MessageTree.from_dict(data) | |
| assert restored.root_id == "root" | |
| assert len(restored.all_nodes()) == 3 | |
| root = restored.get_root() | |
| assert root.state == MessageState.COMPLETED | |
| assert root.session_id == "sess1" | |
| assert set(root.children_ids) == {"c1", "c2"} | |
| async def test_node_round_trip(self): | |
| """MessageNode should survive serialization round-trip.""" | |
| node = MessageNode( | |
| node_id="n1", | |
| incoming=_make_incoming(msg_id="n1"), | |
| status_message_id="s1", | |
| state=MessageState.COMPLETED, | |
| parent_id="root", | |
| session_id="sess_test", | |
| error_message="test error", | |
| ) | |
| data = node.to_dict() | |
| restored = MessageNode.from_dict(data) | |
| assert restored.node_id == "n1" | |
| assert restored.state == MessageState.COMPLETED | |
| assert restored.session_id == "sess_test" | |
| assert restored.error_message == "test error" | |
| assert restored.parent_id == "root" | |
| class TestTreeQueueManagerConcurrency: | |
| """Concurrency tests for TreeQueueManager.""" | |
| async def test_concurrent_create_trees(self): | |
| """Creating multiple trees concurrently should all succeed.""" | |
| mgr = TreeQueueManager() | |
| async def create(i: int): | |
| return await mgr.create_tree( | |
| node_id=f"root_{i}", | |
| incoming=_make_incoming(msg_id=f"root_{i}"), | |
| status_message_id=f"status_{i}", | |
| ) | |
| trees = await asyncio.gather(*[create(i) for i in range(10)]) | |
| assert len(trees) == 10 | |
| assert mgr.get_tree_count() == 10 | |
| async def test_add_to_tree_concurrent(self): | |
| """Adding replies to a tree concurrently should all succeed.""" | |
| mgr = TreeQueueManager() | |
| await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| async def add_reply(i: int): | |
| return await mgr.add_to_tree( | |
| parent_node_id="root", | |
| node_id=f"reply_{i}", | |
| incoming=_make_incoming(msg_id=f"reply_{i}"), | |
| status_message_id=f"s_reply_{i}", | |
| ) | |
| results = await asyncio.gather(*[add_reply(i) for i in range(5)]) | |
| assert len(results) == 5 | |
| tree = mgr.get_tree("root") | |
| assert tree is not None | |
| assert len(tree.all_nodes()) == 6 # root + 5 replies | |
| async def test_add_to_tree_invalid_parent(self): | |
| """Adding to a nonexistent parent should raise ValueError.""" | |
| mgr = TreeQueueManager() | |
| await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| with pytest.raises(ValueError, match="not found"): | |
| await mgr.add_to_tree( | |
| parent_node_id="nonexistent", | |
| node_id="reply", | |
| incoming=_make_incoming(), | |
| status_message_id="s1", | |
| ) | |
| async def test_enqueue_and_process(self): | |
| """Enqueue should process immediately if tree is free.""" | |
| mgr = TreeQueueManager() | |
| await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| processed = [] | |
| async def processor(node_id, node): | |
| processed.append(node_id) | |
| queued = await mgr.enqueue("root", processor) | |
| # Should process immediately (not queued) | |
| assert queued is False | |
| # Wait for the async task to complete | |
| await asyncio.sleep(0.1) | |
| assert "root" in processed | |
| async def test_enqueue_queues_when_busy(self): | |
| """Enqueue should queue when tree is already processing.""" | |
| mgr = TreeQueueManager() | |
| await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| _, _ = await mgr.add_to_tree("root", "c1", _make_incoming(msg_id="c1"), "s1") | |
| processing_started = asyncio.Event() | |
| release = asyncio.Event() | |
| async def slow_processor(node_id, node): | |
| processing_started.set() | |
| await release.wait() | |
| # Start processing root (will block) | |
| queued_first = await mgr.enqueue("root", slow_processor) | |
| assert queued_first is False | |
| await processing_started.wait() | |
| # Now tree is busy, second enqueue should be queued | |
| queued_second = await mgr.enqueue("c1", slow_processor) | |
| assert queued_second is True | |
| # Release the blocker so things clean up | |
| release.set() | |
| await asyncio.sleep(0.2) | |
| async def test_cancel_tree(self): | |
| """cancel_tree should cancel in-progress and queued nodes.""" | |
| mgr = TreeQueueManager() | |
| tree = await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| _, _ = await mgr.add_to_tree("root", "c1", _make_incoming(msg_id="c1"), "s1") | |
| _, _ = await mgr.add_to_tree("root", "c2", _make_incoming(msg_id="c2"), "s2") | |
| processing_started = asyncio.Event() | |
| async def slow_processor(node_id, node): | |
| processing_started.set() | |
| await asyncio.sleep(10) # Long running | |
| # Start processing root | |
| await mgr.enqueue("root", slow_processor) | |
| await processing_started.wait() | |
| # Queue additional nodes | |
| await mgr.enqueue("c1", slow_processor) | |
| await mgr.enqueue("c2", slow_processor) | |
| # Cancel the tree | |
| cancelled = await mgr.cancel_tree("root") | |
| assert len(cancelled) >= 1 # At least the current + queued | |
| # Tree should no longer be processing | |
| assert tree._is_processing is False | |
| async def test_cancel_nonexistent_tree(self): | |
| """cancel_tree for nonexistent tree returns empty list.""" | |
| mgr = TreeQueueManager() | |
| result = await mgr.cancel_tree("nonexistent") | |
| assert result == [] | |
| async def test_cancel_all(self): | |
| """cancel_all cancels all trees.""" | |
| mgr = TreeQueueManager() | |
| await mgr.create_tree("t1", _make_incoming(msg_id="t1"), "s1") | |
| await mgr.create_tree("t2", _make_incoming(msg_id="t2"), "s2") | |
| # Mark nodes as PENDING (they already are by default) | |
| cancelled = await mgr.cancel_all() | |
| # Nodes were PENDING but not in queue, so cleanup_stale logic applies | |
| # At minimum, it should not raise | |
| assert isinstance(cancelled, list) | |
| async def test_cleanup_stale_nodes(self): | |
| """cleanup_stale_nodes marks PENDING/IN_PROGRESS nodes as ERROR.""" | |
| mgr = TreeQueueManager() | |
| tree = await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| _, _ = await mgr.add_to_tree("root", "c1", _make_incoming(msg_id="c1"), "s1") | |
| # Root is PENDING, c1 is PENDING | |
| count = mgr.cleanup_stale_nodes() | |
| assert count == 2 | |
| root = tree.get_node("root") | |
| assert root is not None | |
| assert root.state == MessageState.ERROR | |
| assert root.error_message is not None | |
| assert "restart" in root.error_message | |
| async def test_mark_node_error_with_propagation(self): | |
| """mark_node_error should propagate to pending children.""" | |
| mgr = TreeQueueManager() | |
| tree = await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| _, _ = await mgr.add_to_tree("root", "c1", _make_incoming(msg_id="c1"), "s1") | |
| _, _ = await mgr.add_to_tree("c1", "c2", _make_incoming(msg_id="c2"), "s2") | |
| affected = await mgr.mark_node_error("root", "something failed") | |
| # root + c1 + c2 should all be marked | |
| assert len(affected) >= 1 | |
| root = tree.get_node("root") | |
| assert root is not None | |
| assert root.state == MessageState.ERROR | |
| async def test_mark_node_error_nonexistent(self): | |
| """mark_node_error for nonexistent node returns empty.""" | |
| mgr = TreeQueueManager() | |
| result = await mgr.mark_node_error("nonexistent", "err") | |
| assert result == [] | |
| async def test_get_tree_for_node(self): | |
| """get_tree_for_node returns the right tree.""" | |
| mgr = TreeQueueManager() | |
| await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| _, _ = await mgr.add_to_tree("root", "c1", _make_incoming(msg_id="c1"), "s1") | |
| tree = mgr.get_tree_for_node("c1") | |
| assert tree is not None | |
| assert tree.root_id == "root" | |
| def test_get_tree_for_node_nonexistent(self): | |
| """get_tree_for_node returns None for unknown nodes.""" | |
| mgr = TreeQueueManager() | |
| assert mgr.get_tree_for_node("nonexistent") is None | |
| async def test_enqueue_no_tree(self): | |
| """Enqueue for a node not in any tree returns False.""" | |
| mgr = TreeQueueManager() | |
| async def dummy(nid, node): | |
| pass | |
| result = await mgr.enqueue("nonexistent", dummy) | |
| assert result is False | |
| async def test_serialization_round_trip(self): | |
| """TreeQueueManager should survive serialization round-trip.""" | |
| mgr = TreeQueueManager() | |
| await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| _, _ = await mgr.add_to_tree("root", "c1", _make_incoming(msg_id="c1"), "s1") | |
| data = mgr.to_dict() | |
| restored = TreeQueueManager.from_dict(data) | |
| assert restored.get_tree_count() == 1 | |
| assert restored.get_node("c1") is not None | |
| async def test_rapid_messages_all_queued(self): | |
| """Rapid sequential enqueues should all be queued without loss.""" | |
| mgr = TreeQueueManager() | |
| tree = await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| # Add 10 child nodes | |
| for i in range(10): | |
| await mgr.add_to_tree( | |
| "root", f"c{i}", _make_incoming(msg_id=f"c{i}"), f"s{i}" | |
| ) | |
| blocker = asyncio.Event() | |
| async def blocking_processor(nid, node): | |
| await blocker.wait() | |
| # Start processing root (blocks) | |
| await mgr.enqueue("root", blocking_processor) | |
| await asyncio.sleep(0.05) # Let task start | |
| # Rapidly enqueue all children | |
| results = [] | |
| for i in range(10): | |
| r = await mgr.enqueue(f"c{i}", blocking_processor) | |
| results.append(r) | |
| # All should be queued (True) | |
| assert all(r is True for r in results) | |
| assert tree.get_queue_size() == 10 | |
| # Cleanup | |
| blocker.set() | |
| await asyncio.sleep(0.1) | |
| async def test_concurrent_trees_independent(self): | |
| """Processing in one tree shouldn't affect another.""" | |
| mgr = TreeQueueManager() | |
| await mgr.create_tree("t1", _make_incoming(msg_id="t1"), "s1") | |
| await mgr.create_tree("t2", _make_incoming(msg_id="t2"), "s2") | |
| processed = [] | |
| async def processor(nid, node): | |
| processed.append(nid) | |
| # Process both trees | |
| await mgr.enqueue("t1", processor) | |
| await mgr.enqueue("t2", processor) | |
| await asyncio.sleep(0.2) | |
| assert "t1" in processed | |
| assert "t2" in processed | |
| async def test_callbacks_invoked(self): | |
| """Queue update and node started callbacks should fire.""" | |
| queue_updates = [] | |
| node_starts = [] | |
| async def on_queue_update(tree): | |
| queue_updates.append(tree.root_id) | |
| async def on_node_started(tree, node_id): | |
| node_starts.append(node_id) | |
| mgr = TreeQueueManager( | |
| queue_update_callback=on_queue_update, | |
| node_started_callback=on_node_started, | |
| ) | |
| await mgr.create_tree("root", _make_incoming(msg_id="root"), "s_root") | |
| _, _ = await mgr.add_to_tree("root", "c1", _make_incoming(msg_id="c1"), "s1") | |
| blocker = asyncio.Event() | |
| async def slow_proc(nid, node): | |
| if nid == "root": | |
| blocker.set() | |
| await asyncio.sleep(0.1) | |
| # Process root then c1 should be dequeued | |
| await mgr.enqueue("root", slow_proc) | |
| await blocker.wait() | |
| await mgr.enqueue("c1", slow_proc) | |
| await asyncio.sleep(0.5) | |
| # c1 was dequeued from queue, so callbacks should have fired | |
| assert len(queue_updates) >= 1 or len(node_starts) >= 1 | |