mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 01:40:15 +08:00
补充插件分组环境测试覆盖
This commit is contained in:
561
tests_v4/benchmark_grouped_environment_stress.py
Normal file
561
tests_v4/benchmark_grouped_environment_stress.py
Normal file
@@ -0,0 +1,561 @@
|
||||
# ruff: noqa: E402
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import stat
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import textwrap
|
||||
import time
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
SRC_NEW_DIR = PROJECT_ROOT / "src-new"
|
||||
if str(PROJECT_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(PROJECT_ROOT))
|
||||
if str(SRC_NEW_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SRC_NEW_DIR))
|
||||
|
||||
try:
|
||||
import psutil
|
||||
except ImportError: # pragma: no cover - optional dependency
|
||||
psutil = None
|
||||
|
||||
from astrbot_sdk.protocol.messages import InitializeOutput, PeerInfo
|
||||
from astrbot_sdk.runtime.bootstrap import SupervisorRuntime
|
||||
from astrbot_sdk.runtime.loader import PluginEnvironmentManager
|
||||
from astrbot_sdk.runtime.peer import Peer
|
||||
|
||||
from tests_v4.helpers import make_transport_pair
|
||||
|
||||
DEFAULT_MULTIPLIERS = [1, 2, 4, 8, 12]
|
||||
DEFAULT_CONFLICT_COUNT = 1
|
||||
DEFAULT_COMPATIBLE_COUNT = 5
|
||||
SAMPLE_INTERVAL_SECONDS = 0.05
|
||||
EXACT_PIN_PATTERN = re.compile(r"^([A-Za-z0-9_.-]+)==([^\s;]+)$")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ProcessTreeSnapshot:
|
||||
collector: str
|
||||
process_count: int
|
||||
total_rss_bytes: int
|
||||
total_rss_mb: float
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class BenchmarkCaseResult:
|
||||
multiplier: int
|
||||
conflict_plugins: int
|
||||
compatible_plugins: int
|
||||
total_plugins: int
|
||||
group_count: int
|
||||
skipped_plugins: int
|
||||
startup_duration_ms: float
|
||||
steady_rss_mb: float
|
||||
peak_rss_mb: float
|
||||
process_count: int
|
||||
expected_groups: int
|
||||
|
||||
|
||||
class SyntheticGroupedEnvManager(PluginEnvironmentManager):
|
||||
"""用于 benchmark 的分组环境管理器。
|
||||
|
||||
这个实现保留真实的 supervisor 启动流程和插件分组规划,但把下面两类成
|
||||
本较高、且对本地压力测试不稳定的动作替换掉:
|
||||
|
||||
- `uv pip compile` 改为直接生成可重复的伪 lockfile
|
||||
- `uv venv` / `uv pip sync` 改为为每个分组创建一个指向当前解释器的路径
|
||||
|
||||
这样得到的结果更接近“分组规划 + worker 启动”的资源开销,而不是被
|
||||
外网索引、包下载或磁盘安装速度主导。
|
||||
"""
|
||||
|
||||
def __init__(self, repo_root: Path) -> None:
|
||||
super().__init__(repo_root, uv_binary="synthetic-uv")
|
||||
self._original_is_compatible = self._planner._is_compatible
|
||||
self._planner._is_compatible = self._synthetic_is_compatible
|
||||
self._planner._compile_lockfile = self._synthetic_compile_lockfile
|
||||
self._group_manager.prepare = self._synthetic_prepare_environment
|
||||
|
||||
def _synthetic_is_compatible(self, plugins) -> bool:
|
||||
requirement_lines = self._planner._collect_requirement_lines(plugins)
|
||||
if not requirement_lines:
|
||||
return True
|
||||
|
||||
merged = self._planner._merge_exact_requirements(requirement_lines)
|
||||
if merged is not None:
|
||||
return True
|
||||
|
||||
if all(EXACT_PIN_PATTERN.fullmatch(line) for line in requirement_lines):
|
||||
return False
|
||||
return self._original_is_compatible(plugins)
|
||||
|
||||
@staticmethod
|
||||
def _synthetic_compile_lockfile(
|
||||
*,
|
||||
source_path: Path,
|
||||
output_path: Path,
|
||||
python_version: str,
|
||||
) -> None:
|
||||
lines = []
|
||||
for raw_line in source_path.read_text(encoding="utf-8").splitlines():
|
||||
line = raw_line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
lines.append(line)
|
||||
output = [f"# synthetic lockfile for python {python_version}"]
|
||||
output.extend(sorted(dict.fromkeys(lines)))
|
||||
output_path.write_text("\n".join(output).rstrip() + "\n", encoding="utf-8")
|
||||
|
||||
@staticmethod
|
||||
def _synthetic_prepare_environment(group) -> Path:
|
||||
group.python_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if group.python_path.exists():
|
||||
return group.python_path
|
||||
|
||||
target = Path(sys.executable).resolve()
|
||||
if os.name == "nt":
|
||||
shutil.copy2(target, group.python_path)
|
||||
current_mode = group.python_path.stat().st_mode
|
||||
group.python_path.chmod(current_mode | stat.S_IEXEC)
|
||||
else:
|
||||
os.symlink(target, group.python_path)
|
||||
return group.python_path
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Benchmark v4 grouped plugin environments with configurable counts "
|
||||
"for conflicting and compatible plugins."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--multipliers",
|
||||
nargs="+",
|
||||
type=int,
|
||||
default=DEFAULT_MULTIPLIERS,
|
||||
help="Scale factors applied to the base plugin counts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--conflict-count",
|
||||
type=int,
|
||||
default=DEFAULT_CONFLICT_COUNT,
|
||||
help="Base count of conflicting plugins before applying multipliers.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--compatible-count",
|
||||
type=int,
|
||||
default=DEFAULT_COMPATIBLE_COUNT,
|
||||
help="Base count of compatible plugins before applying multipliers.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-json",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Optional path for the JSON benchmark report.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--keep-temp-dir",
|
||||
action="store_true",
|
||||
help="Keep the generated temporary workspace for inspection.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def write_benchmark_plugin(
|
||||
*,
|
||||
plugins_dir: Path,
|
||||
plugin_name: str,
|
||||
command_name: str,
|
||||
requirement: str,
|
||||
) -> None:
|
||||
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(requirement, encoding="utf-8")
|
||||
(plugin_dir / "plugin.yaml").write_text(
|
||||
textwrap.dedent(
|
||||
f"""\
|
||||
_schema_version: 2
|
||||
name: {plugin_name}
|
||||
display_name: {plugin_name}
|
||||
desc: grouped environment benchmark plugin
|
||||
author: codex
|
||||
version: 0.1.0
|
||||
runtime:
|
||||
python: "{sys.version_info.major}.{sys.version_info.minor}"
|
||||
components:
|
||||
- class: commands.main:BenchmarkCommand
|
||||
type: command
|
||||
name: {command_name}
|
||||
description: {command_name}
|
||||
"""
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(commands_dir / "main.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(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",
|
||||
)
|
||||
|
||||
|
||||
def create_plugin_matrix(
|
||||
*,
|
||||
plugins_dir: Path,
|
||||
multiplier: int,
|
||||
conflict_base_count: int,
|
||||
compatible_base_count: int,
|
||||
) -> tuple[int, int]:
|
||||
conflict_count = conflict_base_count * multiplier
|
||||
compatible_count = compatible_base_count * multiplier
|
||||
|
||||
for index in range(compatible_count):
|
||||
write_benchmark_plugin(
|
||||
plugins_dir=plugins_dir,
|
||||
plugin_name=f"compatible_{index:03d}",
|
||||
command_name=f"compatible_{index:03d}",
|
||||
requirement="shared-demo==1.0.0\n",
|
||||
)
|
||||
|
||||
for index in range(conflict_count):
|
||||
write_benchmark_plugin(
|
||||
plugins_dir=plugins_dir,
|
||||
plugin_name=f"conflict_{index:03d}",
|
||||
command_name=f"conflict_{index:03d}",
|
||||
requirement=f"shared-demo==2.0.{index}\n",
|
||||
)
|
||||
|
||||
return conflict_count, compatible_count
|
||||
|
||||
|
||||
def _snapshot_with_psutil(root_pid: int) -> ProcessTreeSnapshot:
|
||||
assert psutil is not None
|
||||
root = psutil.Process(root_pid)
|
||||
processes = [root] + root.children(recursive=True)
|
||||
total_rss = 0
|
||||
seen = 0
|
||||
for process in processes:
|
||||
try:
|
||||
total_rss += process.memory_info().rss
|
||||
seen += 1
|
||||
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
||||
continue
|
||||
return ProcessTreeSnapshot(
|
||||
collector="psutil",
|
||||
process_count=seen,
|
||||
total_rss_bytes=total_rss,
|
||||
total_rss_mb=round(total_rss / 1024 / 1024, 2),
|
||||
)
|
||||
|
||||
|
||||
def _snapshot_with_ps(root_pid: int) -> ProcessTreeSnapshot:
|
||||
result = subprocess.run(
|
||||
["ps", "-axo", "pid,ppid,rss"],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
children_by_parent: dict[int, list[int]] = {}
|
||||
rss_by_pid: dict[int, int] = {}
|
||||
for line in result.stdout.splitlines()[1:]:
|
||||
parts = line.strip().split(None, 2)
|
||||
if len(parts) != 3:
|
||||
continue
|
||||
pid, ppid, rss_kb = parts
|
||||
pid_int = int(pid)
|
||||
ppid_int = int(ppid)
|
||||
rss_by_pid[pid_int] = int(rss_kb) * 1024
|
||||
children_by_parent.setdefault(ppid_int, []).append(pid_int)
|
||||
|
||||
queue = [root_pid]
|
||||
seen: set[int] = set()
|
||||
total_rss = 0
|
||||
while queue:
|
||||
pid = queue.pop(0)
|
||||
if pid in seen:
|
||||
continue
|
||||
seen.add(pid)
|
||||
total_rss += rss_by_pid.get(pid, 0)
|
||||
queue.extend(children_by_parent.get(pid, []))
|
||||
|
||||
return ProcessTreeSnapshot(
|
||||
collector="ps",
|
||||
process_count=len(seen),
|
||||
total_rss_bytes=total_rss,
|
||||
total_rss_mb=round(total_rss / 1024 / 1024, 2),
|
||||
)
|
||||
|
||||
|
||||
def collect_process_tree_snapshot(root_pid: int) -> ProcessTreeSnapshot:
|
||||
if psutil is not None:
|
||||
try:
|
||||
return _snapshot_with_psutil(root_pid)
|
||||
except (PermissionError, psutil.Error):
|
||||
pass
|
||||
return _snapshot_with_ps(root_pid)
|
||||
|
||||
|
||||
async def sample_peak_rss(root_pid: int, stop_event: asyncio.Event) -> float:
|
||||
peak_bytes = 0
|
||||
while True:
|
||||
snapshot = await asyncio.to_thread(collect_process_tree_snapshot, root_pid)
|
||||
peak_bytes = max(peak_bytes, snapshot.total_rss_bytes)
|
||||
try:
|
||||
await asyncio.wait_for(stop_event.wait(), timeout=SAMPLE_INTERVAL_SECONDS)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
return round(peak_bytes / 1024 / 1024, 2)
|
||||
|
||||
|
||||
async def start_core_peer() -> tuple[Peer, Any]:
|
||||
left, right = make_transport_pair()
|
||||
core = Peer(
|
||||
transport=left,
|
||||
peer_info=PeerInfo(name="benchmark-core", role="core", version="v4"),
|
||||
)
|
||||
core.set_initialize_handler(
|
||||
lambda _message: asyncio.sleep(
|
||||
0,
|
||||
result=InitializeOutput(
|
||||
peer=PeerInfo(name="benchmark-core", role="core", version="v4"),
|
||||
capabilities=[],
|
||||
metadata={},
|
||||
),
|
||||
)
|
||||
)
|
||||
await core.start()
|
||||
return core, right
|
||||
|
||||
|
||||
async def run_case(
|
||||
case_root: Path,
|
||||
multiplier: int,
|
||||
*,
|
||||
conflict_base_count: int,
|
||||
compatible_base_count: int,
|
||||
) -> BenchmarkCaseResult:
|
||||
plugins_dir = case_root / "plugins"
|
||||
conflict_count, compatible_count = create_plugin_matrix(
|
||||
plugins_dir=plugins_dir,
|
||||
multiplier=multiplier,
|
||||
conflict_base_count=conflict_base_count,
|
||||
compatible_base_count=compatible_base_count,
|
||||
)
|
||||
env_manager = SyntheticGroupedEnvManager(case_root)
|
||||
core, supervisor_transport = await start_core_peer()
|
||||
runtime = SupervisorRuntime(
|
||||
transport=supervisor_transport,
|
||||
plugins_dir=plugins_dir,
|
||||
env_manager=env_manager,
|
||||
)
|
||||
|
||||
peak_stop_event = asyncio.Event()
|
||||
peak_task = asyncio.create_task(sample_peak_rss(os.getpid(), peak_stop_event))
|
||||
started_at = time.perf_counter()
|
||||
try:
|
||||
await runtime.start()
|
||||
await core.wait_until_remote_initialized()
|
||||
startup_duration_ms = round((time.perf_counter() - started_at) * 1000, 2)
|
||||
plan_result = env_manager._plan_result
|
||||
if plan_result is None:
|
||||
raise RuntimeError("benchmark plan result missing after runtime start")
|
||||
|
||||
steady_snapshot = await asyncio.to_thread(
|
||||
collect_process_tree_snapshot, os.getpid()
|
||||
)
|
||||
peak_rss_mb = max(
|
||||
steady_snapshot.total_rss_mb,
|
||||
await _finish_peak_sampler(peak_stop_event, peak_task),
|
||||
)
|
||||
expected_groups = conflict_count + (1 if compatible_count else 0)
|
||||
return BenchmarkCaseResult(
|
||||
multiplier=multiplier,
|
||||
conflict_plugins=conflict_count,
|
||||
compatible_plugins=compatible_count,
|
||||
total_plugins=conflict_count + compatible_count,
|
||||
group_count=len(plan_result.groups),
|
||||
skipped_plugins=len(runtime.skipped_plugins),
|
||||
startup_duration_ms=startup_duration_ms,
|
||||
steady_rss_mb=steady_snapshot.total_rss_mb,
|
||||
peak_rss_mb=peak_rss_mb,
|
||||
process_count=steady_snapshot.process_count,
|
||||
expected_groups=expected_groups,
|
||||
)
|
||||
finally:
|
||||
peak_stop_event.set()
|
||||
with contextlib.suppress(Exception):
|
||||
await peak_task
|
||||
await stop_runtime_concurrently(runtime, core)
|
||||
|
||||
|
||||
async def _finish_peak_sampler(
|
||||
stop_event: asyncio.Event, peak_task: asyncio.Task[float]
|
||||
) -> float:
|
||||
stop_event.set()
|
||||
return await peak_task
|
||||
|
||||
|
||||
async def stop_runtime_concurrently(runtime: SupervisorRuntime, core: Peer) -> None:
|
||||
session_stops = [
|
||||
session.stop() for session in list(runtime.worker_sessions.values())
|
||||
]
|
||||
if session_stops:
|
||||
await asyncio.gather(*session_stops, return_exceptions=True)
|
||||
await runtime.peer.stop()
|
||||
await core.stop()
|
||||
|
||||
|
||||
def render_table(results: list[BenchmarkCaseResult]) -> str:
|
||||
headers = [
|
||||
"倍数",
|
||||
"冲突",
|
||||
"兼容",
|
||||
"总插件",
|
||||
"分组",
|
||||
"预期分组",
|
||||
"启动(ms)",
|
||||
"稳态RSS(MB)",
|
||||
"峰值RSS(MB)",
|
||||
"进程数",
|
||||
]
|
||||
rows = [
|
||||
[
|
||||
str(item.multiplier),
|
||||
str(item.conflict_plugins),
|
||||
str(item.compatible_plugins),
|
||||
str(item.total_plugins),
|
||||
str(item.group_count),
|
||||
str(item.expected_groups),
|
||||
f"{item.startup_duration_ms:.2f}",
|
||||
f"{item.steady_rss_mb:.2f}",
|
||||
f"{item.peak_rss_mb:.2f}",
|
||||
str(item.process_count),
|
||||
]
|
||||
for item in results
|
||||
]
|
||||
widths = [
|
||||
max(len(headers[index]), *(len(row[index]) for row in rows))
|
||||
for index in range(len(headers))
|
||||
]
|
||||
lines = [
|
||||
" ".join(headers[index].ljust(widths[index]) for index in range(len(headers))),
|
||||
" ".join("-" * widths[index] for index in range(len(headers))),
|
||||
]
|
||||
lines.extend(
|
||||
" ".join(row[index].ljust(widths[index]) for index in range(len(headers)))
|
||||
for row in rows
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
async def run_benchmark(args: argparse.Namespace) -> list[BenchmarkCaseResult]:
|
||||
temp_dir_context: Any
|
||||
if args.keep_temp_dir:
|
||||
workspace_root = PROJECT_ROOT / ".tmp-benchmark-grouped-env"
|
||||
workspace_root.mkdir(parents=True, exist_ok=True)
|
||||
temp_dir_context = contextlib.nullcontext(str(workspace_root))
|
||||
else:
|
||||
temp_dir_context = tempfile.TemporaryDirectory(prefix="astrbot-grouped-bench-")
|
||||
|
||||
results: list[BenchmarkCaseResult] = []
|
||||
with temp_dir_context as temp_dir:
|
||||
workspace_root = Path(temp_dir)
|
||||
for multiplier in args.multipliers:
|
||||
case_root = workspace_root / f"case_{multiplier:02d}"
|
||||
if case_root.exists():
|
||||
shutil.rmtree(case_root)
|
||||
case_root.mkdir(parents=True, exist_ok=True)
|
||||
results.append(
|
||||
await run_case(
|
||||
case_root,
|
||||
multiplier,
|
||||
conflict_base_count=args.conflict_count,
|
||||
compatible_base_count=args.compatible_count,
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def write_json_report(
|
||||
*,
|
||||
output_path: Path,
|
||||
conflict_base_count: int,
|
||||
compatible_base_count: int,
|
||||
results: list[BenchmarkCaseResult],
|
||||
) -> None:
|
||||
payload = {
|
||||
"generated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"base_ratio": {
|
||||
"conflict_plugins": conflict_base_count,
|
||||
"compatible_plugins": compatible_base_count,
|
||||
},
|
||||
"multipliers": [item.multiplier for item in results],
|
||||
"results": [asdict(item) for item in results],
|
||||
}
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(
|
||||
json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
async def async_main() -> int:
|
||||
args = parse_args()
|
||||
if args.conflict_count < 0 or args.compatible_count < 0:
|
||||
raise SystemExit("conflict-count 和 compatible-count 不能为负数")
|
||||
if any(multiplier <= 0 for multiplier in args.multipliers):
|
||||
raise SystemExit("multipliers 必须全部大于 0")
|
||||
results = await run_benchmark(args)
|
||||
print(render_table(results))
|
||||
if args.output_json is not None:
|
||||
write_json_report(
|
||||
output_path=args.output_json,
|
||||
conflict_base_count=args.conflict_count,
|
||||
compatible_base_count=args.compatible_count,
|
||||
results=results,
|
||||
)
|
||||
print(f"\nJSON 报告已写入: {args.output_json}")
|
||||
|
||||
mismatched = [
|
||||
item
|
||||
for item in results
|
||||
if item.group_count != item.expected_groups or item.skipped_plugins != 0
|
||||
]
|
||||
return 1 if mismatched else 0
|
||||
|
||||
|
||||
def main() -> None:
|
||||
raise SystemExit(asyncio.run(async_main()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
339
tests_v4/test_grouped_environment_smoke.py
Normal file
339
tests_v4/test_grouped_environment_smoke.py
Normal file
@@ -0,0 +1,339 @@
|
||||
"""grouped env 的真实 smoke 测试。
|
||||
|
||||
运行示例:
|
||||
python -m pytest tests_v4/test_grouped_environment_smoke.py -m "slow and integration" -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import textwrap
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from astrbot_sdk.protocol.messages import InitializeOutput, PeerInfo
|
||||
from astrbot_sdk.runtime.bootstrap import SupervisorRuntime
|
||||
from astrbot_sdk.runtime.loader import PluginEnvironmentManager
|
||||
from astrbot_sdk.runtime.peer import Peer
|
||||
|
||||
from tests_v4.helpers import make_transport_pair
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.integration]
|
||||
|
||||
UV_BINARY = shutil.which("uv")
|
||||
|
||||
|
||||
async def start_test_core_peer(transport) -> Peer:
|
||||
"""Provide an initialize responder for supervisor startup."""
|
||||
core = Peer(
|
||||
transport=transport,
|
||||
peer_info=PeerInfo(name="grouped-env-core", role="core", version="v4"),
|
||||
)
|
||||
core.set_initialize_handler(
|
||||
lambda _message: asyncio.sleep(
|
||||
0,
|
||||
result=InitializeOutput(
|
||||
peer=PeerInfo(name="grouped-env-core", role="core", version="v4"),
|
||||
capabilities=[],
|
||||
metadata={},
|
||||
),
|
||||
)
|
||||
)
|
||||
await core.start()
|
||||
return core
|
||||
|
||||
|
||||
def _build_local_wheel(
|
||||
*,
|
||||
packages_root: Path,
|
||||
wheelhouse: Path,
|
||||
project_name: str,
|
||||
version: str,
|
||||
module_name: str,
|
||||
) -> Path:
|
||||
package_root = packages_root / f"{project_name}-{version}"
|
||||
source_dir = package_root / "src" / module_name
|
||||
source_dir.mkdir(parents=True, exist_ok=True)
|
||||
(source_dir / "__init__.py").write_text(
|
||||
f'__version__ = "{version}"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
(package_root / "pyproject.toml").write_text(
|
||||
textwrap.dedent(
|
||||
f"""\
|
||||
[build-system]
|
||||
requires = ["setuptools>=80", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "{project_name}"
|
||||
version = "{version}"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
"""
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
result = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"pip",
|
||||
"wheel",
|
||||
str(package_root),
|
||||
"--no-build-isolation",
|
||||
"--no-deps",
|
||||
"-w",
|
||||
str(wheelhouse),
|
||||
],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"failed to build local wheel {project_name}=={version}:\n"
|
||||
f"stdout:\n{result.stdout}\n\nstderr:\n{result.stderr}"
|
||||
)
|
||||
|
||||
candidates = sorted(
|
||||
wheelhouse.glob(f"{module_name}-{version}-*.whl"),
|
||||
key=lambda path: path.name,
|
||||
)
|
||||
if not candidates:
|
||||
raise RuntimeError(f"local wheel not found for {project_name}=={version}")
|
||||
return candidates[-1]
|
||||
|
||||
|
||||
def build_local_wheelhouse(root: Path) -> dict[str, Path]:
|
||||
"""Build offline wheels used by the smoke test."""
|
||||
packages_root = root / "packages"
|
||||
wheelhouse = root / "wheelhouse"
|
||||
packages_root.mkdir(parents=True, exist_ok=True)
|
||||
wheelhouse.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return {
|
||||
"alpha-1": _build_local_wheel(
|
||||
packages_root=packages_root,
|
||||
wheelhouse=wheelhouse,
|
||||
project_name="alpha-pkg",
|
||||
version="1.0.0",
|
||||
module_name="alpha_pkg",
|
||||
),
|
||||
"alpha-2": _build_local_wheel(
|
||||
packages_root=packages_root,
|
||||
wheelhouse=wheelhouse,
|
||||
project_name="alpha-pkg",
|
||||
version="2.0.0",
|
||||
module_name="alpha_pkg",
|
||||
),
|
||||
"beta-1": _build_local_wheel(
|
||||
packages_root=packages_root,
|
||||
wheelhouse=wheelhouse,
|
||||
project_name="beta-pkg",
|
||||
version="1.0.0",
|
||||
module_name="beta_pkg",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def write_smoke_plugin(
|
||||
*,
|
||||
plugins_dir: Path,
|
||||
plugin_name: str,
|
||||
command_name: str,
|
||||
requirement_line: str,
|
||||
import_module: str,
|
||||
expected_text: str,
|
||||
) -> Path:
|
||||
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(
|
||||
requirement_line + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
(plugin_dir / "plugin.yaml").write_text(
|
||||
yaml.dump(
|
||||
{
|
||||
"_schema_version": 2,
|
||||
"name": plugin_name,
|
||||
"display_name": plugin_name,
|
||||
"desc": "grouped env smoke plugin",
|
||||
"author": "codex",
|
||||
"version": "0.1.0",
|
||||
"runtime": {
|
||||
"python": f"{sys.version_info.major}.{sys.version_info.minor}"
|
||||
},
|
||||
"components": [
|
||||
{
|
||||
"class": "commands.main:SmokeCommand",
|
||||
"type": "command",
|
||||
"name": command_name,
|
||||
"description": command_name,
|
||||
}
|
||||
],
|
||||
},
|
||||
sort_keys=False,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(commands_dir / "main.py").write_text(
|
||||
textwrap.dedent(
|
||||
f"""\
|
||||
from {import_module} import __version__ as DEP_VERSION
|
||||
|
||||
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 SmokeCommand(CommandComponent):
|
||||
def __init__(self, context: Context):
|
||||
self.context = context
|
||||
|
||||
@filter.command("{command_name}")
|
||||
async def handle(self, event: AstrMessageEvent):
|
||||
yield event.plain_result("{expected_text} " + DEP_VERSION)
|
||||
"""
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
return plugin_dir
|
||||
|
||||
|
||||
async def invoke_command(
|
||||
runtime: SupervisorRuntime, core: Peer, command_name: str
|
||||
) -> str:
|
||||
"""Invoke one remote command and return the emitted text payload."""
|
||||
runtime.capability_router.sent_messages.clear()
|
||||
handler = next(
|
||||
(
|
||||
item
|
||||
for item in core.remote_handlers
|
||||
if getattr(item.trigger, "command", None) == command_name
|
||||
),
|
||||
None,
|
||||
)
|
||||
assert handler is not None, (
|
||||
f"command handler not found: {command_name}; "
|
||||
f"available={[getattr(item.trigger, 'command', None) for item in core.remote_handlers]}"
|
||||
)
|
||||
|
||||
await core.invoke(
|
||||
"handler.invoke",
|
||||
{
|
||||
"handler_id": handler.id,
|
||||
"event": {
|
||||
"text": command_name,
|
||||
"session_id": f"smoke-session-{command_name}",
|
||||
"user_id": "user-1",
|
||||
"platform": "test",
|
||||
},
|
||||
},
|
||||
request_id=f"grouped-env-smoke-{command_name}",
|
||||
)
|
||||
|
||||
sent_messages = list(runtime.capability_router.sent_messages)
|
||||
assert sent_messages, f"command {command_name} did not emit any message"
|
||||
return str(sent_messages[-1].get("text", ""))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
UV_BINARY is None, reason="uv is required for grouped env smoke tests"
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_grouped_environment_smoke_handles_shared_and_conflicting_dependencies():
|
||||
"""Real uv-backed smoke test for shared and conflicting plugin environments."""
|
||||
with tempfile.TemporaryDirectory(prefix="astrbot-grouped-env-smoke-") as temp_dir:
|
||||
root = Path(temp_dir)
|
||||
wheel_paths = build_local_wheelhouse(root)
|
||||
plugins_dir = root / "plugins"
|
||||
write_smoke_plugin(
|
||||
plugins_dir=plugins_dir,
|
||||
plugin_name="plugin_a",
|
||||
command_name="probe_alpha_v1",
|
||||
requirement_line=f"alpha-pkg @ {wheel_paths['alpha-1'].as_uri()}",
|
||||
import_module="alpha_pkg",
|
||||
expected_text="alpha-pkg",
|
||||
)
|
||||
write_smoke_plugin(
|
||||
plugins_dir=plugins_dir,
|
||||
plugin_name="plugin_b",
|
||||
command_name="probe_beta_v1",
|
||||
requirement_line=f"beta-pkg @ {wheel_paths['beta-1'].as_uri()}",
|
||||
import_module="beta_pkg",
|
||||
expected_text="beta-pkg",
|
||||
)
|
||||
write_smoke_plugin(
|
||||
plugins_dir=plugins_dir,
|
||||
plugin_name="plugin_c",
|
||||
command_name="probe_alpha_v2",
|
||||
requirement_line=f"alpha-pkg @ {wheel_paths['alpha-2'].as_uri()}",
|
||||
import_module="alpha_pkg",
|
||||
expected_text="alpha-pkg",
|
||||
)
|
||||
|
||||
env_manager = PluginEnvironmentManager(root)
|
||||
left, right = make_transport_pair()
|
||||
core = await start_test_core_peer(left)
|
||||
runtime = SupervisorRuntime(
|
||||
transport=right,
|
||||
plugins_dir=plugins_dir,
|
||||
env_manager=env_manager,
|
||||
)
|
||||
shared_venv_path = None
|
||||
isolated_venv_path = None
|
||||
|
||||
try:
|
||||
await runtime.start()
|
||||
await core.wait_until_remote_initialized()
|
||||
|
||||
assert sorted(runtime.loaded_plugins) == [
|
||||
"plugin_a",
|
||||
"plugin_b",
|
||||
"plugin_c",
|
||||
]
|
||||
assert runtime.skipped_plugins == {}
|
||||
assert env_manager._plan_result is not None
|
||||
assert len(env_manager._plan_result.groups) == 2
|
||||
|
||||
shared_group = env_manager._plan_result.plugin_to_group["plugin_a"]
|
||||
isolated_group = env_manager._plan_result.plugin_to_group["plugin_c"]
|
||||
assert (
|
||||
shared_group.id
|
||||
== env_manager._plan_result.plugin_to_group["plugin_b"].id
|
||||
)
|
||||
assert shared_group.id != isolated_group.id
|
||||
|
||||
shared_venv_path = shared_group.venv_path
|
||||
isolated_venv_path = isolated_group.venv_path
|
||||
assert shared_venv_path.exists()
|
||||
assert isolated_venv_path.exists()
|
||||
|
||||
alpha_v1_text = await invoke_command(runtime, core, "probe_alpha_v1")
|
||||
beta_v1_text = await invoke_command(runtime, core, "probe_beta_v1")
|
||||
alpha_v2_text = await invoke_command(runtime, core, "probe_alpha_v2")
|
||||
|
||||
assert "alpha-pkg 1.0.0" in alpha_v1_text
|
||||
assert "beta-pkg 1.0.0" in beta_v1_text
|
||||
assert "alpha-pkg 2.0.0" in alpha_v2_text
|
||||
assert alpha_v1_text != alpha_v2_text
|
||||
finally:
|
||||
await runtime.stop()
|
||||
await core.stop()
|
||||
env_manager._planner.cleanup_artifacts([])
|
||||
|
||||
assert shared_venv_path is not None
|
||||
assert isolated_venv_path is not None
|
||||
assert not shared_venv_path.exists()
|
||||
assert not isolated_venv_path.exists()
|
||||
@@ -725,6 +725,77 @@ class TestPluginEnvironmentManager:
|
||||
!= plan.plugin_to_group["plugin_two"].id
|
||||
)
|
||||
|
||||
def test_three_plugins_share_and_isolate_group_envs_then_cleanup(self):
|
||||
"""Two plugins should share one env while a conflicting third gets its own."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
manager = PluginEnvironmentManager(Path(temp_dir), uv_binary="/usr/bin/uv")
|
||||
plugins_dir = Path(temp_dir) / "plugins"
|
||||
spec_a = write_test_plugin(
|
||||
plugins_dir,
|
||||
"plugin_a",
|
||||
requirements="alpha==1.0.0\n",
|
||||
)
|
||||
spec_b = write_test_plugin(
|
||||
plugins_dir,
|
||||
"plugin_b",
|
||||
requirements="beta==1.0.0\n",
|
||||
)
|
||||
spec_c = write_test_plugin(
|
||||
plugins_dir,
|
||||
"plugin_c",
|
||||
requirements="alpha==2.0.0\n",
|
||||
)
|
||||
|
||||
def fake_compile(
|
||||
*, source_path: Path, output_path: Path, python_version: str
|
||||
):
|
||||
content = source_path.read_text(encoding="utf-8")
|
||||
if "alpha==1.0.0" in content and "alpha==2.0.0" in content:
|
||||
raise RuntimeError(
|
||||
"compile lockfile failed with exit code 1: conflict"
|
||||
)
|
||||
output_path.write_text(
|
||||
f"# python={python_version}\n{content}",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def fake_prepare(group: EnvironmentGroup) -> Path:
|
||||
group.python_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
group.python_path.write_text(
|
||||
f"group={group.id}\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
return group.python_path
|
||||
|
||||
manager._planner._compile_lockfile = fake_compile
|
||||
manager._group_manager.prepare = fake_prepare
|
||||
|
||||
plan = manager.plan([spec_a, spec_b, spec_c])
|
||||
shared_group = plan.plugin_to_group["plugin_a"]
|
||||
isolated_group = plan.plugin_to_group["plugin_c"]
|
||||
|
||||
assert len(plan.groups) == 2
|
||||
assert shared_group.id == plan.plugin_to_group["plugin_b"].id
|
||||
assert shared_group.id != isolated_group.id
|
||||
|
||||
path_a = manager.prepare_environment(spec_a)
|
||||
path_b = manager.prepare_environment(spec_b)
|
||||
path_c = manager.prepare_environment(spec_c)
|
||||
|
||||
assert path_a == path_b
|
||||
assert path_a != path_c
|
||||
assert len({path_a, path_b, path_c}) == 2
|
||||
assert shared_group.venv_path.exists()
|
||||
assert isolated_group.venv_path.exists()
|
||||
|
||||
manager._planner.cleanup_artifacts([])
|
||||
|
||||
assert not shared_group.venv_path.exists()
|
||||
assert not isolated_group.venv_path.exists()
|
||||
assert spec_a.plugin_dir.exists()
|
||||
assert spec_b.plugin_dir.exists()
|
||||
assert spec_c.plugin_dir.exists()
|
||||
|
||||
def test_plan_skips_only_plugin_with_invalid_lockfile(self):
|
||||
"""A plugin whose lockfile cannot be compiled should be skipped alone."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
|
||||
@@ -6,6 +6,8 @@ concurrency, and real-world scenarios.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
@@ -46,6 +48,30 @@ from astrbot_sdk.runtime.peer import Peer
|
||||
from tests_v4.helpers import FakeEnvManager, MemoryTransport, make_transport_pair
|
||||
|
||||
|
||||
def write_runtime_env_plugin(
|
||||
plugins_dir: Path,
|
||||
name: str,
|
||||
*,
|
||||
requirements: str = "",
|
||||
) -> 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": f"{sys.version_info.major}.{sys.version_info.minor}"
|
||||
},
|
||||
"components": [],
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(plugin_dir / "requirements.txt").write_text(requirements, encoding="utf-8")
|
||||
return plugin_dir
|
||||
|
||||
|
||||
async def start_test_core_peer(transport: MemoryTransport) -> Peer:
|
||||
"""Provide an initialize responder so transport-pair startup tests do not deadlock."""
|
||||
core = Peer(
|
||||
@@ -825,6 +851,105 @@ class TestSupervisorRuntimePluginLoading:
|
||||
await runtime.stop()
|
||||
await core.stop()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_loads_three_plugins_with_shared_and_isolated_group_envs(self):
|
||||
"""SupervisorRuntime should reuse one env for two plugins and isolate the third."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
plugins_dir = Path(temp_dir) / "plugins"
|
||||
write_runtime_env_plugin(
|
||||
plugins_dir,
|
||||
"plugin_a",
|
||||
requirements="alpha==1.0.0\n",
|
||||
)
|
||||
write_runtime_env_plugin(
|
||||
plugins_dir,
|
||||
"plugin_b",
|
||||
requirements="beta==1.0.0\n",
|
||||
)
|
||||
write_runtime_env_plugin(
|
||||
plugins_dir,
|
||||
"plugin_c",
|
||||
requirements="alpha==2.0.0\n",
|
||||
)
|
||||
|
||||
manager = PluginEnvironmentManager(Path(temp_dir), uv_binary="/usr/bin/uv")
|
||||
prepared_groups: list[str] = []
|
||||
|
||||
def fake_compile(
|
||||
*, source_path: Path, output_path: Path, python_version: str
|
||||
):
|
||||
content = source_path.read_text(encoding="utf-8")
|
||||
if "alpha==1.0.0" in content and "alpha==2.0.0" in content:
|
||||
raise RuntimeError(
|
||||
"compile lockfile failed with exit code 1: conflict"
|
||||
)
|
||||
output_path.write_text(
|
||||
f"# python={python_version}\n{content}",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def fake_prepare(group) -> Path:
|
||||
prepared_groups.append(group.id)
|
||||
group.python_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if group.python_path.exists():
|
||||
return group.python_path
|
||||
target = Path(sys.executable).resolve()
|
||||
if os.name == "nt":
|
||||
shutil.copy2(target, group.python_path)
|
||||
else:
|
||||
os.symlink(target, group.python_path)
|
||||
return group.python_path
|
||||
|
||||
manager._planner._compile_lockfile = fake_compile
|
||||
manager._group_manager.prepare = fake_prepare
|
||||
|
||||
left, right = make_transport_pair()
|
||||
core = await start_test_core_peer(left)
|
||||
runtime = SupervisorRuntime(
|
||||
transport=right,
|
||||
plugins_dir=plugins_dir,
|
||||
env_manager=manager,
|
||||
)
|
||||
shared_venv_path = None
|
||||
isolated_venv_path = None
|
||||
|
||||
try:
|
||||
await runtime.start()
|
||||
await core.wait_until_remote_initialized()
|
||||
|
||||
assert sorted(runtime.loaded_plugins) == [
|
||||
"plugin_a",
|
||||
"plugin_b",
|
||||
"plugin_c",
|
||||
]
|
||||
assert manager._plan_result is not None
|
||||
assert len(manager._plan_result.groups) == 2
|
||||
|
||||
shared_group = manager._plan_result.plugin_to_group["plugin_a"]
|
||||
isolated_group = manager._plan_result.plugin_to_group["plugin_c"]
|
||||
|
||||
assert (
|
||||
shared_group.id
|
||||
== manager._plan_result.plugin_to_group["plugin_b"].id
|
||||
)
|
||||
assert shared_group.id != isolated_group.id
|
||||
assert prepared_groups.count(shared_group.id) == 2
|
||||
assert prepared_groups.count(isolated_group.id) == 1
|
||||
|
||||
shared_venv_path = shared_group.venv_path
|
||||
isolated_venv_path = isolated_group.venv_path
|
||||
assert shared_venv_path.exists()
|
||||
assert isolated_venv_path.exists()
|
||||
finally:
|
||||
await runtime.stop()
|
||||
await core.stop()
|
||||
manager._planner.cleanup_artifacts([])
|
||||
|
||||
assert shared_venv_path is not None
|
||||
assert isolated_venv_path is not None
|
||||
assert not shared_venv_path.exists()
|
||||
assert not isolated_venv_path.exists()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_invalid_plugins(self):
|
||||
"""SupervisorRuntime 应该跳过无效插件并记录原因。"""
|
||||
|
||||
Reference in New Issue
Block a user