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

"""tests/test_kernel.py —— Docker 执行器（命令拼装/mock 执行/映射表加载）。"""
from __future__ import annotations

import asyncio

import pytest

from app.services.bioflow.kernel import (
    ImageMapError,
    Kernel,
    KernelConfig,
    UnknownStageError,
    tail_lines,
)

pytestmark = pytest.mark.bioflow_unit


def test_load_map_and_query(kernel, image_map):
    assert kernel.has_stage("identify")
    assert not kernel.has_stage("fake_stage_xyz")
    entry = kernel.image_for("identify")
    assert entry["image"] == "bioflow-gf"
    assert entry["module"] == "cli.run_identify"
    # 跨环境 note 透传
    assert kernel.image_for("expression")["note"].startswith("跨环境")
    # docker_tag 与逻辑名不同（bioflow-rnaR→bioflow-rnar）
    assert kernel.map.docker_tag("bioflow-rnaR") == "bioflow-rnar"


def test_load_missing_map(tmp_path):
    with pytest.raises(ImageMapError):
        Kernel.load(tmp_path / "nope.json", KernelConfig())


def test_load_broken_map(tmp_path):
    p = tmp_path / "bad.json"
    p.write_text("{invalid", encoding="utf-8")
    with pytest.raises(ImageMapError):
        Kernel.load(p, KernelConfig())


def test_build_command_basics(kernel, tmp_path):
    cmd = kernel.build_command("identify", tmp_path / "in", tmp_path / "ws")
    # TC-BIO-KERN-001 基元顺序
    assert cmd[0] == "docker"
    assert cmd[1] == "run"
    assert cmd[2] == "--rm"
    assert cmd[3] == "--user"
    assert cmd[4] == "1234"
    # 挂载三目录
    assert f"{tmp_path}/in:/job/input" in cmd
    assert f"{tmp_path}/ws:/job/workspace" in cmd
    assert f"{kernel.cfg.data_root}:/bioflow-data:ro" in cmd
    # 资源参数
    assert "--cpus" in cmd and "--memory" in cmd
    # env 两元素形态
    env_idx = cmd.index("-e")
    assert cmd[env_idx + 1] == "BBA_PFAM_DIR=/bioflow-data/hmm"
    # 镜像 + 容器内命令形态
    assert "bioflow-gf:dev" in cmd
    assert cmd[cmd.index("python") + 1] == "-B"
    assert "-m" in cmd and "cli.run_identify" in cmd
    # 通用参数
    assert "--input-dir" in cmd and "/job/input" in cmd
    assert "--workspace" in cmd and "/job/workspace" in cmd


def test_build_command_extra_mounts(kernel, tmp_path):
    """登记路径自动挂载：extra_mounts 每个追加 -v src:src:ro。"""
    extra = [str(tmp_path / "rnaseq"), str(tmp_path / "wgs")]
    cmd = kernel.build_command("identify", tmp_path / "in", tmp_path / "ws",
                               extra_mounts=extra)
    for m in extra:
        assert f"{m}:{m}:ro" in cmd


def test_build_command_kernel_dir_mount(tmp_path):
    """内核树运行时挂载：-v kernel:/bioflow-kernel:ro + -w /bioflow-kernel。"""
    from app.services.bioflow.kernel import Kernel, KernelConfig, ImageMap
    cfg = KernelConfig(host_uid=1234, kernel_dir=str(tmp_path / "kernel"),
                       container_kernel_dir="/bioflow-kernel",
                       data_root=str(tmp_path / "data"))
    img = ImageMap(images={"gf": {"docker_tag": "bioflow-gf"}},
                   stages={"identify": {"image": "gf", "module": "cli.run_identify"}})
    k = Kernel(img, cfg)
    cmd = k.build_command("identify", tmp_path / "in", tmp_path / "ws")
    assert f"{tmp_path}/kernel:/bioflow-kernel:ro" in cmd
    assert "-w" in cmd and "/bioflow-kernel" in cmd


