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

"""tests/test_extensions.py —— 扩展生命周期状态机与存储（契约 §1/§2）。

覆盖：创建布局（extension.json/raw_code/raw_test_data/converted_code）/
状态机合法流转 / 非法流转 409 语义 / build 规则版回退（无 LLM key）/
trial 执行（fake executor，成/败）/ validate（无金标+有金标）/
repair（轮次上限/LLM 不可用规则回退）/
confirm（nl_triggers 编辑/parameters 补齐）/ publish / run（仅 published、
run_result 落盘、输出登记）/ 列表 published 优先 / 删除守卫。
"""
from __future__ import annotations

import json
from pathlib import Path

import pytest

from app.services.bioflow import extension_spec as spec
from app.services.bioflow import extensions as ext

pytestmark = pytest.mark.bioflow_unit

_CODE = "#!/usr/bin/env python\nimport sys\nprint('hello')\n"


@pytest.fixture()
def store(tmp_path, monkeypatch):
    monkeypatch.delenv("DOCGEN_LLM_API_KEY", raising=False)
    monkeypatch.delenv("BIOFLOW_EXTENSION_API_KEY", raising=False)
    return ext.ExtensionStore(base_dir=str(tmp_path / "data" / "extensions"))


def _create(store, *, name="测序质量统计", command="python quality_stats.py --fastq input.fastq --out results/",
            code=None, data=None, golden=None):
    return store.create_extension(
        display_name=name,
        command=command,
        runtime_image="test4py311",  # 旧 conda 环境名 → 映射 bioflow-gf
        code_files=code or [("quality_stats.py", _CODE.encode())],
        data_files=data or [],
        golden=golden,
    )


def _fake_executor(exit_code=0, log="run ok\n", outputs=("result.txt",)):
    """sync 可调用 executor：写输出 + 返回 ExtensionRunResult（兼容 _execute 三类形态）。"""

    class FakeExec:
        def __init__(self):
            self.calls = []

        def execute(self, spec_in):
            self.calls.append(spec_in)
            out_dir = Path(spec_in.output_dir)
            out_dir.mkdir(parents=True, exist_ok=True)
            for name in outputs:
                (out_dir / name).write_text("data", encoding="utf-8")
            if spec_in.log_path:
                Path(spec_in.log_path).parent.mkdir(parents=True, exist_ok=True)
                Path(spec_in.log_path).write_text(log, encoding="utf-8")
            return ext.ExtensionRunResult(
                run_id="", exit_code=exit_code, log_tail=log,
                outputs=list(outputs),
                error="" if exit_code == 0 else "命令失败",
            )

    return FakeExec()


# ---- 创建与布局（契约 §1）----

async def test_create_layout(tmp_path):
    s = ext.ExtensionStore(base_dir=str(tmp_path / "ext"))
    e = _create(s)
    assert e["status"] == "draft"
    assert e["ext_id"].startswith("ext_")
    assert e["runtime_image"] == "bioflow-gf"  # test4py311 → 映射
    assert e["schema_version"] == spec.SPEC_VERSION
    root = s.ext_dir(e["ext_id"])
    assert (root / "extension.json").is_file()
    assert (root / "raw_code" / "quality_stats.py").read_text(encoding="utf-8") == _CODE
    assert (root / "converted_code" / "quality_stats.py").is_file()
    assert (root / "raw_test_data").is_dir()
    assert (root / "validation").is_dir()
    s.delete(e["ext_id"])


def test_create_requires_command_and_code(tmp_path):
    s = ext.ExtensionStore(base_dir=str(tmp_path / "ext_req"))
    with pytest.raises(ext.ExtensionError, match="display_name"):
        s.create_extension(display_name="  ", command="python a.py", code_files=[("a.py", b"x")])
    with pytest.raises(ext.ExtensionError, match="command"):
        s.create_extension(display_name="n", command="", code_files=[("a.py", b"x")])
    with pytest.raises(ext.ExtensionError, match="代码包"):
        s.create_extension(display_name="n", command="python a.py", code_files=[])


