Files
PlantAI-dev/tests/unit/test_wgs_index_backend.py
2026-08-03 18:42:47 +08:00

114 lines
6.0 KiB
Python

import json
import stat
from pathlib import Path
from apps.web import pipeline_backend
from gfa_core.mas.stage_capabilities import CAPABILITIES
from gfa_core.mas.stage_catalog import EXECUTABLE_STAGES, STAGE_DEPENDENCIES
from gfa_core.mas.tools.stage_identity import resolve_stage_identity
from gfa_core.mas.tools.stage_status import get_stage_status
def _manifest() -> dict:
components = (".0123", ".amb", ".ann", ".bwt.2bit.64", ".pac")
artifacts = {
"reference_fasta": {"role": "reference_fasta", "sha256": "a" * 64},
"reference_fai": {"role": "reference_fai", "sha256": "b" * 64},
"reference_dict": {"role": "reference_dict", "sha256": "c" * 64},
}
artifacts.update({f"bwa_mem2_index_component.{part}": {"role": "bwa_mem2_index_component", "component": part, "sha256": "d" * 64} for part in components})
return {"stage": "wgs_index", "status": "success", "prefix": "Cohort", "job_id": "index-job", "wgs_reference_id": "wgs-reference:" + "e" * 64, "wgs_analysis_id": "wgs-analysis:" + "f" * 64, "artifacts": artifacts}
def test_wgs_index_backend_requires_exact_roles_and_fresh_manifest(monkeypatch, tmp_path: Path) -> None:
reference = tmp_path / "input" / "reference.fasta"; reference.parent.mkdir(parents=True)
reference.write_text(">chr1\nACGT\n", encoding="utf-8")
workspace = tmp_path / "workspace"; final = workspace / "final"; final.mkdir(parents=True)
manifest_path = final / "Cohort.wgs_index.manifest.json"
manifest_path.write_text(json.dumps({"status": "success"}), encoding="utf-8")
captured = {}
def fake_run(command, *_args, **_kwargs):
captured["command"] = list(command)
manifest_path.write_text(json.dumps(_manifest()), encoding="utf-8")
return 0
monkeypatch.setattr(pipeline_backend, "run_cmd_logged", fake_run)
message, returned = pipeline_backend.run_stage(
stage="wgs_index", job=pipeline_backend.JobInfo("index-job", str(reference.parent), str(workspace), str(workspace / "logs"), "now"),
prefix="Cohort", wgs_reference_fasta=str(reference), wgs_samtools="sam", wgs_bwa_mem2="bwa", wgs_gatk="gatk", wgs_threads=3,
)
assert "completed" in message.lower() and returned == manifest_path
assert captured["command"][:3] == [pipeline_backend.sys.executable, "-m", "apps.cli.run_wgs_index"]
assert captured["command"][captured["command"].index("--reference-fasta") + 1] == str(reference)
assert not pipeline_backend._valid_wgs_index_manifest({**_manifest(), "artifacts": {}}, job_id="index-job", prefix="Cohort")
identityless = _manifest(); identityless.pop("wgs_reference_id")
assert not pipeline_backend._valid_wgs_index_manifest(identityless, job_id="index-job", prefix="Cohort")
def test_wgs_index_is_a_registered_root_with_exact_reference_identity_and_preflight(tmp_path: Path) -> None:
assert STAGE_DEPENDENCIES["wgs_index"] == [] and "wgs_index" in EXECUTABLE_STAGES
assert {item.role for item in CAPABILITIES["wgs_index"].provides} == {"reference_fasta", "reference_fai", "reference_dict", "bwa_mem2_index_component"}
identity = resolve_stage_identity("wgs_index", "Cohort", {"wgsReferenceId": "ref", "wgsAnalysisId": "analysis"})
assert identity.wgs_reference_id == "ref" and identity.wgs_analysis_id == "analysis"
missing = get_stage_status("index-job", "Cohort", "wgs_index", workspace_root=tmp_path)
assert "wgs_reference_id" in missing.errors[0]
try:
pipeline_backend._resolve_wgs_reference_fasta(tmp_path, "")
except ValueError as exc:
assert "explicit" in str(exc)
else:
raise AssertionError("missing root reference must fail preflight")
def _backend_tool(path: Path, kind: str) -> Path:
path.write_text(
"#!/usr/bin/env python3\n"
"import sys\n"
f"KIND = {kind!r}\n"
"args = sys.argv[1:]\n"
"if '--version' in args or args == ['version']:\n"
" print(KIND + ' 1.0'); raise SystemExit(0)\n"
"if KIND == 'samtools':\n"
" ref = args[args.index('faidx') + 1]; open(ref + '.fai', 'w').write('chr1\\t8\\t6\\t8\\t9\\n')\n"
"elif KIND == 'gatk':\n"
" out = args[args.index('-O') + 1]; open(out, 'w').write('@HD\\tVN:1.6\\n@SQ\\tSN:chr1\\tLN:8\\n')\n"
"else:\n"
" ref = args[-1]\n"
" for suffix in ('.0123', '.amb', '.ann', '.bwt.2bit.64', '.pac'):\n"
" open(ref + suffix, 'w').write(suffix + '\\n')\n",
encoding="utf-8",
)
path.chmod(path.stat().st_mode | stat.S_IXUSR)
return path
def test_wgs_index_backend_accepts_valid_cli_reuse_with_same_run_id(tmp_path: Path) -> None:
reference = tmp_path / "input" / "reference.fasta"
reference.parent.mkdir(parents=True)
reference.write_text(">chr1\nACGTACGT\n", encoding="utf-8")
tools = {
"samtools": _backend_tool(tmp_path / "samtools", "samtools"),
"bwa_mem2": _backend_tool(tmp_path / "bwa-mem2", "bwa-mem2"),
"gatk": _backend_tool(tmp_path / "gatk", "gatk"),
}
workspace = tmp_path / "workspace"
job = pipeline_backend.JobInfo(
"index-job", str(reference.parent), str(workspace), str(workspace / "logs"), "now",
)
params = {
"stage": "wgs_index", "job": job, "prefix": "Cohort",
"wgs_reference_fasta": str(reference), "wgs_samtools": str(tools["samtools"]),
"wgs_bwa_mem2": str(tools["bwa_mem2"]), "wgs_gatk": str(tools["gatk"]),
"wgs_threads": 2,
}
first_message, manifest_path = pipeline_backend.run_stage(**params)
assert "completed" in first_message.lower() and manifest_path is not None
first = json.loads(manifest_path.read_text(encoding="utf-8"))
second_message, second_path = pipeline_backend.run_stage(**params)
assert "completed" in second_message.lower() and second_path == manifest_path
second = json.loads(second_path.read_text(encoding="utf-8"))
assert second["run_id"] == first["run_id"]
assert second["wgs_reference_id"] == first["wgs_reference_id"]
assert second["wgs_analysis_id"] == first["wgs_analysis_id"]
assert second.get("reuse", {}).get("count") == 1