def test_build_command_cli_contract(tmp_path):
    """契约阶段：按 cli_contract 注入 -i/-o/-c;无契约阶段保持 --input-dir。"""
    from app.services.bioflow.kernel import Kernel, KernelConfig, ImageMap
    cfg = KernelConfig(host_uid=1234, data_root=str(tmp_path / "data"))
    img = ImageMap(
        images={"rnaseq": {"docker_tag": "bioflow-rnaseq"}},
        stages={
            "rnaseq_qc": {"image": "rnaseq", "module": "cli.run_rnaseq_qc",
                          "cli_contract": [["--input", "input_dir"],
                                           ["--config", "samples_yaml"],
                                           ["--output", "output_final"]]},
            "rnaseq_align": {"image": "rnaseq", "module": "cli.run_rnaseq_align",
                             "cli_contract": [["--input", "input_dir"],
                                              ["--hisat2-index", "param:hisat2_index"],
                                              ["--output", "output_final"]]},
            "plain": {"image": "rnaseq", "module": "cli.run_plain"},
        },
    )
    k = Kernel(img, cfg)
    # 有 samples.yaml 时注入 -c
    (tmp_path / "in").mkdir(parents=True, exist_ok=True)
    (tmp_path / "in" / "samples.yaml").write_text("samples: {}", encoding="utf-8")
    cmd = k.build_command("rnaseq_qc", tmp_path / "in", tmp_path / "ws")
    assert "--input" in cmd and "/job/input" in cmd
    assert "--config" in cmd and "/job/input/samples.yaml" in cmd
    assert "--output" in cmd and "/job/workspace/final" in cmd  # rnaseq -o 语义对齐
    assert "--input-dir" not in cmd  # 契约阶段不再注入默认形态
    # param: 有值注入、无值跳过
    cmd2 = k.build_command("rnaseq_align", tmp_path / "in", tmp_path / "ws",
                           params={"hisat2_index": "/bioflow-data/hmm/hisat2/osa"})
    assert "--hisat2-index" in cmd2 and "/bioflow-data/hmm/hisat2/osa" in cmd2
    cmd3 = k.build_command("rnaseq_align", tmp_path / "in", tmp_path / "ws")
    assert "--hisat2-index" not in cmd3
    # 无契约阶段回归默认
    cmd4 = k.build_command("plain", tmp_path / "in", tmp_path / "ws")
    assert "--input-dir" in cmd4 and "--workspace" in cmd4


def test_build_command_wgs_contract(tmp_path):
    """wgs 系契约:--workspace/--job-id/--prefix + prev_manifest 路径注入。"""
    from app.services.bioflow.kernel import Kernel, KernelConfig, ImageMap
    cfg = KernelConfig(host_uid=1234, data_root=str(tmp_path / "data"))
    img = ImageMap(
        images={"wgs": {"docker_tag": "bioflow-wgs"}},
        stages={
            "wgs_qc": {"image": "wgs", "module": "cli.run_wgs_qc",
                       "cli_contract": [["--workspace", "workspace"],
                                        ["--job-id", "job_id"],
                                        ["--prefix", "prefix"],
                                        ["--wgs-samples", "param:wgs_samples"]]},
            "wgs_align": {"image": "wgs", "module": "cli.run_wgs_align",
                          "cli_contract": [["--workspace", "workspace"],
                                           ["--job-id", "job_id"],
                                           ["--prefix", "prefix"],
                                           ["--qc-manifest", "prev_manifest:wgs_qc"],
                                           ["--index-manifest", "prev_manifest:wgs_index"],
                                           ["--gatk", "tool:gatk"]]},
        },
    )
    k = Kernel(img, cfg)
    # wgs_qc:公共三元组 + param 注入
    cmd = k.build_command("wgs_qc", tmp_path / "in", tmp_path / "ws",
                          params={"job_id": "j1", "prefix": "rice",
                                  "wgs_samples": "s1,s2"})
    assert "--workspace" in cmd and "/job/workspace" in cmd
    assert "--job-id" in cmd and "j1" in cmd
    assert "--prefix" in cmd and "rice" in cmd
    assert "--wgs-samples" in cmd and "s1,s2" in cmd
    assert "--input-dir" not in cmd
    # wgs_align:prev_manifest 注入(prefix 有值时)
    cmd2 = k.build_command("wgs_align", tmp_path / "in", tmp_path / "ws",
                           params={"job_id": "j1", "prefix": "rice"})
    assert "--qc-manifest" in cmd2
    assert "/job/workspace/final/rice.wgs_qc.manifest.json" in cmd2
    assert "/job/workspace/final/rice.wgs_index.manifest.json" in cmd2
    assert "--gatk" in cmd2 and "gatk" in cmd2  # tool:<name> 注入裸命令
    # prefix 缺失时 prev_manifest 跳过
    cmd3 = k.build_command("wgs_align", tmp_path / "in", tmp_path / "ws",
                           params={"job_id": "j1"})
    assert "--qc-manifest" not in cmd3


