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

260 lines
8.6 KiB
Python

from __future__ import annotations
import hashlib
import json
from pathlib import Path
import pytest
SAMPLES = "wgs-samples:" + "a" * 64
REFERENCE = "wgs-reference:" + "b" * 64
MARKDUP_ANALYSIS = "wgs-analysis:" + "c" * 64
INDEX_ANALYSIS = "wgs-analysis:" + "d" * 64
MARKDUP_RUN = "4" * 32
INDEX_RUN = "2" * 32
def _record(workspace: Path, path: Path, role: str, sample_id: str = "") -> dict:
record = {
"artifact": role,
"role": role,
"path": str(path.relative_to(workspace)),
"full_path": str(path.resolve()),
"size": path.stat().st_size,
"size_bytes": path.stat().st_size,
"sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
}
if sample_id:
record["sample_id"] = sample_id
return record
def _authorities(workspace: Path) -> tuple[Path, Path]:
final = workspace / "final"
markdup_bundle = final / ".wgs_markdup_runs" / "Cohort" / MARKDUP_RUN
markdup_bundle.mkdir(parents=True)
markdup_artifacts = {}
for sample in ("S1", "S2", "S3"):
bam = markdup_bundle / f"Cohort.wgs_markdup.{sample}.markdup.bam"
bai = Path(f"{bam}.bai")
bam.write_text(f"{sample}\n", encoding="utf-8")
bai.write_bytes(b"BAI")
markdup_artifacts[f"markdup_bam.{sample}"] = _record(
workspace, bam, "markdup_bam", sample
)
markdup_artifacts[f"markdup_bam_index.{sample}"] = _record(
workspace, bai, "markdup_bam_index", sample
)
markdup = final / "Cohort.wgs_markdup.manifest.json"
markdup.write_text(
json.dumps(
{
"manifest_version": "1.0",
"stage": "wgs_markdup",
"status": "success",
"prefix": "Cohort",
"job_id": "job-1",
"run_id": MARKDUP_RUN,
"workspace": str(workspace.resolve()),
"wgs_sample_set_id": SAMPLES,
"wgs_reference_id": REFERENCE,
"wgs_analysis_id": MARKDUP_ANALYSIS,
"artifacts": markdup_artifacts,
}
),
encoding="utf-8",
)
index_bundle = final / ".wgs_index_runs" / "Cohort" / INDEX_RUN
index_bundle.mkdir(parents=True)
fai = index_bundle / "Cohort.wgs_index.reference.fasta.fai"
fai.write_text("chr1\t8\t0\t0\t0\nchr2\t2\t0\t0\t0\n", encoding="utf-8")
index = final / "Cohort.wgs_index.manifest.json"
index.write_text(
json.dumps(
{
"manifest_version": "1.0",
"stage": "wgs_index",
"status": "success",
"prefix": "Cohort",
"job_id": "job-1",
"run_id": INDEX_RUN,
"workspace": str(workspace.resolve()),
"wgs_reference_id": REFERENCE,
"wgs_analysis_id": INDEX_ANALYSIS,
"artifacts": {
"reference_fai": _record(workspace, fai, "reference_fai")
},
}
),
encoding="utf-8",
)
return markdup, index
def _samtools(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
tool = tmp_path / "samtools"
tool.write_text(
"""#!/usr/bin/env python3
import os, pathlib, sys
if sys.argv[1:] == ['--version']:
print('samtools 1.20')
raise SystemExit(0)
cmd = sys.argv[1]
bam = pathlib.Path(sys.argv[-1])
sample = bam.name.split('.')[2]
if os.environ.get('FAIL_COMMAND') == f'{cmd}:{sample}':
print('controlled failure', file=sys.stderr)
raise SystemExit(23)
if cmd == 'flagstat':
print('2 + 0 in total (QC-passed reads + QC-failed reads)')
print('2 + 0 mapped (100.00% : N/A)')
raise SystemExit(0)
if cmd == 'depth':
# Two covered bases out of a ten-base, two-contig reference.
print('chr1\\t1\\t2')
print('chr1\\t2\\t1')
raise SystemExit(0)
raise SystemExit(2)
""",
encoding="utf-8",
)
tool.chmod(0o755)
return tool
def _argv(workspace: Path, markdup: Path, index: Path, samtools: Path) -> list[str]:
return [
"--workspace", str(workspace),
"--job-id", "job-1",
"--prefix", "Cohort",
"--markdup-manifest", str(markdup),
"--index-manifest", str(index),
"--samtools", str(samtools),
"--coverage-threshold", "1",
"--threads", "2",
"--max-jobs", "2",
]
def test_wgs_bam_qc_exposes_import_safe_cli() -> None:
from apps.cli.run_wgs_bam_qc import build_parser, main
assert callable(build_parser)
assert callable(main)
def test_two_contig_denominator_includes_uncovered_positions_and_keeps_low_depth_samples(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
from apps.cli.run_wgs_bam_qc import main
workspace = tmp_path / "workspace"
markdup, index = _authorities(workspace)
samtools = _samtools(tmp_path, monkeypatch)
assert main(_argv(workspace, markdup, index, samtools)) == 0
manifest = json.loads(
(workspace / "final" / "Cohort.wgs_bam_qc.manifest.json").read_text()
)
assert manifest["status"] == "success"
assert manifest["parameters"]["metric_name"] == "coverage_ge_1x"
assert manifest["metrics"]["input_sample_ids"] == ["S1", "S2", "S3"]
assert manifest["metrics"]["output_sample_ids"] == ["S1", "S2", "S3"]
assert manifest["metrics"]["total_reference_bases"] == 10
assert manifest["metrics"]["coverage_ge_1x"] == pytest.approx(0.2)
assert len(manifest["artifacts"]) == 8
assert {a["role"] for a in manifest["artifacts"].values()} == {
"bam_flagstat", "bam_depth_summary", "bam_qc_summary", "bam_qc_report"
}
depth_paths = [
Path(a["full_path"])
for a in manifest["artifacts"].values()
if a["role"] == "bam_depth_summary"
]
assert len(depth_paths) == 3
for path in depth_paths:
row = path.read_text().splitlines()[1].split("\t")
assert row[0] in {"S1", "S2", "S3"}
assert row[1:] == ["10", "3", "0.3", "2", "0.2"]
report = next(
Path(a["full_path"])
for a in manifest["artifacts"].values()
if a["role"] == "bam_qc_report"
)
assert "coverage_ge_1x" in report.read_text()
assert "No sample was excluded" in report.read_text()
@pytest.mark.parametrize("failure", ("flagstat:S2", "depth:S1"))
def test_each_samtools_failure_fails_stage_without_replaying_old_success(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
failure: str,
) -> None:
from apps.cli.run_wgs_bam_qc import main
workspace = tmp_path / "workspace"
markdup, index = _authorities(workspace)
samtools = _samtools(tmp_path, monkeypatch)
authority = workspace / "final" / "Cohort.wgs_bam_qc.manifest.json"
authority.write_text(
json.dumps({"stage": "wgs_bam_qc", "status": "success", "sentinel": True}),
encoding="utf-8",
)
old_hash = hashlib.sha256(authority.read_bytes()).hexdigest()
monkeypatch.setenv("FAIL_COMMAND", failure)
assert main(_argv(workspace, markdup, index, samtools)) == 1
assert hashlib.sha256(authority.read_bytes()).hexdigest() == old_hash
sidecars = list(
(workspace / "final").glob("Cohort.wgs_bam_qc.failed.*.manifest.json")
)
assert len(sidecars) == 1
assert json.loads(sidecars[0].read_text())["artifacts"] == {}
def test_reference_identity_mismatch_fails_closed(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
from apps.cli.run_wgs_bam_qc import main
workspace = tmp_path / "workspace"
markdup, index = _authorities(workspace)
payload = json.loads(index.read_text())
payload["wgs_reference_id"] = "wgs-reference:" + "e" * 64
index.write_text(json.dumps(payload))
samtools = _samtools(tmp_path, monkeypatch)
assert main(_argv(workspace, markdup, index, samtools)) == 1
failure = json.loads(
(workspace / "final" / "Cohort.wgs_bam_qc.manifest.json").read_text()
)
assert failure["artifacts"] == {}
def test_unsafe_prefix_writes_safe_unique_failure_sidecar(
tmp_path: Path,
) -> None:
from apps.cli.run_wgs_bam_qc import main
workspace = tmp_path / "workspace"
assert main([
"--workspace", str(workspace),
"--job-id", "job-1",
"--prefix", "../escape",
"--markdup-manifest", "unused",
"--index-manifest", "unused",
"--samtools", "unused",
"--coverage-threshold", "1",
"--threads", "1",
"--max-jobs", "1",
]) == 1
sidecars = list(
(workspace / "final").glob("wgs_bam_qc.failed.*.manifest.json")
)
assert len(sidecars) == 1
assert sidecars[0].parent == workspace / "final"
failure = json.loads(sidecars[0].read_text())
assert failure["status"] == "failed"
assert failure["artifacts"] == {}