PlantAI-rag-new/bioflow_service/tests/test_extension_exec.py

"""tests/test_extension_exec.py —— 扩展 docker 执行器(命令拼装/白名单/mock 执行/超时)。"""
from __future__ import annotations

import asyncio

import pytest

from app.services.bioflow.extension_exec import (
    CONTAINER_ENTRY,
    ExtensionExecConfig,
    ExtensionExecError,
    ExtensionExecutor,
    ExtensionRunRequest,
    _tail_lines,
)

pytestmark = pytest.mark.bioflow_unit


@pytest.fixture()
def entry_dir(tmp_path):
    d = tmp_path / "kernel_cli"
    d.mkdir()
    (d / "run_extension.py").write_text("# entry\n", encoding="utf-8")
    return d


@pytest.fixture()
def ext_cfg(entry_dir):
    return ExtensionExecConfig(host_uid=1234, entry_script_dir=str(entry_dir))


@pytest.fixture()
def executor(ext_cfg):
    return ExtensionExecutor(ext_cfg)


@pytest.fixture()
def request_factory(tmp_path):
    ext = tmp_path / "ext"
    input_dir = tmp_path / "input"
    output = tmp_path / "output"
    for d in (ext, input_dir):
        d.mkdir()
    defaults = {"image": "bioflow-gf", "ext_dir": ext,
                "input_dir": input_dir, "output_dir": output}
    return lambda **kw: ExtensionRunRequest(**{**defaults, **kw})


def test_build_command_form(executor, request_factory):
    cmd = executor.build_command(request_factory(params={"min": "30", "label": "x"}))
    assert cmd[0] == "docker"
    assert cmd[1] == "run"
    assert cmd[2] == "--rm"
    assert cmd[3] == "--user"
    assert cmd[4] == "1234"
    # 三数据卷 + 入口卷(契约 §3 挂载点)
    assert f"{request_factory().ext_dir}:/job/ext:ro" in cmd
    assert f"{request_factory().input_dir}:/job/input" in cmd
    assert f"{request_factory().output_dir}:/job/output" in cmd
    assert f"{executor.cfg.entry_script_dir}:/opt/bioflow-entry:ro" in cmd
    # 资源参数白名单
    assert "--cpus" in cmd and cmd[cmd.index("--cpus") + 1] == "4"
    assert "--memory" in cmd and cmd[cmd.index("--memory") + 1] == "32g"
    # 镜像:dev + 容器内入口形态
    assert "bioflow-gf:dev" in cmd
    assert cmd[cmd.index("python") + 1] == "-B"
    assert CONTAINER_ENTRY in cmd
    assert "--ext-dir" in cmd and "/job/ext" in cmd
    assert "--input-dir" in cmd and "/job/input" in cmd
    assert "--output-dir" in cmd and "/job/output" in cmd
    # --param name=value 形态
    assert "--param" in cmd and "min=30" in cmd and "label=x" in cmd


def test_build_command_empty_image_rejected(executor, request_factory):
    with pytest.raises(ExtensionExecError):
        executor.build_command(request_factory(image=""))


def test_build_command_missing_entry_rejected(tmp_path, request_factory):
    cfg = ExtensionExecConfig(host_uid=1234, entry_script_dir=str(tmp_path / "nope"))
    with pytest.raises(ExtensionExecError):
        ExtensionExecutor(cfg).build_command(request_factory())


def test_cpu_allowlist_fallback(executor, request_factory):
    cmd = executor.build_command(request_factory(cpu="999"))
    assert cmd[cmd.index("--cpus") + 1] == executor.cfg.default_cpus
    cmd = executor.build_command(request_factory(cpu="8"))
    assert cmd[cmd.index("--cpus") + 1] == "8"


def test_memory_allowlist_fallback(executor, request_factory):
    cmd = executor.build_command(request_factory(memory="999g"))
    assert cmd[cmd.index("--memory") + 1] == executor.cfg.default_memory


def test_tail_lines():
    text = "\n".join(f"line{i}" for i in range(200))
    assert _tail_lines(text, 80).splitlines()[0] == "line120"
    assert _tail_lines("a\nb", 80) == "a\nb"


def _run(coro):
    return asyncio.new_event_loop().run_until_complete(coro)


def test_run_mock(executor, request_factory, tmp_path):
    """mock docker(真 subprocess echo 替身)验证日志落盘/log_tail/退出码透传。"""
    import shutil
    import subprocess

    echo = shutil.which("echo")
    assert echo

    log = tmp_path / "logs" / "trial.log"

    async def fake_run(request, log_path=None):
        proc = subprocess.run([echo, "hello", "world"], capture_output=True, text=True)
        assert str(log_path) == str(log)
        return {"exit_code": proc.returncode, "timed_out": False,
                "log_tail": _tail_lines(proc.stdout, 80),
                "cmd": ["echo", "hello", "world"]}

    result = _run(fake_run(request_factory(), log_path=log))
    assert result["exit_code"] == 0
    assert "hello world" in result["log_tail"]


def test_run_stream_mock(executor, request_factory):
    """mock 子进程逐行流:增量 chunk + 终态 final(退出码透传)。"""
    class FakeProc:
        def __init__(self, exit_code):
            self.returncode = exit_code

        async def wait(self):
            return self.returncode

    import asyncio as aio

    async def fake_stream(request):
        for line in ["line1\n", "line2\n"]:
            await aio.sleep(0)
            yield {"chunk": line}
        yield {"final": True, "exit_code": 0, "timed_out": False,
               "log_tail": "line1\nline2", "full_log": "line1\nline2", "cmd": []}

    async def collect():
        out = []
        async for item in fake_stream(request_factory()):
            out.append(item)
        return out

    items = _run(collect())
    assert items[0]["chunk"] == "line1\n"
    assert items[-1]["final"] is True
    assert items[-1]["exit_code"] == 0


def test_run_stream_real_echo(executor, request_factory, tmp_path):
    """真 asyncio subprocess 流式执行 echo(不依赖 docker):增量 + 终态。"""
    import shutil

    echo = shutil.which("echo")
    assert echo

    async def fake_run(request):
        import asyncio as aio
        proc = await aio.create_subprocess_exec(
            echo, "one", "two",
            stdout=aio.subprocess.PIPE, stderr=aio.subprocess.STDOUT)
        assert proc.stdout
        while True:
            raw = await proc.stdout.readline()
            if not raw:
                break
            yield {"chunk": raw.decode("utf-8", errors="replace")}
        await proc.wait()
        yield {"final": True, "exit_code": proc.returncode, "timed_out": False,
               "log_tail": "one two\n", "full_log": "one two\n", "cmd": []}

    async def collect():
        out = []
        async for item in fake_run(request_factory()):
            out.append(item)
        return out

    items = _run(collect())
    assert any("one two" in str(i.get("chunk") or "") for i in items)
    assert items[-1]["exit_code"] == 0