From 980efbc23f2aa7e08db02c5a4abe38241384bd25 Mon Sep 17 00:00:00 2001 From: united_pooh Date: Fri, 13 Mar 2026 19:44:03 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A1=A5=E5=85=85=E6=8F=92=E4=BB=B6=E5=88=86?= =?UTF-8?q?=E7=BB=84=E7=8E=AF=E5=A2=83=E6=B5=8B=E8=AF=95=E8=A6=86=E7=9B=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../benchmark_grouped_environment_stress.py | 561 ++++++++++++++++++ tests_v4/test_grouped_environment_smoke.py | 339 +++++++++++ tests_v4/test_loader.py | 71 +++ tests_v4/test_runtime_integration.py | 125 ++++ 4 files changed, 1096 insertions(+) create mode 100644 tests_v4/benchmark_grouped_environment_stress.py create mode 100644 tests_v4/test_grouped_environment_smoke.py diff --git a/tests_v4/benchmark_grouped_environment_stress.py b/tests_v4/benchmark_grouped_environment_stress.py new file mode 100644 index 000000000..e21de4087 --- /dev/null +++ b/tests_v4/benchmark_grouped_environment_stress.py @@ -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() diff --git a/tests_v4/test_grouped_environment_smoke.py b/tests_v4/test_grouped_environment_smoke.py new file mode 100644 index 000000000..7b1994e61 --- /dev/null +++ b/tests_v4/test_grouped_environment_smoke.py @@ -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() diff --git a/tests_v4/test_loader.py b/tests_v4/test_loader.py index 08420068a..d37366260 100644 --- a/tests_v4/test_loader.py +++ b/tests_v4/test_loader.py @@ -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: diff --git a/tests_v4/test_runtime_integration.py b/tests_v4/test_runtime_integration.py index 8c2b473e7..64ad09847 100644 --- a/tests_v4/test_runtime_integration.py +++ b/tests_v4/test_runtime_integration.py @@ -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 应该跳过无效插件并记录原因。"""