286 lines
10 KiB
Python
286 lines
10 KiB
Python
"""DockerExecutor 单元测试。
|
|
|
|
mock subprocess(`docker exec` 命令的实际跑由部署机 smoke 验,RUN.md 有 5 条命令)。
|
|
覆盖关键路径:
|
|
- 信任域 dispatch:host 工具直通 / container 工具走 docker exec
|
|
- argv 形态:--user / --workdir / setsid / bash -c / python <script>
|
|
- tmp .py:写到 host 侧 `.zcbot_tmp/<task_id>/`,执行完 unlink,无残留
|
|
- timeout / cancel:Popen.kill() 兜底
|
|
- schemas() / has_tool() 透传 host
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
from uuid import uuid4
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from core.executor import ExecCtx, ToolResult
|
|
from core.executor_docker import DockerExecutor, TMP_SUBDIR
|
|
from core.executor_host import HostExecutor
|
|
|
|
|
|
class FakePool:
|
|
"""SandboxPool 替身:ensure 返固定容器名,mark_active 记录调用。"""
|
|
|
|
def __init__(self):
|
|
self.ensure_calls = []
|
|
self.mark_active_calls = []
|
|
|
|
def ensure(self, user_id):
|
|
name = f"zcbot-sandbox-{user_id}"
|
|
self.ensure_calls.append(user_id)
|
|
return name
|
|
|
|
def mark_active(self, user_id):
|
|
self.mark_active_calls.append(user_id)
|
|
|
|
|
|
class FakeTool:
|
|
"""tools.base.Tool 替身:execute 返串,schema 暴露 name + 空 parameters。"""
|
|
|
|
def __init__(self, name, output="ok"):
|
|
self.name = name
|
|
self._output = output
|
|
self.execute_calls = []
|
|
|
|
@property
|
|
def schema(self):
|
|
return {"type": "function", "function": {"name": self.name}}
|
|
|
|
def execute(self, **kwargs):
|
|
self.execute_calls.append(kwargs)
|
|
return self._output
|
|
|
|
|
|
def make_executor(tools_dict=None):
|
|
"""构造 DockerExecutor + FakePool + tmp user_root。返回 (executor, pool, tmp_dir)。"""
|
|
tmp = tempfile.mkdtemp()
|
|
user_root = Path(tmp) / "users" / "u1"
|
|
user_root.mkdir(parents=True)
|
|
working_dir = user_root / "demo"
|
|
working_dir.mkdir()
|
|
|
|
if tools_dict is None:
|
|
tools_dict = {
|
|
"read": FakeTool("read", "READ_OUT"),
|
|
"shell": FakeTool("shell"), # host shell 不应被调用
|
|
"run_python": FakeTool("run_python"),
|
|
}
|
|
host = HostExecutor(tools_dict)
|
|
pool = FakePool()
|
|
executor = DockerExecutor(
|
|
host=host,
|
|
pool=pool,
|
|
user_id=uuid4(),
|
|
user_root=user_root,
|
|
working_dir=working_dir,
|
|
)
|
|
return executor, pool, Path(tmp)
|
|
|
|
|
|
def make_ctx(executor):
|
|
return ExecCtx(
|
|
user_id=executor.user_id,
|
|
task_id=uuid4(),
|
|
working_dir=executor.working_dir,
|
|
cancel_check=None,
|
|
)
|
|
|
|
|
|
class TestHostPassthrough(unittest.TestCase):
|
|
"""非 container tool 直通 host backend,不调 pool / subprocess。"""
|
|
|
|
def test_read_passthrough_to_host(self):
|
|
executor, pool, _ = make_executor()
|
|
ctx = make_ctx(executor)
|
|
result = executor.call_tool("read", {"file": "x"}, ctx)
|
|
self.assertEqual(result.content, "READ_OUT")
|
|
self.assertEqual(result.exit_code, 0)
|
|
self.assertEqual(pool.ensure_calls, [])
|
|
self.assertEqual(pool.mark_active_calls, [])
|
|
|
|
def test_schemas_and_has_tool_from_host(self):
|
|
executor, _, _ = make_executor()
|
|
names = [s["function"]["name"] for s in executor.schemas()]
|
|
self.assertIn("read", names)
|
|
self.assertIn("shell", names)
|
|
self.assertTrue(executor.has_tool("shell"))
|
|
self.assertFalse(executor.has_tool("nope"))
|
|
|
|
|
|
class TestShellExec(unittest.TestCase):
|
|
"""shell 调用走 docker exec subprocess,argv 形态正确。"""
|
|
|
|
def test_shell_invokes_docker_exec(self):
|
|
executor, pool, _ = make_executor()
|
|
ctx = make_ctx(executor)
|
|
|
|
proc = MagicMock()
|
|
proc.communicate.return_value = ("hello\n", "")
|
|
proc.returncode = 0
|
|
|
|
with patch("core.executor_docker.subprocess.Popen", return_value=proc) as popen:
|
|
result = executor.call_tool("shell", {"command": "echo hello"}, ctx)
|
|
|
|
self.assertIn("[stdout]\nhello", result.content)
|
|
self.assertIn("[exit 0]", result.content)
|
|
self.assertEqual(result.exit_code, 0)
|
|
|
|
argv = popen.call_args[0][0]
|
|
self.assertEqual(argv[:2], ["docker", "exec"])
|
|
self.assertIn("--user", argv)
|
|
self.assertIn("--workdir", argv)
|
|
# workdir 应是 /workspace/demo(working_dir 相对 user_root)
|
|
self.assertEqual(argv[argv.index("--workdir") + 1], "/workspace/demo")
|
|
# container name = zcbot-sandbox-<uid>
|
|
container_idx = argv.index(f"zcbot-sandbox-{executor.user_id}")
|
|
# setsid bash -c 必须出现且紧跟 container 之后
|
|
self.assertEqual(argv[container_idx + 1:], ["setsid", "bash", "-c", "echo hello"])
|
|
|
|
self.assertEqual(pool.ensure_calls, [executor.user_id])
|
|
self.assertEqual(pool.mark_active_calls, [executor.user_id])
|
|
|
|
def test_shell_bad_args(self):
|
|
executor, _, _ = make_executor()
|
|
ctx = make_ctx(executor)
|
|
result = executor.call_tool("shell", {"command": ""}, ctx)
|
|
self.assertIn("[Error]", result.content)
|
|
self.assertEqual(result.exit_code, 2)
|
|
|
|
def test_shell_timeout(self):
|
|
executor, pool, _ = make_executor()
|
|
ctx = make_ctx(executor)
|
|
|
|
import subprocess as real_subprocess
|
|
proc = MagicMock()
|
|
# 第一次 communicate 抛 TimeoutExpired,第二次(kill 后)返空
|
|
proc.communicate.side_effect = [
|
|
real_subprocess.TimeoutExpired(cmd="docker", timeout=0.5),
|
|
("", "killed\n"),
|
|
]
|
|
proc.returncode = -9
|
|
|
|
with patch("core.executor_docker.subprocess.Popen", return_value=proc), \
|
|
patch("core.executor_docker.time.monotonic", side_effect=[0, 100]):
|
|
result = executor.call_tool("shell", {"command": "sleep 9999", "timeout": 1}, ctx)
|
|
|
|
self.assertIn("timed out after 1s", result.content)
|
|
self.assertEqual(result.exit_code, 124)
|
|
proc.kill.assert_called_once()
|
|
|
|
def test_shell_cancel(self):
|
|
executor, _, _ = make_executor()
|
|
ctx = ExecCtx(
|
|
user_id=executor.user_id,
|
|
task_id=uuid4(),
|
|
working_dir=executor.working_dir,
|
|
cancel_check=lambda: True, # 立即 cancel
|
|
)
|
|
|
|
import subprocess as real_subprocess
|
|
proc = MagicMock()
|
|
proc.communicate.side_effect = [
|
|
real_subprocess.TimeoutExpired(cmd="docker", timeout=0.5),
|
|
("", ""),
|
|
]
|
|
proc.returncode = -15
|
|
|
|
with patch("core.executor_docker.subprocess.Popen", return_value=proc):
|
|
result = executor.call_tool("shell", {"command": "sleep 9999"}, ctx)
|
|
|
|
self.assertIn("cancelled by user", result.content)
|
|
self.assertEqual(result.exit_code, 130)
|
|
proc.kill.assert_called_once()
|
|
|
|
|
|
class TestRunPython(unittest.TestCase):
|
|
"""run_python:tmp .py 落 user_root/.zcbot_tmp/<task_id>/,跑完 unlink。"""
|
|
|
|
def test_run_python_tmp_script(self):
|
|
executor, pool, tmp_root = make_executor()
|
|
ctx = make_ctx(executor)
|
|
|
|
proc = MagicMock()
|
|
proc.communicate.return_value = ("42\n", "")
|
|
proc.returncode = 0
|
|
|
|
captured_argv = []
|
|
|
|
def _popen(argv, **kwargs):
|
|
captured_argv.append(argv)
|
|
return proc
|
|
|
|
with patch("core.executor_docker.subprocess.Popen", side_effect=_popen):
|
|
result = executor.call_tool(
|
|
"run_python", {"code": "print(42)"}, ctx
|
|
)
|
|
|
|
self.assertIn("[stdout]\n42", result.content)
|
|
self.assertEqual(result.exit_code, 0)
|
|
|
|
argv = captured_argv[0]
|
|
# 末尾形态:setsid python /workspace/.zcbot_tmp/<task_id>/<rand>.py
|
|
self.assertEqual(argv[-3], "setsid")
|
|
self.assertEqual(argv[-2], "python")
|
|
self.assertTrue(argv[-1].startswith(f"/workspace/{TMP_SUBDIR}/{ctx.task_id}/"))
|
|
self.assertTrue(argv[-1].endswith(".py"))
|
|
# PYTHONIOENCODING / PYTHONPATH 注入
|
|
env_kvs = [argv[i + 1] for i, a in enumerate(argv) if a == "-e"]
|
|
self.assertIn("PYTHONIOENCODING=utf-8", env_kvs)
|
|
self.assertIn("PYTHONPATH=/workspace", env_kvs)
|
|
|
|
# host 侧 tmp 已 unlink(目录可能仍在,无所谓 —— ensure 容器时会重新 mkdir)
|
|
tmp_subroot = executor.user_root / TMP_SUBDIR / str(ctx.task_id)
|
|
leftover = list(tmp_subroot.glob("*.py")) if tmp_subroot.exists() else []
|
|
self.assertEqual(leftover, [], f"tmp .py not cleaned up: {leftover}")
|
|
|
|
def test_run_python_bad_code_type(self):
|
|
executor, _, _ = make_executor()
|
|
ctx = make_ctx(executor)
|
|
result = executor.call_tool("run_python", {"code": 123}, ctx)
|
|
self.assertIn("[Error]", result.content)
|
|
self.assertEqual(result.exit_code, 2)
|
|
|
|
def test_run_python_cleans_tmp_on_exception(self):
|
|
"""Popen 抛异常时 tmp .py 仍要被清理(finally 兜底)。"""
|
|
executor, _, _ = make_executor()
|
|
ctx = make_ctx(executor)
|
|
|
|
with patch(
|
|
"core.executor_docker.subprocess.Popen",
|
|
side_effect=RuntimeError("boom"),
|
|
):
|
|
result = executor.call_tool("run_python", {"code": "x"}, ctx)
|
|
|
|
self.assertIn("[Error executing run_python via docker]", result.content)
|
|
self.assertEqual(result.exit_code, 1)
|
|
tmp_subroot = executor.user_root / TMP_SUBDIR / str(ctx.task_id)
|
|
leftover = list(tmp_subroot.glob("*.py")) if tmp_subroot.exists() else []
|
|
self.assertEqual(leftover, [])
|
|
|
|
|
|
class TestUnknownTool(unittest.TestCase):
|
|
def test_unknown_tool_goes_to_host(self):
|
|
executor, _, _ = make_executor(tools_dict={}) # 空 host → 啥都没
|
|
ctx = make_ctx(executor)
|
|
result = executor.call_tool("nope", {}, ctx)
|
|
self.assertIn("unknown tool", result.content)
|
|
self.assertEqual(result.exit_code, 2)
|
|
|
|
def test_container_tool_not_registered_on_host(self):
|
|
"""caps.enable_run_python=False:host 没装 run_python,docker 也应拒。"""
|
|
executor, _, _ = make_executor(tools_dict={"read": FakeTool("read")})
|
|
ctx = make_ctx(executor)
|
|
result = executor.call_tool("run_python", {"code": "x"}, ctx)
|
|
self.assertIn("unknown tool", result.content)
|
|
self.assertEqual(result.exit_code, 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|