mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 01:40:15 +08:00
306 lines
11 KiB
Python
306 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import sys
|
|
import tempfile
|
|
import textwrap
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from astrbot_sdk.protocol.messages import InitializeOutput, PeerInfo
|
|
from astrbot_sdk.runtime.bootstrap import PluginWorkerRuntime, SupervisorRuntime
|
|
from astrbot_sdk.runtime.capability_router import CapabilityRouter
|
|
from astrbot_sdk.runtime.peer import Peer
|
|
from astrbot_sdk.runtime.transport import (
|
|
WebSocketClientTransport,
|
|
WebSocketServerTransport,
|
|
)
|
|
|
|
from tests_v4.helpers import FakeEnvManager, make_transport_pair
|
|
|
|
|
|
def write_websocket_plugin(plugin_root: Path) -> None:
|
|
(plugin_root / "commands").mkdir(parents=True, exist_ok=True)
|
|
(plugin_root / "commands" / "__init__.py").write_text("", encoding="utf-8")
|
|
(plugin_root / "requirements.txt").write_text("", encoding="utf-8")
|
|
(plugin_root / "plugin.yaml").write_text(
|
|
textwrap.dedent(
|
|
f"""\
|
|
_schema_version: 2
|
|
name: websocket_plugin
|
|
display_name: WebSocket Plugin
|
|
desc: websocket test
|
|
author: tester
|
|
version: 0.1.0
|
|
runtime:
|
|
python: "{sys.version_info.major}.{sys.version_info.minor}"
|
|
components:
|
|
- class: commands.sample:MyPlugin
|
|
type: command
|
|
name: hello
|
|
description: hello
|
|
"""
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
(plugin_root / "commands" / "sample.py").write_text(
|
|
textwrap.dedent(
|
|
"""\
|
|
from astrbot_sdk import Context, MessageEvent, Star, on_command
|
|
|
|
|
|
class MyPlugin(Star):
|
|
@on_command("hello")
|
|
async def hello(self, event: MessageEvent, ctx: Context):
|
|
await event.reply(f"ws:{event.text}")
|
|
"""
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def write_benchmark_plugin(plugins_dir: Path, index: int) -> None:
|
|
plugin_name = f"plugin_{index:03d}"
|
|
command_name = f"bench_{index:03d}"
|
|
plugin_dir = plugins_dir / plugin_name
|
|
commands_dir = plugin_dir / "commands"
|
|
commands_dir.mkdir(parents=True, exist_ok=True)
|
|
(commands_dir / "__init__.py").write_text("", encoding="utf-8")
|
|
(plugin_dir / "requirements.txt").write_text("", encoding="utf-8")
|
|
(plugin_dir / "plugin.yaml").write_text(
|
|
textwrap.dedent(
|
|
f"""\
|
|
_schema_version: 2
|
|
name: {plugin_name}
|
|
display_name: {plugin_name}
|
|
desc: benchmark plugin {index}
|
|
author: tester
|
|
version: 0.1.0
|
|
runtime:
|
|
python: "{sys.version_info.major}.{sys.version_info.minor}"
|
|
components:
|
|
- class: commands.plugin_{index:03d}:BenchmarkCommand{index:03d}
|
|
type: command
|
|
name: {command_name}
|
|
description: {command_name}
|
|
"""
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
(commands_dir / f"plugin_{index:03d}.py").write_text(
|
|
textwrap.dedent(
|
|
f"""\
|
|
from astrbot_sdk.api.components.command import CommandComponent
|
|
from astrbot_sdk.api.event import AstrMessageEvent, filter
|
|
from astrbot_sdk.api.star.context import Context
|
|
|
|
|
|
class BenchmarkCommand{index:03d}(CommandComponent):
|
|
def __init__(self, context: Context):
|
|
self.context = context
|
|
|
|
@filter.command("{command_name}")
|
|
async def handle(self, event: AstrMessageEvent):
|
|
yield event.plain_result("{plugin_name}:{command_name}")
|
|
"""
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
class StartClientMigrationTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_websocket_plugin_worker_supports_handshake_and_handler_invoke(
|
|
self,
|
|
) -> None:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugin_root = Path(temp_dir) / "websocket_plugin"
|
|
write_websocket_plugin(plugin_root)
|
|
|
|
server_transport = WebSocketServerTransport(
|
|
host="127.0.0.1", port=0, path="/ws"
|
|
)
|
|
runtime = PluginWorkerRuntime(
|
|
plugin_dir=plugin_root, transport=server_transport
|
|
)
|
|
runtime_task = asyncio.create_task(runtime.start())
|
|
|
|
try:
|
|
for _ in range(100):
|
|
if server_transport.port != 0:
|
|
break
|
|
await asyncio.sleep(0.02)
|
|
|
|
core_router = CapabilityRouter()
|
|
core_peer = Peer(
|
|
transport=WebSocketClientTransport(url=server_transport.url),
|
|
peer_info=PeerInfo(
|
|
name="websocket-core", role="core", version="v4"
|
|
),
|
|
)
|
|
core_peer.set_initialize_handler(
|
|
lambda _message: asyncio.sleep(
|
|
0,
|
|
result=InitializeOutput(
|
|
peer=PeerInfo(
|
|
name="websocket-core", role="core", version="v4"
|
|
),
|
|
capabilities=core_router.descriptors(),
|
|
metadata={},
|
|
),
|
|
)
|
|
)
|
|
core_peer.set_invoke_handler(
|
|
lambda message, cancel_token: core_router.execute(
|
|
message.capability,
|
|
message.input,
|
|
stream=message.stream,
|
|
cancel_token=cancel_token,
|
|
request_id=message.id,
|
|
)
|
|
)
|
|
await core_peer.start()
|
|
try:
|
|
await asyncio.wait_for(runtime_task, timeout=5)
|
|
await core_peer.wait_until_remote_initialized()
|
|
handler_id = core_peer.remote_handlers[0].id
|
|
self.assertEqual(
|
|
core_peer.remote_metadata["plugin_id"], "websocket_plugin"
|
|
)
|
|
self.assertEqual(
|
|
core_peer.remote_handlers[0].trigger.command, "hello"
|
|
)
|
|
|
|
await core_peer.invoke(
|
|
"handler.invoke",
|
|
{
|
|
"handler_id": handler_id,
|
|
"event": {
|
|
"text": "hello-websocket",
|
|
"session_id": "session-ws",
|
|
"user_id": "user-1",
|
|
"platform": "test",
|
|
},
|
|
},
|
|
request_id="call-ws",
|
|
)
|
|
|
|
self.assertEqual(
|
|
[item.get("text") for item in core_router.sent_messages],
|
|
["ws:hello-websocket"],
|
|
)
|
|
finally:
|
|
await core_peer.stop()
|
|
finally:
|
|
if not runtime_task.done():
|
|
runtime_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await runtime_task
|
|
await runtime.stop()
|
|
|
|
|
|
class BenchmarkMigrationTest(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
self.left, self.right = make_transport_pair()
|
|
self.core = Peer(
|
|
transport=self.left,
|
|
peer_info=PeerInfo(name="benchmark-core", role="core", version="v4"),
|
|
)
|
|
self.core.set_initialize_handler(
|
|
lambda _message: asyncio.sleep(
|
|
0,
|
|
result=InitializeOutput(
|
|
peer=PeerInfo(name="benchmark-core", role="core", version="v4"),
|
|
capabilities=[],
|
|
metadata={},
|
|
),
|
|
)
|
|
)
|
|
await self.core.start()
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
await self.core.stop()
|
|
|
|
async def test_benchmark_style_runtime_report_covers_multi_plugin_workers(
|
|
self,
|
|
) -> None:
|
|
plugin_count = 8
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
plugins_dir = Path(temp_dir) / "plugins"
|
|
for index in range(plugin_count):
|
|
write_benchmark_plugin(plugins_dir, index)
|
|
|
|
runtime = SupervisorRuntime(
|
|
transport=self.right,
|
|
plugins_dir=plugins_dir,
|
|
env_manager=FakeEnvManager(),
|
|
)
|
|
started_at = time.perf_counter()
|
|
try:
|
|
await runtime.start()
|
|
await self.core.wait_until_remote_initialized()
|
|
measured_at = time.perf_counter()
|
|
|
|
worker_pids = sorted(
|
|
process.pid
|
|
for session in runtime.worker_sessions.values()
|
|
if (
|
|
process := getattr(
|
|
getattr(session.peer, "transport", None), "_process", None
|
|
)
|
|
)
|
|
is not None
|
|
and process.returncode is None
|
|
)
|
|
report = {
|
|
"plugin_count": plugin_count,
|
|
"loaded_plugin_count": len(runtime.loaded_plugins),
|
|
"loaded_plugins": sorted(runtime.loaded_plugins),
|
|
"aggregated_handler_ids": list(
|
|
self.core.remote_metadata["aggregated_handler_ids"]
|
|
),
|
|
"startup_total_duration_ms": round(
|
|
(measured_at - started_at) * 1000, 2
|
|
),
|
|
"worker_pids": worker_pids,
|
|
}
|
|
|
|
handler_id = next(
|
|
item.id
|
|
for item in self.core.remote_handlers
|
|
if item.id.startswith("plugin_002:")
|
|
)
|
|
await self.core.invoke(
|
|
"handler.invoke",
|
|
{
|
|
"handler_id": handler_id,
|
|
"event": {
|
|
"text": "/bench_002",
|
|
"session_id": "bench-session",
|
|
"user_id": "user-1",
|
|
"platform": "test",
|
|
},
|
|
},
|
|
request_id="bench-call",
|
|
)
|
|
|
|
self.assertEqual(report["loaded_plugin_count"], plugin_count)
|
|
self.assertEqual(len(report["worker_pids"]), plugin_count)
|
|
self.assertEqual(len(report["aggregated_handler_ids"]), plugin_count)
|
|
self.assertEqual(
|
|
report["loaded_plugins"],
|
|
[f"plugin_{index:03d}" for index in range(plugin_count)],
|
|
)
|
|
self.assertGreaterEqual(report["startup_total_duration_ms"], 0)
|
|
self.assertIn(
|
|
"plugin_002:bench_002",
|
|
runtime.capability_router.sent_messages[-1]["text"],
|
|
)
|
|
finally:
|
|
await runtime.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|