def test_create_zip_path_traversal_sanitized(store):
    e = _create(store, code=[("../../evil.py", b"x"), ("ok.py", b"y")])
    root = store.ext_dir(e["ext_id"])
    assert (root / "raw_code" / "evil.py").is_file()
    assert not (root.parent / "evil.py").exists()


def test_load_missing_404(store):
    with pytest.raises(ext.ExtensionNotFoundError):
        store.load_extension("ext_nope")


# ---- 状态机 ----

async def test_status_transitions_valid(store):
    e = _create(store)
    assert (await store.build_extension(e["ext_id"]))["status"] == "built"


async def test_invalid_transition(store):
    e = _create(store)
    with pytest.raises(ext.ExtensionStatusError):
        await store.trial_run(e["ext_id"])  # draft → trial 非法
    await store.build_extension(e["ext_id"])
    with pytest.raises(ext.ExtensionStatusError):
        store.publish(e["ext_id"])  # built → publish 非法（须 validate/confirm）


async def test_build_failed_converges(store):
    e = _create(store)
    # 命令无 .py 文件可引用 + 无 LLM → 规则版仍能给出模板（不会 build_failed）；
    # 构造一个规则版也会失败的场景：空原始命令被创建层拦截。这里直接验证
    # 规则版回退路径本身可用（无 LLM key 时 build 成功且 source=rules）。
    out = await store.build_extension(e["ext_id"])
    assert out["status"] == "built"
    assert out["last_build"]["source"] == "rules"
    assert "{output_dir}" in out["command_template"]
    spec.validate_command_template(out["command_template"])


async def test_llm_invalid_output_falls_back_rules(store, monkeypatch):
    """LLM 返回白名单外占位符 → 视为转换失败 → 规则版回退（build 仍成功）。"""
    e = _create(store)
    monkeypatch.setattr("app.services.bioflow.extension_llm.llm_config",
                        lambda: {"base_url": "http://x", "api_key": "k", "model": "m", "timeout": 1.0})

    async def fake_complete_json(system, user, **kw):
        return {"command_template": "python {code_dir}/a.py --out {not_allowed}",
                "requirements": {}, "nl_triggers": ["t"]}

    monkeypatch.setattr("app.services.bioflow.llm.complete_json", fake_complete_json)
    out = await store.build_extension(e["ext_id"])
    assert out["status"] == "built"
    assert out["last_build"]["source"] == "rules"
    spec.validate_command_template(out["command_template"])


# ---- trial ----

async def test_trial_pass(store):
    e = _create(store, data=[("input.fastq", b">s1\nACGT\n")])
    await store.build_extension(e["ext_id"])
    store.executor = _fake_executor(0, "ok\n", ("result.txt",))
    out = await store.trial_run(e["ext_id"])
    assert out["status"] == "trial_passed"
    assert out["trial"]["ok"] is True
    assert out["trial"]["outputs"] == ["result.txt"]
    vdir = store.validation_dir(e["ext_id"])
    assert (vdir / "trial_1.log").is_file()
    assert (vdir / "trial_1.outputs" / "result.txt").is_file()


async def test_trial_fail(store):
    e = _create(store)
    await store.build_extension(e["ext_id"])
    store.executor = _fake_executor(2, "boom\n", ())
    out = await store.trial_run(e["ext_id"])
    assert out["status"] == "trial_failed"
    assert out["trial"]["ok"] is False
    assert "命令失败" in out["trial"]["error"]


async def test_trial_retry_increments(store):
    e = _create(store)
    await store.build_extension(e["ext_id"])
    store.executor = _fake_executor(1, "fail\n", ())
    await store.trial_run(e["ext_id"])
    store.executor = _fake_executor(0, "ok\n", ("r.txt",))
    out = await store.trial_run(e["ext_id"])
    assert out["trial"]["n"] == 2
    assert out["status"] == "trial_passed"


async def test_trial_without_executor_503_semantics(store):
    e = _create(store)
    await store.build_extension(e["ext_id"])
    with pytest.raises(ext.ExtensionRunError):
        await store.trial_run(e["ext_id"])  # 占位执行器 → 明确错误（路由层 503）


# ---- validate ----