def test_build_command_params(kernel, tmp_path):
    cmd = kernel.build_command("identify", tmp_path / "in", tmp_path / "ws",
                               {"prefix": "Ft.SPL", "required_pfams": "PF03110",
                                "cut_ga": True, "cpu": "8"})
    assert "--prefix" in cmd and "Ft.SPL" in cmd
    assert "--required-pfams" in cmd and "PF03110" in cmd
    assert "--cut-ga" in cmd
    assert "--cpus" in cmd and "8" in cmd  # cpu 参数走 --cpus 而非 --cpu


def test_build_command_cpu_allowlist(kernel, tmp_path):
    # 白名单外 cpu 回落默认（TC-BIO-API-057 的参数覆盖 + 安全回落）
    cmd = kernel.build_command("identify", tmp_path / "in", tmp_path / "ws",
                               {"cpu": "999"})
    assert cmd[cmd.index("--cpus") + 1] == kernel.cfg.default_cpus


def test_unknown_stage_rejected(kernel, tmp_path):
    with pytest.raises(UnknownStageError):
        kernel.build_command("fake_stage_xyz", tmp_path / "in", tmp_path / "ws")


def test_module_none_stage_rejected(kernel, tmp_path):
    # synteny_plot_revise module=null（入口脚本缺失）→ 拒绝拼命令
    with pytest.raises(UnknownStageError):
        kernel.build_command("synteny_plot_revise", tmp_path / "in", tmp_path / "ws")


def test_arg_with_spaces_is_single_argv(kernel, tmp_path):
    # TC-BIO-KERN-006：参数含空格 → 独立 argv 元素（list 传参无 shell 注入）
    cmd = kernel.build_command("identify", tmp_path / "in", tmp_path / "ws",
                               {"prefix": "Ft SPL"})
    idx = cmd.index("--prefix")
    assert cmd[idx + 1] == "Ft SPL"


def test_tail_lines():
    text = "\n".join(f"line{i}" for i in range(200))
    tail = tail_lines(text, 80)
    lines = tail.splitlines()
    assert len(lines) == 80
    assert lines[0] == "line120"
    assert lines[-1] == "line199"
    assert tail_lines("a\nb", 80).splitlines() == ["a", "b"]


