219 lines
9.3 KiB
Python
219 lines
9.3 KiB
Python
 
|
"""run_extension 容器入口单测 —— 纯 stdlib unittest(容器内无 pytest)。
|
||
|
|
|
||
|
|
覆盖:模板渲染白名单/输出契约(退出码 3)/参数注入/退出码透传/超时/extension.json
|
||
|
|
异常。真 subprocess 执行(容器入口语义,宿主机 python3 直跑,不需要 docker)。
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
import textwrap
|
||
|
|
import unittest
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
ENTRY = REPO_ROOT / "bioflow_kernel" / "cli" / "run_extension.py"
|
||
|
|
|
||
|
|
if not ENTRY.is_file():
|
||
|
|
raise SystemExit(f"入口脚本缺失: {ENTRY}")
|
||
|
|
|
||
|
|
|
||
|
|
def _write_extension(ext_dir: Path, template: str, **extra) -> None:
|
||
|
|
payload = {
|
||
|
|
"schema_version": "extension_stage_v1",
|
||
|
|
"ext_id": "ext_test",
|
||
|
|
"display_name": "测试扩展",
|
||
|
|
"status": "published",
|
||
|
|
"command_template": template,
|
||
|
|
"original_command": "",
|
||
|
|
"runtime_image": "bioflow-gf",
|
||
|
|
"requirements": {"pip": [], "apt": []},
|
||
|
|
"parameters": [],
|
||
|
|
"input_roles": {"input_file": "fastq"},
|
||
|
|
"nl_triggers": [],
|
||
|
|
"analysis_type": "table_filter",
|
||
|
|
"golden": None,
|
||
|
|
"repair_count": 0,
|
||
|
|
}
|
||
|
|
payload.update(extra)
|
||
|
|
(ext_dir / "extension.json").write_text(json.dumps(payload), encoding="utf-8")
|
||
|
|
|
||
|
|
|
||
|
|
def _run_entry(ext_dir: Path, input_dir: Path, output_dir: Path,
|
||
|
|
params: list[str] | None = None, timeout: float = 120.0,
|
||
|
|
extra_args: list[str] | None = None) -> subprocess.CompletedProcess:
|
||
|
|
cmd = [sys.executable, "-B", str(ENTRY),
|
||
|
|
"--ext-dir", str(ext_dir),
|
||
|
|
"--input-dir", str(input_dir),
|
||
|
|
"--output-dir", str(output_dir)]
|
||
|
|
cmd += extra_args or []
|
||
|
|
for item in params or []:
|
||
|
|
cmd += ["--param", item]
|
||
|
|
# 测试宿主无裸 python(生产镜像内由 /opt/conda/bin 提供):注入 shim 验证执行语义
|
||
|
|
import tempfile
|
||
|
|
|
||
|
|
shim_dir = tempfile.mkdtemp(prefix="ext-shim-")
|
||
|
|
shim = Path(shim_dir) / "python"
|
||
|
|
shim.write_text(f"#!/bin/sh\nexec {sys.executable} \"$@\"\n", encoding="utf-8")
|
||
|
|
shim.chmod(0o755)
|
||
|
|
env = dict(__import__("os").environ)
|
||
|
|
env["PATH"] = shim_dir + ":" + env.get("PATH", "")
|
||
|
|
return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=env)
|
||
|
|
|
||
|
|
|
||
|
|
class RenderTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
import importlib.util
|
||
|
|
|
||
|
|
spec = importlib.util.spec_from_file_location("run_extension", ENTRY)
|
||
|
|
self.mod = importlib.util.module_from_spec(spec)
|
||
|
|
spec.loader.exec_module(self.mod)
|
||
|
|
|
||
|
|
def _ctx(self, template: str, **extra):
|
||
|
|
ext = {"command_template": template, "input_roles": {"input_file": "fastq"}}
|
||
|
|
ext.update(extra)
|
||
|
|
return self.mod.build_render_context(
|
||
|
|
ext,
|
||
|
|
run_dir=Path("/job/output"),
|
||
|
|
input_dir=Path("/job/input"),
|
||
|
|
input_files=[Path("/job/input/sample.fastq")],
|
||
|
|
output_dir=Path("/job/output"),
|
||
|
|
code_dir=Path("/job/ext/converted_code"),
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_basic_context(self):
|
||
|
|
ctx = self._ctx("python {code_dir}/x.py --fastq {input_file} --out {output_dir}")
|
||
|
|
self.assertEqual(ctx["code_dir"], "/job/ext/converted_code")
|
||
|
|
self.assertEqual(ctx["input_file"], "/job/input/sample.fastq")
|
||
|
|
self.assertEqual(ctx["output_dir"], "/job/output")
|
||
|
|
self.assertEqual(ctx["work_dir"], "/job/output")
|
||
|
|
|
||
|
|
def test_input_hmm_suffix(self):
|
||
|
|
ctx = self._ctx("hmmsearch {input_hmm} > {output_dir}/out.txt")
|
||
|
|
# 无 .hmm 输入时回退 input_file(与 executor.py 同语义)
|
||
|
|
self.assertEqual(ctx["input_hmm"], "/job/input/sample.fastq")
|
||
|
|
|
||
|
|
def test_template_render_unknown_placeholder_rejected(self):
|
||
|
|
mod = self.mod
|
||
|
|
with self.assertRaises(mod.ExtensionRenderError):
|
||
|
|
mod.render_command_template(
|
||
|
|
{"command_template": "run {not_a_placeholder}"},
|
||
|
|
template="run {not_a_placeholder}",
|
||
|
|
run_dir=Path("/job/output"),
|
||
|
|
input_dir=Path("/job/input"),
|
||
|
|
input_files=[Path("/job/input/a.fastq")],
|
||
|
|
output_dir=Path("/job/output"),
|
||
|
|
code_dir=Path("/job/ext/converted_code"),
|
||
|
|
)
|
||
|
|
|
||
|
|
def test_validate_template_output_dir_required(self):
|
||
|
|
self.assertIn("output_dir", self.mod._validate_template("python {code_dir}/x.py") or "")
|
||
|
|
self.assertIsNone(self.mod._validate_template("python {code_dir}/x.py > {output_dir}/r.txt"))
|
||
|
|
|
||
|
|
def test_validate_template_whitelist(self):
|
||
|
|
err = self.mod._validate_template("run {unknown_placeholder} > {output_dir}/x")
|
||
|
|
self.assertIn("unknown_placeholder", err or "")
|
||
|
|
self.assertIsNone(self.mod._validate_template("x {param_alpha} {param_beta} {work_dir} {output_dir}"))
|
||
|
|
|
||
|
|
def test_shell_parameter_value_quoting(self):
|
||
|
|
self.assertEqual(self.mod._shell_parameter_value("a-b_c/d.e"), "a-b_c/d.e")
|
||
|
|
self.assertEqual(self.mod._shell_parameter_value("a b"), "'a b'")
|
||
|
|
|
||
|
|
def test_safe_parameter_name(self):
|
||
|
|
self.assertEqual(self.mod._safe_parameter_name("my-param"), "my_param")
|
||
|
|
self.assertEqual(self.mod._safe_parameter_name("2fast"), "p_2fast")
|
||
|
|
|
||
|
|
|
||
|
|
class EntryTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
import tempfile
|
||
|
|
|
||
|
|
self._tmp = tempfile.TemporaryDirectory()
|
||
|
|
self.root = Path(self._tmp.name)
|
||
|
|
self.ext = self.root / "ext"
|
||
|
|
self.code = self.ext / "converted_code"
|
||
|
|
self.input_dir = self.root / "input"
|
||
|
|
self.output = self.root / "output"
|
||
|
|
for d in (self.code, self.input_dir):
|
||
|
|
d.mkdir(parents=True)
|
||
|
|
self.script = self.code / "hello.py"
|
||
|
|
self.script.write_text(
|
||
|
|
textwrap.dedent("""\
|
||
|
|
import pathlib, sys
|
||
|
|
print("hello from ext", flush=True)
|
||
|
|
pathlib.Path(sys.argv[1]).mkdir(parents=True, exist_ok=True)
|
||
|
|
pathlib.Path(sys.argv[1], "result.tsv").write_text("a\\tb\\n")
|
||
|
|
"""), encoding="utf-8")
|
||
|
|
(self.input_dir / "sample.fastq").write_text(">s1\\nACGT\\n", encoding="utf-8")
|
||
|
|
|
||
|
|
def tearDown(self):
|
||
|
|
self._tmp.cleanup()
|
||
|
|
|
||
|
|
def _template(self) -> str:
|
||
|
|
return ("python {code_dir}/hello.py {output_dir} --fastq {input_file} "
|
||
|
|
"--min {param_min} --label {param_label}")
|
||
|
|
|
||
|
|
def test_success_and_output_contract(self):
|
||
|
|
_write_extension(self.ext, self._template())
|
||
|
|
result = _run_entry(self.ext, self.input_dir, self.output,
|
||
|
|
params=["min=30", "label=run_1"])
|
||
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
||
|
|
self.assertIn("hello from ext", result.stdout)
|
||
|
|
self.assertTrue((self.output / "result.tsv").is_file())
|
||
|
|
|
||
|
|
def test_output_contract_empty_exit_3(self):
|
||
|
|
(self.script).write_text("print('no output written')\n", encoding="utf-8")
|
||
|
|
_write_extension(self.ext, "python {code_dir}/hello.py --out {output_dir}")
|
||
|
|
result = _run_entry(self.ext, self.input_dir, self.output)
|
||
|
|
self.assertEqual(result.returncode, 3)
|
||
|
|
self.assertIn("输出契约", result.stderr)
|
||
|
|
|
||
|
|
def test_missing_extension_json_exit_2(self):
|
||
|
|
result = _run_entry(self.ext, self.input_dir, self.output)
|
||
|
|
self.assertEqual(result.returncode, 2)
|
||
|
|
|
||
|
|
def test_invalid_template_exit_2(self):
|
||
|
|
_write_extension(self.ext, "python {code_dir}/hello.py --unknown {nope} --out {output_dir}")
|
||
|
|
result = _run_entry(self.ext, self.input_dir, self.output)
|
||
|
|
self.assertEqual(result.returncode, 2)
|
||
|
|
self.assertIn("白名单外", result.stderr)
|
||
|
|
|
||
|
|
def test_exit_code_passthrough(self):
|
||
|
|
(self.code / "fail.sh").write_text("#!/bin/sh\nexit 7\n", encoding="utf-8")
|
||
|
|
(self.code / "fail.sh").chmod(0o755)
|
||
|
|
_write_extension(self.ext, "sh {code_dir}/fail.sh > {output_dir}/o.txt")
|
||
|
|
result = _run_entry(self.ext, self.input_dir, self.output)
|
||
|
|
self.assertEqual(result.returncode, 7)
|
||
|
|
self.assertIn("code 7", result.stderr)
|
||
|
|
|
||
|
|
def test_template_missing_param_exit_2(self):
|
||
|
|
_write_extension(self.ext, "python {code_dir}/hello.py --x {param_missing} --out {output_dir}")
|
||
|
|
result = _run_entry(self.ext, self.input_dir, self.output)
|
||
|
|
self.assertEqual(result.returncode, 2)
|
||
|
|
self.assertIn("模板渲染失败", result.stderr)
|
||
|
|
|
||
|
|
def test_param_with_spaces_single_arg(self):
|
||
|
|
# 参数含空格 → shlex.quote 保证独立 argv(等价 kernel list-form 无 shell 注入)
|
||
|
|
(self.code / "args.py").write_text(
|
||
|
|
"import sys; print(repr(sys.argv[1]))\n", encoding="utf-8")
|
||
|
|
_write_extension(self.ext, "python {code_dir}/args.py {param_label} --out {output_dir}")
|
||
|
|
result = _run_entry(self.ext, self.input_dir, self.output,
|
||
|
|
params=["label=hello world"])
|
||
|
|
self.assertEqual(result.returncode, 3, result.stderr) # 未写输出 → 契约校验
|
||
|
|
self.assertIn("'hello world'", result.stdout)
|
||
|
|
|
||
|
|
def test_timeout_returns_124(self):
|
||
|
|
(self.code / "sleep.py").write_text(
|
||
|
|
"import time\ntime.sleep(600)\n", encoding="utf-8")
|
||
|
|
_write_extension(self.ext, "python {code_dir}/sleep.py --out {output_dir}")
|
||
|
|
result = _run_entry(self.ext, self.input_dir, self.output,
|
||
|
|
timeout=10.0, extra_args=["--timeout", "2"])
|
||
|
|
self.assertEqual(result.returncode, 124)
|
||
|
|
self.assertIn("超时", result.stderr)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main(verbosity=2)
|