async def test_validate_without_golden(store):
    e = _create(store)
    await store.build_extension(e["ext_id"])
    store.executor = _fake_executor(0, "ok\n")
    await store.trial_run(e["ext_id"])
    out = store.validate(e["ext_id"])
    assert out["status"] == "validated"
    assert out["validation"]["compared"] is False


async def test_validate_golden_pass(store):
    e = _create(store, data=[("sample.fastq", b">s\nACGT\n")],
                golden={"test_input": "sample.fastq", "expected_outputs": ["result.txt"]})
    await store.build_extension(e["ext_id"])
    store.executor = _fake_executor(0, "ok\n", ("result.txt",))
    await store.trial_run(e["ext_id"])
    out = store.validate(e["ext_id"])
    assert out["status"] == "validated"
    assert out["validation"]["compared"] is True
    assert out["validation"]["ok"] is True


async def test_validate_golden_missing_output(store):
    e = _create(store, golden={"test_input": "sample.fastq", "expected_outputs": ["gold.tsv"]})
    await store.build_extension(e["ext_id"])
    store.executor = _fake_executor(0, "ok\n", ("other.txt",))
    await store.trial_run(e["ext_id"])
    out = store.validate(e["ext_id"])
    assert out["status"] == "trial_failed"
    assert out["validation"]["missing_outputs"] == ["gold.tsv"]


# ---- repair ----

async def test_repair_limit(store):
    e = _create(store)
    await store.build_extension(e["ext_id"])
    store.executor = _fake_executor(1, "No module named 'pandas'\n", ())
    await store.trial_run(e["ext_id"])
    store.max_repair_rounds = 1
    out = await store.repair(e["ext_id"])  # 第 1 轮成功（规则版并入 pandas）
    assert out["repair_count"] == 1
    assert "pandas" in out["requirements"]["pip"]
    store.executor = _fake_executor(1, "still fail\n", ())
    await store.trial_run(e["ext_id"])
    with pytest.raises(ext.ExtensionStatusError, match="上限"):
        await store.repair(e["ext_id"])


async def test_repair_requires_failure(store):
    e = _create(store)
    await store.build_extension(e["ext_id"])
    with pytest.raises(ext.ExtensionStatusError):
        await store.repair(e["ext_id"])  # built（未失败）不可修复


# ---- confirm / publish ----

async def test_full_flow_to_publish(store):
    e = _create(store)
    ext_id = e["ext_id"]
    await store.build_extension(ext_id)
    store.executor = _fake_executor(0, "ok\n")
    await store.trial_run(ext_id)
    store.validate(ext_id)
    out = store.confirm(ext_id, nl_triggers=["质量统计", "fastq 统计"])
    assert out["status"] == "confirming"
    assert out["nl_triggers"] == ["质量统计", "fastq 统计"]
    pub = store.publish(ext_id)
    assert pub["status"] == "published"
    # 已发布不可再发布流转（幂等返回）
    assert store.publish(ext_id)["status"] == "published"


async def test_confirm_parameters_reconcile(store):
    e = _create(store, command="python quality_stats.py --fastq input.fastq --cpu {param_cpu} --out results/")
    await store.build_extension(e["ext_id"])
    store.executor = _fake_executor(0)
    await store.trial_run(e["ext_id"])
    store.validate(e["ext_id"])
    # 用户编辑后的 parameters 丢掉 {param_cpu} → 自动补齐
    out = store.confirm(e["ext_id"], parameters=[])
    names = {p["name"] for p in out["parameters"]}
    assert "cpu" in names


# ---- run ----

async def test_run_requires_published(store):
    e = _create(store)
    await store.build_extension(e["ext_id"])
    with pytest.raises(ext.ExtensionStatusError, match="仅已发布"):
        await store.run(e["ext_id"], params={}, files=[])


