""" Tests for runtime/handler_dispatcher.py - HandlerDispatcher implementation. """ from __future__ import annotations import asyncio from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest from astrbot_sdk.api.event import AstrMessageEvent from astrbot_sdk.api.message import Comp, MessageChain from astrbot_sdk.context import CancelToken, Context from astrbot_sdk.events import MessageEvent, PlainTextResult from astrbot_sdk.protocol.descriptors import ( CommandTrigger, HandlerDescriptor, ) from astrbot_sdk.protocol.messages import InvokeMessage from astrbot_sdk.runtime.handler_dispatcher import HandlerDispatcher from astrbot_sdk.runtime.loader import LoadedHandler class MockPeer: """Mock peer for testing.""" def __init__(self): self.sent_messages: list[dict[str, Any]] = [] self.platform = self # platform.send 通过 self.send 调用 # CapabilityProxy 需要的属性 self.remote_capability_map: dict[str, Any] = {} async def send(self, session_id: str, text: str) -> None: """模拟 platform.send 方法""" self.sent_messages.append({"session_id": session_id, "text": text}) async def invoke( self, name: str, payload: dict[str, Any], stream: bool = False ) -> dict[str, Any]: """模拟 peer.invoke 方法,用于 CapabilityProxy""" if name == "platform.send": await self.send( payload.get("session", payload.get("session_id", "")), payload.get("text", ""), ) return {} if name == "platform.send_chain": self.sent_messages.append( { "session_id": payload.get("session", ""), "chain": payload.get("chain", []), } ) return {} return {} def create_mock_handler( handler_id: str = "test.handler", command: str = "hello", ) -> LoadedHandler: """Create a mock loaded handler.""" descriptor = HandlerDescriptor( id=handler_id, trigger=CommandTrigger(command=command), ) async def handler_func(event: MessageEvent, ctx: Context): await event.reply("Hello!") return None handler_func.__func__ = handler_func # Simulate bound method return LoadedHandler( descriptor=descriptor, callable=handler_func, owner=MagicMock(), legacy_context=None, ) def create_invoke_message( message_id: str = "msg_001", handler_id: str = "test.handler", event_data: dict[str, Any] | None = None, args: dict[str, Any] | None = None, ) -> InvokeMessage: """Create a mock invoke message.""" input_data = {"handler_id": handler_id, "event": event_data or {}} if args: input_data["args"] = args return InvokeMessage( id=message_id, capability="handler.invoke", input=input_data, ) def create_message_event() -> MessageEvent: """Create a mock message event.""" return MessageEvent( session_id="session-1", user_id="user-1", platform="test", text="hello world", ) class TestHandlerDispatcherInit: """Tests for HandlerDispatcher initialization.""" def test_init(self): """HandlerDispatcher should initialize with handlers.""" peer = MockPeer() handler = create_mock_handler() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) assert dispatcher._plugin_id == "test_plugin" assert dispatcher._peer is peer assert "test.handler" in dispatcher._handlers assert dispatcher._active == {} def test_handlers_indexed_by_id(self): """HandlerDispatcher should index handlers by id.""" peer = MockPeer() handlers = [ create_mock_handler("handler.one", "cmd1"), create_mock_handler("handler.two", "cmd2"), ] dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=handlers, ) assert "handler.one" in dispatcher._handlers assert "handler.two" in dispatcher._handlers class TestHandlerDispatcherInvoke: """Tests for HandlerDispatcher.invoke method.""" @pytest.mark.asyncio async def test_invoke_calls_handler(self): """invoke should call the registered handler.""" peer = MockPeer() sent_messages = [] # 记录 platform.send 的调用 original_send = peer.send async def track_send(session_id: str, text: str) -> None: sent_messages.append({"session_id": session_id, "text": text}) await original_send(session_id, text) peer.send = track_send handler_called = [] async def handler_func(e: MessageEvent, ctx: Context): handler_called.append(e) await e.reply("response") descriptor = HandlerDescriptor( id="test.handler", trigger=CommandTrigger(command="hello"), ) handler = LoadedHandler( descriptor=descriptor, callable=handler_func, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) event = create_message_event() # MessageEvent 使用 to_payload() 而不是 model_dump() message = InvokeMessage( id="msg_001", capability="handler.invoke", input={"handler_id": "test.handler", "event": event.to_payload()}, ) cancel_token = CancelToken() result = await dispatcher.invoke(message, cancel_token) assert result == {} assert len(handler_called) == 1 # 验证 reply 通过 platform.send 发送 assert any("response" in m.get("text", "") for m in sent_messages) @pytest.mark.asyncio async def test_invoke_missing_handler_raises(self): """invoke should raise LookupError for missing handler.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) message = InvokeMessage( id="msg_001", capability="handler.invoke", input={"handler_id": "nonexistent.handler", "event": {}}, ) cancel_token = CancelToken() with pytest.raises(LookupError, match="handler not found"): await dispatcher.invoke(message, cancel_token) @pytest.mark.asyncio async def test_invoke_with_legacy_args(self): """invoke should pass legacy args to handler.""" peer = MockPeer() received_args = [] async def handler_func(event: MessageEvent, ctx: Context, name: str): received_args.append(name) return None descriptor = HandlerDescriptor( id="test.handler", trigger=CommandTrigger(command="hello"), ) handler = LoadedHandler( descriptor=descriptor, callable=handler_func, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) message = InvokeMessage( id="msg_001", capability="handler.invoke", input={ "handler_id": "test.handler", "event": { "type": "message", "session_id": "s1", "user_id": "u1", "platform": "test", }, "args": {"name": "test_name"}, }, ) cancel_token = CancelToken() await dispatcher.invoke(message, cancel_token) assert "test_name" in received_args @pytest.mark.asyncio async def test_invoke_tracks_active_task(self): """invoke should track active task.""" peer = MockPeer() async def slow_handler(event: MessageEvent, ctx: Context): await asyncio.sleep(0.1) return None descriptor = HandlerDescriptor( id="slow.handler", trigger=CommandTrigger(command="slow"), ) handler = LoadedHandler( descriptor=descriptor, callable=slow_handler, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) message = InvokeMessage( id="msg_001", capability="handler.invoke", input={ "handler_id": "slow.handler", "event": { "type": "message", "session_id": "s1", "user_id": "u1", "platform": "test", }, }, ) cancel_token = CancelToken() # Start invoke in background task = asyncio.create_task(dispatcher.invoke(message, cancel_token)) # Give it time to start await asyncio.sleep(0) # Should have active task during execution # Note: might be empty if task completes quickly await task # After completion, should be cleared assert "msg_001" not in dispatcher._active @pytest.mark.asyncio async def test_invoke_wraps_legacy_astr_message_event(self): """Annotated AstrMessageEvent handlers should receive the compat wrapper.""" peer = MockPeer() received_types = [] replies = [] async def handler_func(event: AstrMessageEvent): received_types.append(type(event)) yield event.plain_result("legacy reply") async def track_send(session_id: str, text: str) -> None: replies.append({"session_id": session_id, "text": text}) peer.send = track_send descriptor = HandlerDescriptor( id="legacy.handler", trigger=CommandTrigger(command="hello"), ) handler = LoadedHandler( descriptor=descriptor, callable=handler_func, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) message = InvokeMessage( id="msg_legacy", capability="handler.invoke", input={ "handler_id": "legacy.handler", "event": { "text": "hello", "session_id": "session-legacy", "user_id": "user-1", "platform": "test", }, }, ) await dispatcher.invoke(message, CancelToken()) assert received_types == [AstrMessageEvent] assert replies == [{"session_id": "session-legacy", "text": "legacy reply"}] @pytest.mark.asyncio async def test_invoke_reply_falls_back_to_peer_send_for_sync_mock(self): """invoke should fall back to peer.send when peer.invoke is a sync mock.""" peer = MagicMock() peer.remote_capability_map = {} sent_messages = [] async def track_send(session_id: str, text: str) -> None: sent_messages.append({"session_id": session_id, "text": text}) peer.send = track_send async def handler_func(event: MessageEvent, ctx: Context): await event.reply("fallback") descriptor = HandlerDescriptor( id="test.handler", trigger=CommandTrigger(command="hello"), ) handler = LoadedHandler( descriptor=descriptor, callable=handler_func, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) message = InvokeMessage( id="msg_sync_mock", capability="handler.invoke", input={ "handler_id": "test.handler", "event": { "text": "hello", "session_id": "session-sync", "user_id": "user-1", "platform": "test", }, }, ) await dispatcher.invoke(message, CancelToken()) assert sent_messages == [{"session_id": "session-sync", "text": "fallback"}] class TestHandlerDispatcherCancel: """Tests for HandlerDispatcher.cancel method.""" @pytest.mark.asyncio async def test_cancel_stops_active_task(self): """cancel should stop the active task.""" peer = MockPeer() cancelled = [] async def slow_handler(event: MessageEvent, ctx: Context): try: await asyncio.sleep(10) # Long sleep except asyncio.CancelledError: cancelled.append(True) raise descriptor = HandlerDescriptor( id="slow.handler", trigger=CommandTrigger(command="slow"), ) handler = LoadedHandler( descriptor=descriptor, callable=slow_handler, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) message = InvokeMessage( id="msg_001", capability="handler.invoke", input={ "handler_id": "slow.handler", "event": { "type": "message", "session_id": "s1", "user_id": "u1", "platform": "test", }, }, ) cancel_token = CancelToken() # Start invoke task = asyncio.create_task(dispatcher.invoke(message, cancel_token)) # Wait for task to be active await asyncio.sleep(0.05) # Cancel await dispatcher.cancel("msg_001") # Task should be cancelled await asyncio.sleep(0.05) assert cancelled or task.cancelled() or task.done() @pytest.mark.asyncio async def test_cancel_unknown_request_does_nothing(self): """cancel should do nothing for unknown request.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) # Should not raise await dispatcher.cancel("unknown_request") class TestHandlerDispatcherBuildArgs: """Tests for HandlerDispatcher._build_args method.""" def test_build_args_event_parameter(self): """_build_args should inject event parameter.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) def handler(event: MessageEvent, ctx: Context): pass event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) args = dispatcher._build_args(handler, event, ctx) assert args[0] is event def test_build_args_ctx_parameter(self): """_build_args should inject ctx/context parameter.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) def handler(event: MessageEvent, ctx: Context): pass event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) args = dispatcher._build_args(handler, event, ctx) assert args[1] is ctx def test_build_args_legacy_args(self): """_build_args should inject legacy args.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) def handler(event: MessageEvent, ctx: Context, custom_arg: str): pass event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) args = dispatcher._build_args(handler, event, ctx, {"custom_arg": "value"}) assert args[2] == "value" def test_build_args_skip_keyword_only(self): """_build_args should skip keyword-only parameters.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) def handler(event: MessageEvent, *, optional: str = "default"): pass event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) args = dispatcher._build_args(handler, event, ctx) # Should only have event assert len(args) == 1 class TestHandlerDispatcherConsumeResult: """Tests for HandlerDispatcher._consume_legacy_result method.""" @pytest.mark.asyncio async def test_consume_plain_text_result(self): """_consume_legacy_result should handle PlainTextResult.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) replies = [] event = create_message_event() # reply_handler 必须是异步的 async def async_reply(text: str) -> None: replies.append(text) event._reply_handler = async_reply result = PlainTextResult(text="plain text") await dispatcher._consume_legacy_result(result, event) assert "plain text" in replies @pytest.mark.asyncio async def test_consume_string(self): """_consume_legacy_result should handle string.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) replies = [] event = create_message_event() # reply_handler 必须是异步的 async def async_reply(text: str) -> None: replies.append(text) event._reply_handler = async_reply await dispatcher._consume_legacy_result("string reply", event) assert "string reply" in replies @pytest.mark.asyncio async def test_consume_dict_with_text(self): """_consume_legacy_result should handle dict with text.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) replies = [] event = create_message_event() # reply_handler 必须是异步的 async def async_reply(text: str) -> None: replies.append(text) event._reply_handler = async_reply await dispatcher._consume_legacy_result({"text": "dict reply"}, event) assert "dict reply" in replies @pytest.mark.asyncio async def test_consume_other_type_ignored(self): """_consume_legacy_result should ignore other types.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) event = create_message_event() event._reply_handler = MagicMock() # Should not raise await dispatcher._consume_legacy_result(123, event) await dispatcher._consume_legacy_result(None, event) @pytest.mark.asyncio async def test_consume_message_chain_uses_platform_send_chain(self): """_consume_legacy_result should preserve rich chains when ctx is available.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) chain = MessageChain( [ Comp.Plain(text="hello"), Comp.Image(file="https://example.com/image.png"), ] ) await dispatcher._consume_legacy_result(chain, event, ctx) assert peer.sent_messages == [ { "session_id": "session-1", "chain": [ {"type": "Plain", "text": "hello"}, {"type": "Image", "file": "https://example.com/image.png"}, ], } ] class TestHandlerDispatcherHandleError: """Tests for HandlerDispatcher._handle_error method.""" @pytest.mark.asyncio async def test_handle_error_with_on_error_method(self): """_handle_error should call owner.on_error if available.""" peer = MockPeer() errors_handled = [] class OwnerWithOnError: async def on_error(self, exc: Exception, event: MessageEvent, ctx: Context): errors_handled.append(exc) owner = OwnerWithOnError() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) exc = ValueError("test error") await dispatcher._handle_error(owner, exc, event, ctx) assert exc in errors_handled @pytest.mark.asyncio async def test_handle_error_without_on_error_method(self): """_handle_error should use Star.on_error if owner has no on_error.""" peer = MockPeer() dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[], ) event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) exc = ValueError("test error") # owner 是 MagicMock,on_error 方法返回 MagicMock 而不是协程 # 但 _handle_error 会 await owner.on_error(...) # 所以我们需要让 owner.on_error 返回一个协程 owner = MagicMock() owner.on_error = AsyncMock() # Should not raise await dispatcher._handle_error(owner, exc, event, ctx) class TestHandlerDispatcherRunHandler: """Tests for HandlerDispatcher._run_handler method.""" @pytest.mark.asyncio async def test_run_handler_sync_function(self): """_run_handler should handle sync function.""" peer = MockPeer() called = [] def sync_handler(event: MessageEvent, ctx: Context): called.append(True) descriptor = HandlerDescriptor( id="sync.handler", trigger=CommandTrigger(command="sync"), ) handler = LoadedHandler( descriptor=descriptor, callable=sync_handler, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) await dispatcher._run_handler(handler, event, ctx) assert called @pytest.mark.asyncio async def test_run_handler_async_function(self): """_run_handler should handle async function.""" peer = MockPeer() called = [] async def async_handler(event: MessageEvent, ctx: Context): called.append(True) descriptor = HandlerDescriptor( id="async.handler", trigger=CommandTrigger(command="async"), ) handler = LoadedHandler( descriptor=descriptor, callable=async_handler, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) await dispatcher._run_handler(handler, event, ctx) assert called @pytest.mark.asyncio async def test_run_handler_async_generator(self): """_run_handler should handle async generator.""" peer = MockPeer() replies = [] async def gen_handler(event: MessageEvent, ctx: Context): yield "first" yield "second" descriptor = HandlerDescriptor( id="gen.handler", trigger=CommandTrigger(command="gen"), ) handler = LoadedHandler( descriptor=descriptor, callable=gen_handler, owner=MagicMock(), legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) event = create_message_event() # reply_handler 必须是异步的 async def async_reply(text: str) -> None: replies.append(text) event._reply_handler = async_reply ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) await dispatcher._run_handler(handler, event, ctx) assert "first" in replies assert "second" in replies @pytest.mark.asyncio async def test_run_handler_with_exception(self): """_run_handler should handle exceptions.""" peer = MockPeer() async def failing_handler(event: MessageEvent, ctx: Context): raise ValueError("handler error") descriptor = HandlerDescriptor( id="failing.handler", trigger=CommandTrigger(command="fail"), ) # owner.on_error 需要是异步的 owner = MagicMock() owner.on_error = AsyncMock() handler = LoadedHandler( descriptor=descriptor, callable=failing_handler, owner=owner, legacy_context=None, ) dispatcher = HandlerDispatcher( plugin_id="test_plugin", peer=peer, handlers=[handler], ) event = create_message_event() ctx = Context(peer=peer, plugin_id="test", cancel_token=CancelToken()) with pytest.raises(ValueError, match="handler error"): await dispatcher._run_handler(handler, event, ctx)