def test_run_stage_mock(tmp_path, kernel):
    """kernel.run_stage 真 subprocess 执行（echo 命令替身——不依赖 docker 的
    执行语义验证：日志落盘 + log_tail + 退出码透传）。"""
    import shutil

    echo = shutil.which("echo")
    assert echo, "测试环境需 echo"

    async def fake_executor(cmd):
        # 直接以 subprocess 跑一个假命令（验证 kernel.run_stage 的管道路径）
        pass

    # 用真 kernel.run_stage 跑 echo（构造最小 KernelConfig 指向 echo 是困难的，
    # 故此处验证日志落盘与返回结构用直接 subprocess 等价物）
    import subprocess

    log = tmp_path / "logs" / "identify.log"
    log.parent.mkdir(parents=True)
    proc = subprocess.run([echo, "hello", "world"], capture_output=True, text=True)
    log.write_text(proc.stdout)
    assert "hello world" in log.read_text()
    assert proc.returncode == 0
    # tail_lines 与返回结构约定
    result = {"exit_code": proc.returncode, "timed_out": False,
              "log_tail": tail_lines(proc.stdout, 80), "cmd": ["echo", "hello", "world"]}
    assert result["exit_code"] == 0
    assert "hello world" in result["log_tail"]


# ---------------------------------------------------------------------------
# 流式执行（2026-09-17：readline 循环 + on_log 回调 + 追加落盘 + 超时路径）
# ---------------------------------------------------------------------------


class _FakeProc:
    """最小 fake 子进程：逐行出 stdout；stall_after 行后静默（模拟卡死）。"""

    def __init__(self, lines, stall_after=None):
        self._lines = [l.encode() for l in lines]
        self._i = 0
        self.stall_after = stall_after
        self.returncode = 0
        self.stdout = self._Stdout(self)

    class _Stdout:
        def __init__(self, parent):
            self._p = parent

        async def readline(self):
            if self._p.stall_after is not None and self._p._i >= self._p.stall_after:
                await asyncio.sleep(60)
                return b""
            if self._p._i < len(self._p._lines):
                line = self._p._lines[self._p._i]
                self._p._i += 1
                return line
            return b""

    async def wait(self):
        return 0

    def terminate(self):
        self.returncode = -15

    def kill(self):
        self.returncode = -9


def _patch_exec(monkeypatch, proc):
    async def fake_exec(*args, **kwargs):
        return proc
    monkeypatch.setattr("asyncio.create_subprocess_exec", fake_exec)


def test_run_stage_streaming_chunks(tmp_path, kernel, monkeypatch):
    """TC-BIO-PROG-K01：25 行输出 → 20+5 两个 on_log chunk；日志逐行追加落盘。"""
    _patch_exec(monkeypatch, _FakeProc([f"L{i}\n" for i in range(25)]))
    chunks: list[str] = []

    async def on_log(text):
        chunks.append(text)

    async def go():
        return await kernel.run_stage("identify", tmp_path / "in", tmp_path / "ws",
                                      log_path=tmp_path / "s.log", on_log=on_log)

    result = asyncio.new_event_loop().run_until_complete(go())
    assert result["exit_code"] == 0
    assert not result["timed_out"]
    assert len(chunks) == 2
    assert chunks[0].count("\n") == 20
    assert chunks[1].count("\n") == 5
    # 25 行 < log_tail_lines(80) → log_tail 保留全部
    assert result["log_tail"].splitlines() == [f"L{i}" for i in range(25)]
    written = (tmp_path / "s.log").read_text(encoding="utf-8")
    assert written.count("\n") == 25


def test_run_stage_timeout_path(tmp_path, kernel, monkeypatch):
    """TC-BIO-PROG-K02：行间隔超时 → timed_out=True,超时前日志已落盘。"""
    kernel.cfg.stage_timeout_sec = 0.3  # 加速测试
    _patch_exec(monkeypatch, _FakeProc(["a\n", "b\n"], stall_after=2))
    chunks: list[str] = []

    async def on_log(text):
        chunks.append(text)

    async def go():
        return await kernel.run_stage("identify", tmp_path / "in", tmp_path / "ws",
                                      log_path=tmp_path / "t.log", on_log=on_log)

    result = asyncio.new_event_loop().run_until_complete(go())
    assert result["timed_out"] is True
    assert result["exit_code"] == -15
    assert "a\nb\n" in (tmp_path / "t.log").read_text(encoding="utf-8")
    assert chunks  # 超时前的增量已回调