async def test_run_full(store):
    e = _create(store)
    ext_id = e["ext_id"]
    await store.build_extension(ext_id)
    store.executor = _fake_executor(0, "run ok\n", ("out.tsv",))
    await store.trial_run(ext_id)
    store.validate(ext_id)
    store.confirm(ext_id)
    store.publish(ext_id)

    # 输入文件登记（须在白名单根内：data_root 或扩展目录）
    data_root = Path(store.base_dir).parent
    input_file = data_root / "input.fastq"
    input_file.write_text(">s\nACGT\n", encoding="utf-8")
    result = await store.run(ext_id, params={"cpu": "4"}, files=[str(input_file)])
    assert result["status"] == "done"
    assert result["exit_code"] == 0
    assert "out.tsv" in result["outputs"]
    run_dir = store.runs_dir(ext_id) / result["run_id"]
    assert (run_dir / "run_result.json").is_file()
    assert (run_dir / "inputs" / "input.fastq").is_file()
    assert (run_dir / "outputs" / "out.tsv").is_file()
    runs = store.list_runs(ext_id)
    assert len(runs) == 1
    assert runs[0]["run_id"] == result["run_id"]


async def test_run_rejects_outside_file(store):
    e = _create(store)
    ext_id = e["ext_id"]
    await store.build_extension(ext_id)
    store.executor = _fake_executor(0)
    await store.trial_run(ext_id)
    store.validate(ext_id)
    store.publish(ext_id)
    with pytest.raises(ext.ExtensionRunError, match="白名单"):
        await store.run(ext_id, params={}, files=["/etc/passwd"])


async def test_run_failed_status(store):
    e = _create(store)
    ext_id = e["ext_id"]
    await store.build_extension(ext_id)
    store.executor = _fake_executor(3, "output contract failed\n", ())
    await store.trial_run(ext_id)  # 空输出 → 失败（假执行器 exit=3）
    assert store.load_extension(ext_id)["status"] == "trial_failed"


# ---- 列表 / 删除 ----

async def test_list_published_first(store):
    a = _create(store, name="扩展 A")
    await store.build_extension(a["ext_id"])
    b = _create(store, name="扩展 B")
    await store.build_extension(b["ext_id"])
    store.executor = _fake_executor(0)
    await store.trial_run(b["ext_id"])
    store.validate(b["ext_id"])
    store.publish(b["ext_id"])
    items = store.list_extensions()
    assert items[0]["ext_id"] == b["ext_id"]  # published 优先
    assert items[1]["ext_id"] == a["ext_id"]


def test_delete_guard_trial_running(store):
    e = _create(store)
    # 手工置 trial_running（不可删除守卫；无需 build——直接构造状态）
    store.save_extension({**store.load_extension(e["ext_id"]), "status": "trial_running"})
    with pytest.raises(ext.ExtensionStatusError, match="试运行"):
        store.delete(e["ext_id"])
    store.save_extension({**store.load_extension(e["ext_id"]), "status": "draft"})
    store.delete(e["ext_id"])
    assert store.load_extension(e["ext_id"], required=False) is None


# ---- 并发写合并 ----

async def test_apply_and_write_preserves_fields(store):
    e = _create(store)
    ext_id = e["ext_id"]
    # 两次锁内读-改-写（模拟交错协程写不同字段）：后写不丢先写字段
    store._apply_and_write(ext_id, lambda ext: ext.__setitem__(
        "command_template", "python {code_dir}/a.py --out {output_dir}"))
    store._apply_and_write(ext_id, lambda ext: ext.__setitem__("nl_triggers", ["触发词"]))
    merged = store.load_extension(ext_id)
    assert merged["command_template"] == "python {code_dir}/a.py --out {output_dir}"
    assert merged["nl_triggers"] == ["触发词"]
    assert merged["original_command"]  # 其余字段保留


# ---- 事件总线发布（契约 §4；run 完成后可回放）----

async def test_run_publishes_events(store, tmp_path):
    from app.services.bioflow.events import EventBus

    bus = EventBus(tmp_path / "events")
    e = _create(store)
    ext_id = e["ext_id"]
    await store.build_extension(ext_id)
    store.executor = _fake_executor(0, "run ok\n", ("o.txt",))
    await store.trial_run(ext_id)
    store.validate(ext_id)
    store.publish(ext_id)

    store._bus_override = bus
    result = await store.run(ext_id, params={}, files=[])
    replay = bus.replay(ext_id)
    types = [ev["type"] for ev in replay]
    assert "extension_run_started" in types
    assert "extension_run_done" in types
    assert result["status"] == "done"