refactor: simplify /stop dispatch, inline commands, trim verbose docstrings
This commit is contained in:
Re-bin
@@ -1,4 +1,4 @@
|
||||
"""Tests for the command system and task cancellation."""
|
||||
"""Tests for /stop task cancellation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -7,117 +7,42 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.commands import (
|
||||
COMMANDS,
|
||||
get_help_text,
|
||||
is_immediate_command,
|
||||
parse_command,
|
||||
)
|
||||
|
||||
def _make_loop():
|
||||
"""Create a minimal AgentLoop with mocked dependencies."""
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.queue import MessageBus
|
||||
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
workspace = MagicMock()
|
||||
workspace.__truediv__ = MagicMock(return_value=MagicMock())
|
||||
|
||||
with patch("nanobot.agent.loop.ContextBuilder"), \
|
||||
patch("nanobot.agent.loop.SessionManager"), \
|
||||
patch("nanobot.agent.loop.SubagentManager") as MockSubMgr:
|
||||
MockSubMgr.return_value.cancel_by_session = AsyncMock(return_value=0)
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
|
||||
return loop, bus
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# commands.py unit tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestParseCommand:
|
||||
def test_slash_command(self):
|
||||
assert parse_command("/stop") == "/stop"
|
||||
|
||||
def test_slash_command_with_args(self):
|
||||
assert parse_command("/new some args") == "/new"
|
||||
|
||||
def test_not_a_command(self):
|
||||
assert parse_command("hello world") is None
|
||||
|
||||
def test_empty_string(self):
|
||||
assert parse_command("") is None
|
||||
|
||||
def test_leading_whitespace(self):
|
||||
assert parse_command(" /help") == "/help"
|
||||
|
||||
def test_uppercase_normalized(self):
|
||||
assert parse_command("/STOP") == "/stop"
|
||||
|
||||
|
||||
class TestIsImmediateCommand:
|
||||
def test_stop_is_immediate(self):
|
||||
assert is_immediate_command("/stop") is True
|
||||
|
||||
def test_new_is_not_immediate(self):
|
||||
assert is_immediate_command("/new") is False
|
||||
|
||||
def test_help_is_not_immediate(self):
|
||||
assert is_immediate_command("/help") is False
|
||||
|
||||
def test_unknown_command(self):
|
||||
assert is_immediate_command("/unknown") is False
|
||||
|
||||
|
||||
class TestGetHelpText:
|
||||
def test_contains_all_commands(self):
|
||||
text = get_help_text()
|
||||
for cmd in COMMANDS:
|
||||
assert cmd in text
|
||||
|
||||
def test_contains_descriptions(self):
|
||||
text = get_help_text()
|
||||
for defn in COMMANDS.values():
|
||||
assert defn.description in text
|
||||
|
||||
def test_starts_with_header(self):
|
||||
assert get_help_text().startswith("🐈")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Task cancellation integration tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestTaskCancellation:
|
||||
"""Tests for /stop cancelling an active task in AgentLoop."""
|
||||
|
||||
def _make_loop(self):
|
||||
"""Create a minimal AgentLoop with mocked dependencies."""
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.queue import MessageBus
|
||||
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
workspace = MagicMock()
|
||||
workspace.__truediv__ = MagicMock(return_value=MagicMock())
|
||||
|
||||
with patch("nanobot.agent.loop.ContextBuilder"), \
|
||||
patch("nanobot.agent.loop.SessionManager"), \
|
||||
patch("nanobot.agent.loop.SubagentManager") as MockSubMgr:
|
||||
MockSubMgr.return_value.cancel_by_session = AsyncMock(return_value=0)
|
||||
loop = AgentLoop(
|
||||
bus=bus,
|
||||
provider=provider,
|
||||
workspace=workspace,
|
||||
)
|
||||
return loop, bus
|
||||
|
||||
class TestHandleStop:
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_no_active_task(self):
|
||||
"""'/stop' when nothing is running returns 'No active task'."""
|
||||
from nanobot.bus.events import InboundMessage
|
||||
|
||||
loop, bus = self._make_loop()
|
||||
msg = InboundMessage(
|
||||
channel="test", sender_id="u1", chat_id="c1", content="/stop"
|
||||
)
|
||||
await loop._handle_immediate_command("/stop", msg)
|
||||
loop, bus = _make_loop()
|
||||
msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="/stop")
|
||||
await loop._handle_stop(msg)
|
||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
||||
assert "No active task" in out.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_cancels_active_task(self):
|
||||
"""'/stop' cancels a running task."""
|
||||
from nanobot.bus.events import InboundMessage
|
||||
|
||||
loop, bus = self._make_loop()
|
||||
session_key = "test:c1"
|
||||
|
||||
loop, bus = _make_loop()
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def slow_task():
|
||||
@@ -128,74 +53,61 @@ class TestTaskCancellation:
|
||||
raise
|
||||
|
||||
task = asyncio.create_task(slow_task())
|
||||
await asyncio.sleep(0) # Let task enter its await
|
||||
loop._active_tasks[session_key] = task
|
||||
await asyncio.sleep(0)
|
||||
loop._active_tasks["test:c1"] = [task]
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="test", sender_id="u1", chat_id="c1", content="/stop"
|
||||
)
|
||||
await loop._handle_immediate_command("/stop", msg)
|
||||
msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="/stop")
|
||||
await loop._handle_stop(msg)
|
||||
|
||||
assert cancelled.is_set()
|
||||
assert task.cancelled()
|
||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
||||
assert "stopped" in out.content.lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_registers_and_clears_task(self):
|
||||
"""_dispatch registers the task in _active_tasks and clears it after."""
|
||||
async def test_stop_cancels_multiple_tasks(self):
|
||||
from nanobot.bus.events import InboundMessage
|
||||
|
||||
loop, bus = _make_loop()
|
||||
events = [asyncio.Event(), asyncio.Event()]
|
||||
|
||||
async def slow(idx):
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
except asyncio.CancelledError:
|
||||
events[idx].set()
|
||||
raise
|
||||
|
||||
tasks = [asyncio.create_task(slow(i)) for i in range(2)]
|
||||
await asyncio.sleep(0)
|
||||
loop._active_tasks["test:c1"] = tasks
|
||||
|
||||
msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="/stop")
|
||||
await loop._handle_stop(msg)
|
||||
|
||||
assert all(e.is_set() for e in events)
|
||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
||||
assert "2 task" in out.content
|
||||
|
||||
|
||||
class TestDispatch:
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_processes_and_publishes(self):
|
||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||
|
||||
loop, bus = self._make_loop()
|
||||
msg = InboundMessage(
|
||||
channel="test", sender_id="u1", chat_id="c1", content="hello"
|
||||
)
|
||||
|
||||
# Mock _process_message to return a simple response
|
||||
loop, bus = _make_loop()
|
||||
msg = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="hello")
|
||||
loop._process_message = AsyncMock(
|
||||
return_value=OutboundMessage(channel="test", chat_id="c1", content="hi")
|
||||
)
|
||||
|
||||
task = asyncio.create_task(loop._dispatch(msg))
|
||||
await task
|
||||
|
||||
# Task should be cleaned up
|
||||
assert msg.session_key not in loop._active_tasks
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatch_handles_cancelled_error(self):
|
||||
"""_dispatch catches CancelledError gracefully."""
|
||||
from nanobot.bus.events import InboundMessage
|
||||
|
||||
loop, bus = self._make_loop()
|
||||
msg = InboundMessage(
|
||||
channel="test", sender_id="u1", chat_id="c1", content="hello"
|
||||
)
|
||||
|
||||
async def mock_process(m, **kwargs):
|
||||
await asyncio.sleep(60)
|
||||
|
||||
loop._process_message = mock_process
|
||||
|
||||
task = asyncio.create_task(loop._dispatch(msg))
|
||||
await asyncio.sleep(0.05) # Let task start
|
||||
|
||||
assert msg.session_key in loop._active_tasks
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
# Task should be cleaned up even after cancel
|
||||
assert msg.session_key not in loop._active_tasks
|
||||
await loop._dispatch(msg)
|
||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
||||
assert out.content == "hi"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_processing_lock_serializes(self):
|
||||
"""Only one message processes at a time due to _processing_lock."""
|
||||
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||
|
||||
loop, bus = self._make_loop()
|
||||
loop, bus = _make_loop()
|
||||
order = []
|
||||
|
||||
async def mock_process(m, **kwargs):
|
||||
@@ -205,27 +117,18 @@ class TestTaskCancellation:
|
||||
return OutboundMessage(channel="test", chat_id="c1", content=m.content)
|
||||
|
||||
loop._process_message = mock_process
|
||||
|
||||
msg1 = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="a")
|
||||
msg2 = InboundMessage(channel="test", sender_id="u1", chat_id="c1", content="b")
|
||||
|
||||
t1 = asyncio.create_task(loop._dispatch(msg1))
|
||||
t2 = asyncio.create_task(loop._dispatch(msg2))
|
||||
await asyncio.gather(t1, t2)
|
||||
|
||||
# Should be serialized: start-a, end-a, start-b, end-b
|
||||
assert order == ["start-a", "end-a", "start-b", "end-b"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSubagentCancellation:
|
||||
"""Tests for /stop cancelling subagents spawned under a session."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_by_session(self):
|
||||
"""cancel_by_session cancels all tasks for that session."""
|
||||
from nanobot.agent.subagent import SubagentManager
|
||||
from nanobot.bus.queue import MessageBus
|
||||
|
||||
@@ -236,28 +139,24 @@ class TestSubagentCancellation:
|
||||
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def slow_subagent():
|
||||
async def slow():
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
|
||||
task = asyncio.create_task(slow_subagent())
|
||||
task = asyncio.create_task(slow())
|
||||
await asyncio.sleep(0)
|
||||
tid = "sub-1"
|
||||
session_key = "test:c1"
|
||||
mgr._running_tasks[tid] = task
|
||||
mgr._session_tasks[session_key] = {tid}
|
||||
mgr._running_tasks["sub-1"] = task
|
||||
mgr._session_tasks["test:c1"] = {"sub-1"}
|
||||
|
||||
count = await mgr.cancel_by_session(session_key)
|
||||
count = await mgr.cancel_by_session("test:c1")
|
||||
assert count == 1
|
||||
assert cancelled.is_set()
|
||||
assert task.cancelled()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_by_session_no_tasks(self):
|
||||
"""cancel_by_session returns 0 when no subagents for session."""
|
||||
from nanobot.agent.subagent import SubagentManager
|
||||
from nanobot.bus.queue import MessageBus
|
||||
|
||||
@@ -265,54 +164,4 @@ class TestSubagentCancellation:
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
mgr = SubagentManager(provider=provider, workspace=MagicMock(), bus=bus)
|
||||
|
||||
count = await mgr.cancel_by_session("nonexistent:session")
|
||||
assert count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_cancels_subagents_via_loop(self):
|
||||
"""/stop on AgentLoop also cancels subagents for that session."""
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
workspace = MagicMock()
|
||||
workspace.__truediv__ = MagicMock(return_value=MagicMock())
|
||||
|
||||
with patch("nanobot.agent.loop.ContextBuilder"), \
|
||||
patch("nanobot.agent.loop.SessionManager"), \
|
||||
patch("nanobot.agent.loop.SubagentManager"):
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
|
||||
|
||||
# Replace subagents with a real SubagentManager
|
||||
from nanobot.agent.subagent import SubagentManager
|
||||
loop.subagents = SubagentManager(
|
||||
provider=provider, workspace=MagicMock(), bus=bus
|
||||
)
|
||||
|
||||
cancelled = asyncio.Event()
|
||||
session_key = "test:c1"
|
||||
|
||||
async def slow_sub():
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
|
||||
task = asyncio.create_task(slow_sub())
|
||||
await asyncio.sleep(0)
|
||||
loop.subagents._running_tasks["sub-1"] = task
|
||||
loop.subagents._session_tasks[session_key] = {"sub-1"}
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="test", sender_id="u1", chat_id="c1", content="/stop"
|
||||
)
|
||||
await loop._handle_immediate_command("/stop", msg)
|
||||
|
||||
assert cancelled.is_set()
|
||||
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
||||
assert "stopped" in out.content.lower() or "background" in out.content.lower()
|
||||
assert await mgr.cancel_by_session("nonexistent") == 0
|
||||
|
||||
Reference in New Issue
Block a user