补充插件分组环境测试覆盖

This commit is contained in:
united_pooh
2026-03-13 19:44:03 +08:00
parent 885e4fe32e
commit 980efbc23f
4 changed files with 1096 additions and 0 deletions

View 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()

View 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()

View File

@@ -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:

View File

@@ -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 应该跳过无效插件并记录原因。"""