mirror of
https://github.com/AstrBotDevs/AstrBot.git
synced 2026-09-24 16:39:52 +08:00
1235 lines
43 KiB
Python
1235 lines
43 KiB
Python
"""
|
|
Tests for runtime/bootstrap.py - Bootstrap and runtime classes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
import tempfile
|
|
from io import StringIO
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
from astrbot_sdk._legacy_api import LegacyContext
|
|
from astrbot_sdk._legacy_runtime import LegacyRuntimeAdapter
|
|
from astrbot_sdk.context import CancelToken
|
|
from astrbot_sdk.errors import AstrBotError
|
|
from astrbot_sdk.protocol.descriptors import (
|
|
CapabilityDescriptor,
|
|
CommandTrigger,
|
|
HandlerDescriptor,
|
|
)
|
|
from astrbot_sdk.protocol.messages import (
|
|
EventMessage,
|
|
InitializeMessage,
|
|
InitializeOutput,
|
|
InvokeMessage,
|
|
PeerInfo,
|
|
)
|
|
from astrbot_sdk.runtime.bootstrap import (
|
|
PluginWorkerRuntime,
|
|
SupervisorRuntime,
|
|
WorkerSession,
|
|
_install_signal_handlers,
|
|
_prepare_stdio_transport,
|
|
_wait_for_shutdown,
|
|
)
|
|
from astrbot_sdk.runtime.capability_router import CapabilityRouter
|
|
from astrbot_sdk.runtime.loader import LoadedHandler, PluginSpec
|
|
from astrbot_sdk.runtime.peer import Peer
|
|
|
|
from tests_v4.helpers import FakeEnvManager, MemoryTransport, make_transport_pair
|
|
|
|
|
|
async def start_test_core_peer(transport: MemoryTransport) -> Peer:
|
|
"""Provide an initialize responder so transport-pair startup tests do not deadlock."""
|
|
core = Peer(
|
|
transport=transport,
|
|
peer_info=PeerInfo(name="test-core", role="core", version="v4"),
|
|
)
|
|
core.set_initialize_handler(
|
|
lambda _message: asyncio.sleep(
|
|
0,
|
|
result=InitializeOutput(
|
|
peer=PeerInfo(name="test-core", role="core", version="v4"),
|
|
capabilities=[],
|
|
metadata={},
|
|
),
|
|
)
|
|
)
|
|
await core.start()
|
|
return core
|
|
|
|
|
|
def write_bootstrap_plugin(
|
|
plugins_dir: Path,
|
|
name: str,
|
|
*,
|
|
python_version: str | None = None,
|
|
) -> Path:
|
|
plugin_dir = plugins_dir / name
|
|
plugin_dir.mkdir(parents=True, exist_ok=True)
|
|
(plugin_dir / "plugin.yaml").write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": name,
|
|
"runtime": {
|
|
"python": python_version
|
|
or f"{sys.version_info.major}.{sys.version_info.minor}"
|
|
},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
(plugin_dir / "requirements.txt").write_text("", encoding="utf-8")
|
|
return plugin_dir
|
|
|
|
|
|
class PlanningFakeEnvManager(FakeEnvManager):
|
|
def __init__(self, *, skipped_plugins: dict[str, str] | None = None) -> None:
|
|
self.skipped_plugins = dict(skipped_plugins or {})
|
|
self.events: list[tuple[str, object]] = []
|
|
self.prepared_paths: dict[str, Path] = {}
|
|
|
|
def plan(self, plugins):
|
|
plugin_names = [plugin.name for plugin in plugins]
|
|
self.events.append(("plan", plugin_names))
|
|
planned_plugins = [
|
|
plugin for plugin in plugins if plugin.name not in self.skipped_plugins
|
|
]
|
|
return SimpleNamespace(
|
|
groups=[],
|
|
plugins=planned_plugins,
|
|
plugin_to_group={},
|
|
skipped_plugins=dict(self.skipped_plugins),
|
|
)
|
|
|
|
def prepare_environment(self, plugin) -> Path:
|
|
self.events.append(("prepare", plugin.name))
|
|
path = Path(sys.executable)
|
|
self.prepared_paths[plugin.name] = path
|
|
return path
|
|
|
|
|
|
class TestInstallSignalHandlers:
|
|
"""Tests for _install_signal_handlers function."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_installs_handlers(self):
|
|
"""_install_signal_handlers should install signal handlers."""
|
|
stop_event = asyncio.Event()
|
|
_install_signal_handlers(stop_event)
|
|
# Just verify it doesn't raise on platforms that support it
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handles_not_implemented(self):
|
|
"""_install_signal_handlers should handle NotImplementedError."""
|
|
stop_event = asyncio.Event()
|
|
with patch.object(
|
|
asyncio.get_running_loop(),
|
|
"add_signal_handler",
|
|
side_effect=NotImplementedError,
|
|
):
|
|
_install_signal_handlers(stop_event)
|
|
|
|
|
|
class TestPrepareStdioTransport:
|
|
"""Tests for _prepare_stdio_transport function."""
|
|
|
|
def test_with_both_streams(self):
|
|
"""_prepare_stdio_transport should use provided streams."""
|
|
stdin = StringIO()
|
|
stdout = StringIO()
|
|
|
|
in_stream, out_stream, original = _prepare_stdio_transport(stdin, stdout)
|
|
|
|
assert in_stream is stdin
|
|
assert out_stream is stdout
|
|
assert original is None
|
|
|
|
def test_without_streams(self):
|
|
"""_prepare_stdio_transport should use sys.stdin/stdout."""
|
|
# 保存原始值
|
|
original_stdin = sys.stdin
|
|
original_stdout = sys.stdout
|
|
|
|
try:
|
|
in_stream, out_stream, original = _prepare_stdio_transport(None, None)
|
|
|
|
# in_stream 应该是原始的 sys.stdin
|
|
assert in_stream is original_stdin
|
|
# out_stream 应该是原始的 sys.stdout(在修改前)
|
|
assert out_stream is original_stdout
|
|
# original 也应该是原始的 sys.stdout
|
|
assert original is original_stdout
|
|
# 函数会修改 sys.stdout 为 sys.stderr
|
|
assert sys.stdout is sys.stderr
|
|
finally:
|
|
# 恢复
|
|
sys.stdout = original_stdout
|
|
|
|
def test_redirects_stdout(self):
|
|
"""_prepare_stdio_transport should redirect sys.stdout to stderr."""
|
|
original_stdout = sys.stdout
|
|
|
|
_prepare_stdio_transport(None, None)
|
|
|
|
assert sys.stdout is sys.stderr
|
|
|
|
# Restore
|
|
sys.stdout = original_stdout
|
|
|
|
|
|
class TestWaitForShutdown:
|
|
"""Tests for _wait_for_shutdown function."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_waits_for_stop_event(self):
|
|
"""_wait_for_shutdown should wait for stop_event."""
|
|
peer = MagicMock()
|
|
|
|
# wait_closed 应该返回一个永不完成的协程
|
|
async def never_complete():
|
|
await asyncio.sleep(3600)
|
|
|
|
peer.wait_closed = MagicMock(return_value=never_complete())
|
|
|
|
stop_event = asyncio.Event()
|
|
|
|
async def set_event():
|
|
await asyncio.sleep(0.05)
|
|
stop_event.set()
|
|
|
|
asyncio.create_task(set_event())
|
|
|
|
await _wait_for_shutdown(peer, stop_event)
|
|
|
|
assert stop_event.is_set()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_waits_for_peer_closed(self):
|
|
"""_wait_for_shutdown should wait for peer.wait_closed()."""
|
|
peer = MagicMock()
|
|
|
|
# wait_closed 应该返回一个会完成的协程
|
|
async def complete_soon():
|
|
await asyncio.sleep(0.05)
|
|
|
|
peer.wait_closed = MagicMock(return_value=complete_soon())
|
|
|
|
stop_event = asyncio.Event()
|
|
|
|
await _wait_for_shutdown(peer, stop_event)
|
|
|
|
peer.wait_closed.assert_called_once()
|
|
|
|
|
|
class TestWorkerSessionInit:
|
|
"""Tests for WorkerSession initialization."""
|
|
|
|
def test_init(self):
|
|
"""WorkerSession should store all parameters."""
|
|
plugin = PluginSpec(
|
|
name="test",
|
|
plugin_dir=Path("/tmp"),
|
|
manifest_path=Path("/tmp/plugin.yaml"),
|
|
requirements_path=Path("/tmp/requirements.txt"),
|
|
python_version="3.12",
|
|
manifest_data={},
|
|
)
|
|
router = CapabilityRouter()
|
|
env_manager = FakeEnvManager()
|
|
|
|
session = WorkerSession(
|
|
plugin=plugin,
|
|
repo_root=Path("/repo"),
|
|
env_manager=env_manager,
|
|
capability_router=router,
|
|
)
|
|
|
|
assert session.plugin == plugin
|
|
assert session.capability_router == router
|
|
assert session.peer is None
|
|
assert session.handlers == []
|
|
|
|
|
|
class TestWorkerSessionMethods:
|
|
"""Tests for WorkerSession methods."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_invoke_handler_without_peer_raises(self):
|
|
"""invoke_handler should raise if peer is None."""
|
|
plugin = PluginSpec(
|
|
name="test",
|
|
plugin_dir=Path("/tmp"),
|
|
manifest_path=Path("/tmp/plugin.yaml"),
|
|
requirements_path=Path("/tmp/requirements.txt"),
|
|
python_version="3.12",
|
|
manifest_data={},
|
|
)
|
|
session = WorkerSession(
|
|
plugin=plugin,
|
|
repo_root=Path("/tmp"),
|
|
env_manager=FakeEnvManager(),
|
|
capability_router=CapabilityRouter(),
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="not running"):
|
|
await session.invoke_handler("handler.id", {}, request_id="req-1")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_without_peer_does_nothing(self):
|
|
"""cancel should do nothing if peer is None."""
|
|
plugin = PluginSpec(
|
|
name="test",
|
|
plugin_dir=Path("/tmp"),
|
|
manifest_path=Path("/tmp/plugin.yaml"),
|
|
requirements_path=Path("/tmp/requirements.txt"),
|
|
python_version="3.12",
|
|
manifest_data={},
|
|
)
|
|
session = WorkerSession(
|
|
plugin=plugin,
|
|
repo_root=Path("/tmp"),
|
|
env_manager=FakeEnvManager(),
|
|
capability_router=CapabilityRouter(),
|
|
)
|
|
|
|
# Should not raise
|
|
await session.cancel("req-1")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_initialize(self):
|
|
"""_handle_initialize should return InitializeOutput."""
|
|
plugin = PluginSpec(
|
|
name="test_plugin",
|
|
plugin_dir=Path("/tmp"),
|
|
manifest_path=Path("/tmp/plugin.yaml"),
|
|
requirements_path=Path("/tmp/requirements.txt"),
|
|
python_version="3.12",
|
|
manifest_data={},
|
|
)
|
|
router = CapabilityRouter()
|
|
session = WorkerSession(
|
|
plugin=plugin,
|
|
repo_root=Path("/tmp"),
|
|
env_manager=FakeEnvManager(),
|
|
capability_router=router,
|
|
)
|
|
|
|
message = InitializeMessage(
|
|
id="init-1",
|
|
protocol_version="1.0",
|
|
peer=PeerInfo(name="test", role="plugin"),
|
|
)
|
|
|
|
output = await session._handle_initialize(message)
|
|
|
|
assert output.peer.name == "astrbot-supervisor"
|
|
assert output.peer.role == "core"
|
|
assert len(output.capabilities) > 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_capability_invoke(self):
|
|
"""_handle_capability_invoke should route to capability_router."""
|
|
plugin = PluginSpec(
|
|
name="test",
|
|
plugin_dir=Path("/tmp"),
|
|
manifest_path=Path("/tmp/plugin.yaml"),
|
|
requirements_path=Path("/tmp/requirements.txt"),
|
|
python_version="3.12",
|
|
manifest_data={},
|
|
)
|
|
router = CapabilityRouter()
|
|
session = WorkerSession(
|
|
plugin=plugin,
|
|
repo_root=Path("/tmp"),
|
|
env_manager=FakeEnvManager(),
|
|
capability_router=router,
|
|
)
|
|
|
|
message = InvokeMessage(
|
|
id="invoke-1",
|
|
capability="llm.chat",
|
|
input={"prompt": "hello"},
|
|
)
|
|
token = CancelToken()
|
|
|
|
result = await session._handle_capability_invoke(message, token)
|
|
|
|
assert result["text"] == "Echo: hello"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_plugin_capability_routes_through_worker_session(self):
|
|
"""SupervisorRuntime should expose plugin-provided capabilities via router."""
|
|
transport = MemoryTransport()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=transport,
|
|
plugins_dir=Path(temp_dir),
|
|
)
|
|
session = MagicMock()
|
|
session.invoke_capability = AsyncMock(return_value={"echo": "hi"})
|
|
descriptor = CapabilityDescriptor(
|
|
name="demo.echo",
|
|
description="Echo text",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {"text": {"type": "string"}},
|
|
},
|
|
output_schema={
|
|
"type": "object",
|
|
"properties": {"echo": {"type": "string"}},
|
|
},
|
|
)
|
|
|
|
runtime._register_plugin_capability(descriptor, session, "demo_plugin")
|
|
result = await runtime.capability_router.execute(
|
|
"demo.echo",
|
|
{"text": "hi"},
|
|
stream=False,
|
|
cancel_token=CancelToken(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
assert result == {"echo": "hi"}
|
|
session.invoke_capability.assert_awaited_once_with(
|
|
"demo.echo",
|
|
{"text": "hi"},
|
|
request_id="req-1",
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_plugin_stream_capability_preserves_completed_output(self):
|
|
"""SupervisorRuntime should preserve plugin stream capability finalize output."""
|
|
transport = MemoryTransport()
|
|
|
|
async def stream_events():
|
|
yield EventMessage(id="req-1", phase="delta", data={"chunk": 1})
|
|
yield EventMessage(id="req-1", phase="delta", data={"chunk": 2})
|
|
yield EventMessage(
|
|
id="req-1",
|
|
phase="completed",
|
|
output={"count": 2},
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=transport,
|
|
plugins_dir=Path(temp_dir),
|
|
)
|
|
session = MagicMock()
|
|
session.invoke_capability_stream = MagicMock(return_value=stream_events())
|
|
descriptor = CapabilityDescriptor(
|
|
name="demo.stream",
|
|
description="Stream text",
|
|
supports_stream=True,
|
|
output_schema={
|
|
"type": "object",
|
|
"properties": {"count": {"type": "integer"}},
|
|
"required": ["count"],
|
|
},
|
|
)
|
|
|
|
runtime._register_plugin_capability(descriptor, session, "demo_plugin")
|
|
result = await runtime.capability_router.execute(
|
|
"demo.stream",
|
|
{},
|
|
stream=True,
|
|
cancel_token=CancelToken(),
|
|
request_id="req-1",
|
|
)
|
|
|
|
chunks = []
|
|
async for chunk in result.iterator:
|
|
chunks.append(chunk)
|
|
|
|
assert chunks == [{"chunk": 1}, {"chunk": 2}]
|
|
assert result.finalize(chunks) == {"count": 2}
|
|
|
|
|
|
class TestSupervisorRuntimeInit:
|
|
"""Tests for SupervisorRuntime initialization."""
|
|
|
|
def test_init(self):
|
|
"""SupervisorRuntime should initialize correctly."""
|
|
transport = MemoryTransport()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=transport,
|
|
plugins_dir=Path(temp_dir),
|
|
)
|
|
|
|
assert runtime.transport is transport
|
|
assert runtime.worker_sessions == {}
|
|
assert runtime.handler_to_worker == {}
|
|
assert runtime.capability_to_worker == {}
|
|
assert runtime.active_requests == {}
|
|
assert runtime.loaded_plugins == []
|
|
assert isinstance(runtime.capability_router, CapabilityRouter)
|
|
|
|
def test_registers_internal_capabilities(self):
|
|
"""SupervisorRuntime should register internal capabilities."""
|
|
transport = MemoryTransport()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=transport,
|
|
plugins_dir=Path(temp_dir),
|
|
)
|
|
|
|
# handler.invoke should be registered (but not exposed)
|
|
assert "handler.invoke" in runtime.capability_router._registrations
|
|
# Should not be in descriptors (exposed=False)
|
|
names = [d.name for d in runtime.capability_router.descriptors()]
|
|
assert "handler.invoke" not in names
|
|
|
|
|
|
class TestSupervisorRuntimeMethods:
|
|
"""Tests for SupervisorRuntime methods."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_with_empty_plugins_dir(self):
|
|
"""SupervisorRuntime.start should work with empty plugins dir."""
|
|
left, right = make_transport_pair()
|
|
core = await start_test_core_peer(left)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=right,
|
|
plugins_dir=Path(temp_dir),
|
|
env_manager=FakeEnvManager(),
|
|
)
|
|
|
|
try:
|
|
await runtime.start()
|
|
await core.wait_until_remote_initialized()
|
|
assert runtime.loaded_plugins == []
|
|
assert runtime.skipped_plugins == {}
|
|
finally:
|
|
await runtime.stop()
|
|
await core.stop()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_calls_plan_before_prepare_and_reuses_python_path(self):
|
|
"""SupervisorRuntime should plan plugins before starting worker sessions."""
|
|
left, right = make_transport_pair()
|
|
core = await start_test_core_peer(left)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugins_dir = Path(temp_dir) / "plugins"
|
|
write_bootstrap_plugin(plugins_dir, "plugin_one")
|
|
write_bootstrap_plugin(plugins_dir, "plugin_two")
|
|
env_manager = PlanningFakeEnvManager()
|
|
runtime = SupervisorRuntime(
|
|
transport=right,
|
|
plugins_dir=plugins_dir,
|
|
env_manager=env_manager,
|
|
)
|
|
|
|
try:
|
|
await runtime.start()
|
|
await core.wait_until_remote_initialized()
|
|
|
|
assert env_manager.events[0] == ("plan", ["plugin_one", "plugin_two"])
|
|
assert ("prepare", "plugin_one") in env_manager.events
|
|
assert ("prepare", "plugin_two") in env_manager.events
|
|
assert env_manager.prepared_paths["plugin_one"] == Path(sys.executable)
|
|
assert env_manager.prepared_paths["plugin_two"] == Path(sys.executable)
|
|
assert runtime.loaded_plugins == ["plugin_one", "plugin_two"]
|
|
finally:
|
|
await runtime.stop()
|
|
await core.stop()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_merges_planning_skips_into_runtime_metadata(self):
|
|
"""SupervisorRuntime should expose planning-stage skipped plugins."""
|
|
left, right = make_transport_pair()
|
|
core = await start_test_core_peer(left)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugins_dir = Path(temp_dir) / "plugins"
|
|
write_bootstrap_plugin(plugins_dir, "plugin_one")
|
|
write_bootstrap_plugin(plugins_dir, "plugin_two")
|
|
env_manager = PlanningFakeEnvManager(
|
|
skipped_plugins={"plugin_two": "compile lockfile failed"}
|
|
)
|
|
runtime = SupervisorRuntime(
|
|
transport=right,
|
|
plugins_dir=plugins_dir,
|
|
env_manager=env_manager,
|
|
)
|
|
|
|
try:
|
|
await runtime.start()
|
|
await core.wait_until_remote_initialized()
|
|
|
|
assert runtime.loaded_plugins == ["plugin_one"]
|
|
assert (
|
|
runtime.skipped_plugins["plugin_two"] == "compile lockfile failed"
|
|
)
|
|
assert core.remote_metadata["skipped_plugins"]["plugin_two"] == (
|
|
"compile lockfile failed"
|
|
)
|
|
finally:
|
|
await runtime.stop()
|
|
await core.stop()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_route_handler_invoke_missing_handler(self):
|
|
"""_route_handler_invoke should raise for missing handler."""
|
|
transport = MemoryTransport()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=transport,
|
|
plugins_dir=Path(temp_dir),
|
|
)
|
|
|
|
with pytest.raises(AstrBotError, match="handler not found"):
|
|
await runtime._route_handler_invoke(
|
|
"req-1",
|
|
{"handler_id": "missing.handler", "event": {}},
|
|
CancelToken(),
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_worker_closed_removes_session(self):
|
|
"""_handle_worker_closed should remove session and handlers."""
|
|
left, right = make_transport_pair()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=right,
|
|
plugins_dir=Path(temp_dir),
|
|
env_manager=FakeEnvManager(),
|
|
)
|
|
|
|
# Add fake session
|
|
mock_session = MagicMock()
|
|
mock_session.handlers = [
|
|
HandlerDescriptor(
|
|
id="test.handler", trigger=CommandTrigger(command="test")
|
|
)
|
|
]
|
|
|
|
runtime.worker_sessions["test_plugin"] = mock_session
|
|
runtime.handler_to_worker["test.handler"] = mock_session
|
|
runtime._handler_sources["test.handler"] = "test_plugin"
|
|
runtime.loaded_plugins.append("test_plugin")
|
|
|
|
runtime._handle_worker_closed("test_plugin")
|
|
|
|
assert "test_plugin" not in runtime.worker_sessions
|
|
assert "test.handler" not in runtime.handler_to_worker
|
|
assert "test.handler" not in runtime._handler_sources
|
|
assert "test_plugin" not in runtime.loaded_plugins
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_worker_closed_removes_active_requests(self):
|
|
"""_handle_worker_closed should drop in-flight requests owned by the worker."""
|
|
transport = MemoryTransport()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=transport,
|
|
plugins_dir=Path(temp_dir),
|
|
)
|
|
|
|
mock_session = MagicMock()
|
|
mock_session.handlers = [
|
|
HandlerDescriptor(
|
|
id="test.handler", trigger=CommandTrigger(command="test")
|
|
)
|
|
]
|
|
|
|
runtime.worker_sessions["test_plugin"] = mock_session
|
|
runtime.handler_to_worker["test.handler"] = mock_session
|
|
runtime._handler_sources["test.handler"] = "test_plugin"
|
|
runtime.loaded_plugins.append("test_plugin")
|
|
runtime.active_requests["req-1"] = mock_session
|
|
|
|
runtime._handle_worker_closed("test_plugin")
|
|
|
|
assert "req-1" not in runtime.active_requests
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_worker_closed_unknown_plugin(self):
|
|
"""_handle_worker_closed should handle unknown plugin."""
|
|
transport = MemoryTransport()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=transport,
|
|
plugins_dir=Path(temp_dir),
|
|
)
|
|
|
|
# Should not raise for unknown plugin
|
|
runtime._handle_worker_closed("unknown_plugin")
|
|
|
|
|
|
class TestPluginWorkerRuntimeInit:
|
|
"""Tests for PluginWorkerRuntime initialization."""
|
|
|
|
def test_init_with_valid_plugin(self):
|
|
"""PluginWorkerRuntime should initialize with valid plugin."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
transport = MemoryTransport()
|
|
|
|
# This should work if the plugin is valid
|
|
runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport)
|
|
|
|
assert runtime.plugin.name == "test_plugin"
|
|
assert runtime.peer is not None
|
|
assert runtime.dispatcher is not None
|
|
|
|
|
|
class TestPluginWorkerRuntimeMethods:
|
|
"""Tests for PluginWorkerRuntime methods."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_invoke_wrong_capability(self):
|
|
"""_handle_invoke should raise for wrong capability."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
transport = MemoryTransport()
|
|
runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport)
|
|
|
|
message = InvokeMessage(
|
|
id="invoke-1",
|
|
capability="wrong.capability",
|
|
input={},
|
|
)
|
|
token = CancelToken()
|
|
|
|
with pytest.raises(AstrBotError, match="未找到能力"):
|
|
await runtime._handle_invoke(message, token)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_invoke_plugin_capability(self):
|
|
"""_handle_invoke should dispatch plugin-provided capabilities."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
module_dir = plugin_dir / "demo_plugin"
|
|
module_dir.mkdir()
|
|
(module_dir / "__init__.py").write_text("", encoding="utf-8")
|
|
(module_dir / "component.py").write_text(
|
|
"""
|
|
from astrbot_sdk import Star, provide_capability
|
|
|
|
|
|
class DemoComponent(Star):
|
|
@provide_capability(
|
|
"demo.echo",
|
|
description="Echo text",
|
|
input_schema={"type": "object", "properties": {"text": {"type": "string"}}},
|
|
output_schema={"type": "object", "properties": {"echo": {"type": "string"}}},
|
|
)
|
|
async def echo(self, payload):
|
|
return {"echo": payload["text"]}
|
|
""".strip(),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [
|
|
{"class": "demo_plugin.component:DemoComponent"}
|
|
],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
if str(plugin_dir) not in sys.path:
|
|
sys.path.insert(0, str(plugin_dir))
|
|
|
|
try:
|
|
transport = MemoryTransport()
|
|
runtime = PluginWorkerRuntime(
|
|
plugin_dir=plugin_dir,
|
|
transport=transport,
|
|
)
|
|
message = InvokeMessage(
|
|
id="invoke-cap",
|
|
capability="demo.echo",
|
|
input={"text": "hello"},
|
|
)
|
|
|
|
result = await runtime._handle_invoke(message, CancelToken())
|
|
|
|
assert result == {"echo": "hello"}
|
|
finally:
|
|
if str(plugin_dir) in sys.path:
|
|
sys.path.remove(str(plugin_dir))
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_and_stop_run_compat_context_lifecycle_hooks(self):
|
|
"""PluginWorkerRuntime should execute compat lifecycle hooks around peer startup/shutdown."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
transport = MemoryTransport()
|
|
runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport)
|
|
|
|
seen_hooks = []
|
|
legacy_context = LegacyContext("test_plugin")
|
|
|
|
class CompatHooks:
|
|
async def on_astrbot_loaded(self, context):
|
|
seen_hooks.append(("astrbot", context.plugin_id))
|
|
|
|
async def on_platform_loaded(self, context):
|
|
seen_hooks.append(("platform", context.plugin_id))
|
|
|
|
async def on_plugin_loaded(self, metadata):
|
|
seen_hooks.append(("loaded", metadata["name"]))
|
|
|
|
async def on_plugin_unloaded(self, metadata):
|
|
seen_hooks.append(("unloaded", metadata["name"]))
|
|
|
|
from astrbot_sdk.api.event.filter import (
|
|
on_astrbot_loaded,
|
|
on_platform_loaded,
|
|
on_plugin_loaded,
|
|
on_plugin_unloaded,
|
|
)
|
|
|
|
CompatHooks.on_astrbot_loaded = on_astrbot_loaded()(
|
|
CompatHooks.on_astrbot_loaded
|
|
)
|
|
CompatHooks.on_platform_loaded = on_platform_loaded()(
|
|
CompatHooks.on_platform_loaded
|
|
)
|
|
CompatHooks.on_plugin_loaded = on_plugin_loaded()(
|
|
CompatHooks.on_plugin_loaded
|
|
)
|
|
CompatHooks.on_plugin_unloaded = on_plugin_unloaded()(
|
|
CompatHooks.on_plugin_unloaded
|
|
)
|
|
legacy_context._register_component(CompatHooks())
|
|
legacy_context.bind_runtime_context(runtime._lifecycle_context)
|
|
|
|
runtime.loaded_plugin.handlers.append(
|
|
LoadedHandler(
|
|
descriptor=HandlerDescriptor(
|
|
id="legacy.compat",
|
|
trigger=CommandTrigger(command="compat"),
|
|
),
|
|
callable=AsyncMock(),
|
|
owner=MagicMock(),
|
|
legacy_context=legacy_context,
|
|
)
|
|
)
|
|
runtime.peer.start = AsyncMock()
|
|
runtime.peer.initialize = AsyncMock()
|
|
runtime.peer.stop = AsyncMock()
|
|
|
|
await runtime.start()
|
|
await runtime.stop()
|
|
|
|
assert seen_hooks == [
|
|
("astrbot", "test_plugin"),
|
|
("platform", "test_plugin"),
|
|
("loaded", "test_plugin"),
|
|
("unloaded", "test_plugin"),
|
|
]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_uses_legacy_runtime_boundary_for_startup_hooks(self):
|
|
"""PluginWorkerRuntime startup should delegate compat lifecycle work to the adapter boundary."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_data = {
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
manifest_path.write_text(
|
|
yaml.dump(manifest_data),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
runtime = PluginWorkerRuntime(
|
|
plugin_dir=plugin_dir,
|
|
transport=MemoryTransport(),
|
|
)
|
|
runtime.peer.start = AsyncMock()
|
|
runtime.peer.initialize = AsyncMock()
|
|
runtime.peer.stop = AsyncMock()
|
|
|
|
with patch(
|
|
"astrbot_sdk.runtime.bootstrap.run_legacy_worker_startup_hooks",
|
|
new=AsyncMock(),
|
|
) as startup_hooks:
|
|
await runtime.start()
|
|
|
|
startup_hooks.assert_awaited_once()
|
|
args, kwargs = startup_hooks.await_args
|
|
assert args == ([],)
|
|
assert kwargs["context"] is runtime._lifecycle_context
|
|
assert kwargs["metadata"] == manifest_data
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stop_uses_legacy_runtime_boundary_for_shutdown_hooks(self):
|
|
"""PluginWorkerRuntime shutdown should delegate compat unload hooks to the adapter boundary."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_data = {
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
manifest_path.write_text(
|
|
yaml.dump(manifest_data),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
runtime = PluginWorkerRuntime(
|
|
plugin_dir=plugin_dir,
|
|
transport=MemoryTransport(),
|
|
)
|
|
runtime.peer.stop = AsyncMock()
|
|
|
|
with patch(
|
|
"astrbot_sdk.runtime.bootstrap.run_legacy_worker_shutdown_hooks",
|
|
new=AsyncMock(),
|
|
) as shutdown_hooks:
|
|
await runtime.stop()
|
|
|
|
shutdown_hooks.assert_awaited_once()
|
|
args, kwargs = shutdown_hooks.await_args
|
|
assert args == ([],)
|
|
assert kwargs["context"] is runtime._lifecycle_context
|
|
assert kwargs["metadata"] == manifest_data
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_lifecycle_sync_hook(self):
|
|
"""_run_lifecycle should call sync hooks."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
transport = MemoryTransport()
|
|
runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport)
|
|
|
|
# Add mock instance with sync hook
|
|
called = []
|
|
|
|
class MockInstance:
|
|
def on_start(self, ctx):
|
|
called.append("on_start")
|
|
|
|
runtime.loaded_plugin.instances.append(MockInstance())
|
|
|
|
await runtime._run_lifecycle("on_start")
|
|
|
|
assert "on_start" in called
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_lifecycle_async_hook(self):
|
|
"""_run_lifecycle should call async hooks."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
transport = MemoryTransport()
|
|
runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport)
|
|
|
|
# Add mock instance with async hook
|
|
called = []
|
|
|
|
class MockInstance:
|
|
async def on_stop(self, ctx):
|
|
called.append("on_stop")
|
|
|
|
runtime.loaded_plugin.instances.append(MockInstance())
|
|
|
|
await runtime._run_lifecycle("on_stop")
|
|
|
|
assert "on_stop" in called
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_lifecycle_missing_method(self):
|
|
"""_run_lifecycle should skip missing methods."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
transport = MemoryTransport()
|
|
runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport)
|
|
|
|
# Add mock instance without the method
|
|
class MockInstance:
|
|
pass
|
|
|
|
runtime.loaded_plugin.instances.append(MockInstance())
|
|
|
|
# Should not raise
|
|
await runtime._run_lifecycle("on_start")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_lifecycle_legacy_initialize_and_terminate_aliases(self):
|
|
"""_run_lifecycle should map legacy initialize/terminate for old stars."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
transport = MemoryTransport()
|
|
runtime = PluginWorkerRuntime(plugin_dir=plugin_dir, transport=transport)
|
|
|
|
called = []
|
|
|
|
class MockLegacyInstance:
|
|
@classmethod
|
|
def __astrbot_is_new_star__(cls):
|
|
return False
|
|
|
|
async def initialize(self):
|
|
called.append("initialize")
|
|
|
|
async def terminate(self):
|
|
called.append("terminate")
|
|
|
|
runtime.loaded_plugin.instances.append(MockLegacyInstance())
|
|
|
|
await runtime._run_lifecycle("on_start")
|
|
await runtime._run_lifecycle("on_stop")
|
|
|
|
assert called == ["initialize", "terminate"]
|
|
|
|
def test_bind_legacy_runtime_contexts_reuses_shared_context(self):
|
|
"""legacy lifecycle helpers should see the worker runtime context."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_dir = Path(temp_dir)
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {"python": "3.12"},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
runtime = PluginWorkerRuntime(
|
|
plugin_dir=plugin_dir, transport=MemoryTransport()
|
|
)
|
|
legacy_context = LegacyContext("test_plugin")
|
|
runtime.loaded_plugin.handlers.append(
|
|
SimpleNamespace(
|
|
legacy_runtime=LegacyRuntimeAdapter(legacy_context=legacy_context)
|
|
)
|
|
)
|
|
runtime.loaded_plugin.capabilities.append(
|
|
SimpleNamespace(
|
|
legacy_runtime=LegacyRuntimeAdapter(legacy_context=legacy_context)
|
|
)
|
|
)
|
|
|
|
runtime._bind_legacy_runtime_contexts(runtime._lifecycle_context)
|
|
|
|
assert (
|
|
legacy_context.require_runtime_context() is runtime._lifecycle_context
|
|
)
|
|
|
|
|
|
class TestIntegrationWithTransportPair:
|
|
"""Integration tests using transport pairs."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_supervisor_responds_to_initialize(self):
|
|
"""SupervisorRuntime should respond to initialize messages."""
|
|
left, right = make_transport_pair()
|
|
core = await start_test_core_peer(left)
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
runtime = SupervisorRuntime(
|
|
transport=right,
|
|
plugins_dir=Path(temp_dir),
|
|
env_manager=FakeEnvManager(),
|
|
)
|
|
|
|
try:
|
|
await runtime.start()
|
|
await core.wait_until_remote_initialized()
|
|
assert core.remote_peer is not None
|
|
assert core.remote_peer.name == "astrbot-supervisor"
|
|
assert core.remote_metadata["plugins"] == []
|
|
assert core.remote_metadata["skipped_plugins"] == {}
|
|
|
|
finally:
|
|
await runtime.stop()
|
|
await core.stop()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_worker_session_lifecycle(self):
|
|
"""WorkerSession should start and stop cleanly."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
# Create a minimal plugin
|
|
plugin_dir = Path(temp_dir) / "plugins" / "test_plugin"
|
|
plugin_dir.mkdir(parents=True)
|
|
|
|
manifest_path = plugin_dir / "plugin.yaml"
|
|
requirements_path = plugin_dir / "requirements.txt"
|
|
|
|
manifest_path.write_text(
|
|
yaml.dump(
|
|
{
|
|
"name": "test_plugin",
|
|
"runtime": {
|
|
"python": f"{sys.version_info.major}.{sys.version_info.minor}"
|
|
},
|
|
"components": [],
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
requirements_path.write_text("", encoding="utf-8")
|
|
|
|
plugin = PluginSpec(
|
|
name="test_plugin",
|
|
plugin_dir=plugin_dir,
|
|
manifest_path=manifest_path,
|
|
requirements_path=requirements_path,
|
|
python_version=f"{sys.version_info.major}.{sys.version_info.minor}",
|
|
manifest_data={"name": "test_plugin"},
|
|
)
|
|
|
|
left, right = make_transport_pair()
|
|
|
|
session = WorkerSession(
|
|
plugin=plugin,
|
|
repo_root=Path(temp_dir),
|
|
env_manager=FakeEnvManager(),
|
|
capability_router=CapabilityRouter(),
|
|
)
|
|
|
|
# Note: Full start would require subprocess, skip for unit test
|
|
# Just verify the session can be created and stopped
|
|
await session.stop()
|